from pathlib import Path


WINS_OVER = {
    "A": "C",
    "B": "A",
    "C": "B",
    "X": "C",
    "Y": "A",
    "Z": "B",
}


BASE_SCORE = {
    "A": 1,
    "B": 2,
    "C": 3,
    "X": 1,
    "Y": 2,
    "Z": 3,
}


def score(a: str, b: str) -> int:
    result = BASE_SCORE.get(b)

    if BASE_SCORE[a] == BASE_SCORE[b]:
        # Draw
        result += 3
    elif a == WINS_OVER[b]:
        # Win
        result += 6

    return result


def winning_shape(other: str, goal: str) -> str:
    if goal == "Y":
        # Draw
        return other

    if goal == "X":
        # Win
        return WINS_OVER[other]

    if goal == "Z":
        # Lose
        return next(k for k, v in WINS_OVER.items() if v == other)


# Setup

# Example data
# data = """
# A Y
# B X
# C Z
# """

data = (Path(__file__).parents[1] / "input").read_text()

# Problem solving - part 1
scores = [
    score(*parts) for line in data.splitlines() if (parts := line.split(maxsplit=1))
]

total = sum(scores)
print(total)

# Problem solving - part 2
assert winning_shape("B", "Z") == "C"  # Win against paper -> scissors
assert winning_shape("B", "Y") == "B"  # Draw against paper -> paper
assert winning_shape("B", "X") == "A"  # Lose against paper -> rock

scores = [
    score(parts[0], winning_shape(parts[0], parts[1]))
    for line in data.splitlines()
    if (parts := line.split(maxsplit=1))
]

total = sum(scores)
print(total)