import random

from app.schemas.grade_1.join_corresponding_sums import (
    AdditionExpression,
    JoinCorrespondingSumsProblem,
    SumConnection,
)


def join_corresponding_sums(
    pair_count: int = 3,
    min_sum: int = 2,
    max_sum: int = 10,
    min_addend: int = 1,
    max_addend: int = 9,
    seed: int | None = None,
) -> dict:
    """Generate addition expressions that students connect by equal sums."""
    if pair_count < 1:
        raise ValueError("pair_count must be at least 1")
    if min_sum > max_sum:
        raise ValueError("min_sum must be less than or equal to max_sum")
    if min_addend > max_addend:
        raise ValueError("min_addend must be less than or equal to max_addend")

    expressions_by_total: dict[int, list[tuple[int, int]]] = {}
    for total in range(min_sum, max_sum + 1):
        expressions = []
        for first_addend in range(min_addend, max_addend + 1):
            second_addend = total - first_addend
            if min_addend <= second_addend <= max_addend:
                expressions.append((first_addend, second_addend))

        if len(expressions) >= 2:
            expressions_by_total[total] = expressions

    available_totals = list(expressions_by_total)
    if len(available_totals) < pair_count:
        raise ValueError("sum and addend ranges must contain enough matchable sums")

    rng = random.Random(seed)
    selected_totals = rng.sample(available_totals, pair_count)
    left_items: list[dict] = []
    right_items: list[dict] = []

    for index, total in enumerate(selected_totals, start=1):
        left_expression, right_expression = rng.sample(expressions_by_total[total], 2)
        match_id = f"sum-{total}"
        left_items.append(
            {
                "first_addend": left_expression[0],
                "second_addend": left_expression[1],
                "total": total,
                "match_id": match_id,
            }
        )
        right_items.append(
            {
                "first_addend": right_expression[0],
                "second_addend": right_expression[1],
                "total": total,
                "match_id": match_id,
            }
        )

    rng.shuffle(left_items)
    rng.shuffle(right_items)

    left_expressions = [
        AdditionExpression(position=index + 1, **item)
        for index, item in enumerate(left_items)
    ]
    right_expressions = [
        AdditionExpression(position=index + 1, **item)
        for index, item in enumerate(right_items)
    ]
    right_positions_by_match_id = {
        expression.match_id: expression.position for expression in right_expressions
    }
    answer_key = [
        SumConnection(
            match_id=expression.match_id,
            total=expression.total,
            left_position=expression.position,
            right_position=right_positions_by_match_id[expression.match_id],
        )
        for expression in left_expressions
    ]

    problem = JoinCorrespondingSumsProblem(
        left_expressions=left_expressions,
        right_expressions=right_expressions,
        answer_key=answer_key,
    )

    return problem.model_dump()
