from collections import Counter
from enum import auto
from enum import Enum
from enum import StrEnum
from pathlib import Path


FACE_VALUES = {
    "A": 14,
    "K": 13,
    "Q": 12,
    "J": 11,
    "T": 10,
}


def card_value(card: str) -> int:
    if card.isnumeric():
        return int(card)

    return FACE_VALUES[card]


class HandType(Enum):
    HIGH_CARD = auto()
    ONE_PAIR = auto()
    TWO_PAIR = auto()
    THREE_OF_A_KIND = auto()
    FULL_HOUSE = auto()
    FOUR_OF_A_KIND = auto()
    FIVE_OF_A_KIND = auto()


class Hand:
    def __init__(self, cards: str, bid: int, use_jokers=False):
        self.cards = cards
        self.bid = bid

        self.hand_type: HandType = HandType.HIGH_CARD

        counts = Counter(self.cards)
        jokers = counts["J"]

        for card, count in counts.most_common():
            if count == 5 or (use_jokers and card != "J" and count + jokers >= 5):
                self.hand_type = HandType.FIVE_OF_A_KIND
                if count != 5:
                    jokers -= 5 - count
                break
            # It's always better to use the jokers to improve a hand
            if use_jokers and card == "J":
                continue
            if count == 4 or count + jokers >= 4:
                self.hand_type = HandType.FOUR_OF_A_KIND
                if count != 4:
                    jokers -= 4 - count
                break
            if count == 3 or count + jokers >= 3:
                if self.hand_type == HandType.HIGH_CARD:
                    self.hand_type = HandType.THREE_OF_A_KIND
                elif self.hand_type == HandType.ONE_PAIR:
                    self.hand_type = HandType.FULL_HOUSE
                if count != 3:
                    jokers -= 3 - count
                continue
            if count == 2 or count + jokers >= 2:
                if self.hand_type == HandType.HIGH_CARD:
                    self.hand_type = HandType.ONE_PAIR
                elif self.hand_type == HandType.ONE_PAIR:
                    self.hand_type = HandType.TWO_PAIR
                elif self.hand_type == HandType.THREE_OF_A_KIND:
                    self.hand_type = HandType.FULL_HOUSE
                if count != 2:
                    jokers -= 2 - count
                continue

    def __lt__(self, other: "Hand") -> bool:
        # print(f"Comparing {self} to {other}")
        if self.hand_type == other.hand_type:
            for s, o in zip(self.cards, other.cards):
                if s != o:
                    return card_value(s) < card_value(o)

        return self.hand_type.value < other.hand_type.value

    def __repr__(self) -> str:
        return f"<{self.cards} {self.bid}: {self.hand_type}>"


def part1(input: Path) -> int:
    hands: list[Hand] = []
    with input.open() as f:
        for line in f:
            line = line.strip()
            parts = line.partition(" ")
            hands.append(Hand(parts[0], int(parts[2])))

    answer = 0
    for i, hand in enumerate(sorted(hands)):
        print(hand)
        answer += (i + 1) * hand.bid

    return answer


def part2(input: Path) -> int:
    FACE_VALUES["J"] = 1

    hands: list[Hand] = []
    with input.open() as f:
        for line in f:
            line = line.strip()
            parts = line.partition(" ")
            hands.append(Hand(parts[0], int(parts[2]), use_jokers=True))

    answer = 0
    for i, hand in enumerate(sorted(hands)):
        print(hand)
        answer += (i + 1) * hand.bid

    return answer


if __name__ == "__main__":
    input = Path("input.txt")
    # input = Path("sub-test.txt")
    # input = Path("sample.txt")
    result1 = part1(input)
    print("part1", result1)

    result2 = part2(input)
    print("part2", result2)