from __future__ import annotations

from .fen import ParsedFEN, parse_fen
from .utils import Board, Color, HandPiece, Move, Piece, PieceType, setup_mode_to_code


def assign_piece_ids_with_state(
    current_fen: str, previous_state: ParsedFEN | None = None, move: Move | None = None
) -> ParsedFEN:
    current_state = parse_fen(current_fen)

    if not previous_state:
        return _assign_canonical_ids(current_state)

    if move:
        return _assign_ids_with_explicit_move(current_state, previous_state, move)

    return _assign_ids_with_position_matching(current_state, previous_state)


def _assign_canonical_ids(state: ParsedFEN) -> ParsedFEN:
    result = _clone_state(state)
    mode_code = setup_mode_to_code[state.mode]

    groups = _group_pieces_by_type_and_color(result)
    for key, pieces in groups.items():
        color, ptype = key.split("|")
        sorted_pieces = _sort_pieces_canonically(pieces)
        for index, piece in enumerate(sorted_pieces, start=1):
            piece.id = f"{mode_code}-{color}-{ptype}-{index}"
    return result


def _assign_ids_with_explicit_move(
    current_state: ParsedFEN, previous_state: ParsedFEN, move: Move
) -> ParsedFEN:
    result = _clone_state_without_ids(current_state)
    mode_code = setup_mode_to_code[current_state.mode]

    affected = _identify_affected_pieces(move)
    _preserve_unaffected_piece_ids(result, previous_state, affected)
    _handle_affected_pieces(result, previous_state, move, mode_code)
    _assign_missing_ids(result, mode_code)
    return result


def _identify_affected_pieces(move: Move) -> dict[str, object]:
    captured_pieces: list[dict[str, object]] = []
    hand_piece_types: list[dict[str, object]] = []
    moving_piece: dict[str, object] | None = None
    placement_square: str | None = None

    if move.type in {"route", "tsuke", "capture", "betray"} and move.from_:
        from_rank, from_file, from_tier = map(int, move.from_.split("-"))
        moving_piece = {
            "square": f"{from_rank}-{from_file}",
            "tier": from_tier,
            "type": move.piece,
            "color": move.color,
        }

    if move.type in {"capture", "betray"} and move.captured:
        for captured in move.captured:
            captured_pieces.append(
                {
                    "square": move.to,
                    "tier": captured.tier,
                    "type": captured.type,
                    "color": captured.color,
                }
            )
        if move.type == "betray":
            for captured in move.captured:
                hand_piece_types.append({"color": move.color, "type": captured.type})

    if move.type == "arata":
        hand_piece_types.append({"color": move.color, "type": move.piece})
        placement_square = move.to

    return {
        "movingPiece": moving_piece,
        "capturedPieces": captured_pieces,
        "placementSquare": placement_square,
        "handPieceTypes": hand_piece_types,
    }


def _preserve_unaffected_piece_ids(
    current_state: ParsedFEN,
    previous_state: ParsedFEN,
    affected: dict[str, object],
) -> None:
    def is_piece_affected(square: str, tier: int, ptype: PieceType, color: Color) -> bool:
        moving = affected["movingPiece"]
        if (
            moving
            and moving["square"] == square
            and moving["tier"] == tier
            and moving["type"] == ptype
            and moving["color"] == color
        ):  # type: ignore[index]
            return True
        if any(
            c["square"] == square
            and c["tier"] == tier
            and c["type"] == ptype
            and c["color"] == color  # type: ignore[index]
            for c in affected["capturedPieces"]  # type: ignore[index]
        ):
            return True
        return affected["placementSquare"] == square  # type: ignore[index]

    for rank_idx in range(9):
        for file_idx in range(9):
            square = f"{rank_idx + 1}-{9 - file_idx}"
            current_tower = current_state.board[rank_idx][file_idx]
            previous_tower = previous_state.board[rank_idx][file_idx]

            for tier_idx in range(min(len(current_tower), len(previous_tower))):
                current_piece = current_tower[tier_idx]
                previous_piece = previous_tower[tier_idx]
                if (
                    current_piece
                    and previous_piece
                    and current_piece.type == previous_piece.type
                    and current_piece.color == previous_piece.color
                    and previous_piece.id
                    and not is_piece_affected(
                        square, tier_idx + 1, current_piece.type, current_piece.color
                    )
                ):
                    current_piece.id = previous_piece.id

    for current_hand in current_state.hand:
        affected_hand = any(
            hp["color"] == current_hand.color and hp["type"] == current_hand.type  # type: ignore[index]
            for hp in affected["handPieceTypes"]  # type: ignore[index]
        )
        if not affected_hand:
            prev_hand = next(
                (
                    hp
                    for hp in previous_state.hand
                    if hp.type == current_hand.type and hp.color == current_hand.color
                ),
                None,
            )
            if prev_hand and prev_hand.id:
                current_hand.id = prev_hand.id


def _handle_affected_pieces(
    current_state: ParsedFEN,
    previous_state: ParsedFEN,
    move: Move,
    mode_code: int,
) -> None:
    if move.type in {"route", "tsuke"} or move.type == "capture":
        _handle_regular_move(current_state, previous_state, move)
    elif move.type == "arata":
        _handle_arata_move(current_state, previous_state, move, mode_code)
    elif move.type == "betray":
        _handle_regular_move(current_state, previous_state, move)
        if move.captured:
            for captured in move.captured:
                hand_piece = next(
                    (
                        hp
                        for hp in current_state.hand
                        if hp.type == captured.type and hp.color == move.color
                    ),
                    None,
                )
                if hand_piece and not hand_piece.id:
                    hand_piece.id = _find_next_available_id(
                        current_state, mode_code, move.color, captured.type
                    )


def _handle_regular_move(current_state: ParsedFEN, previous_state: ParsedFEN, move: Move) -> None:
    if not move.from_:
        return
    from_rank, from_file, from_tier = map(int, move.from_.split("-"))
    from_square = f"{from_rank}-{from_file}"
    moved_piece = _find_piece_at_square(
        previous_state.board, from_square, move.piece, move.color, from_tier
    )
    if not moved_piece or not moved_piece.id:
        return
    to_rank, to_file, to_tier = map(int, move.to.split("-"))
    to_square = f"{to_rank}-{to_file}"
    new_piece = _find_piece_at_square(
        current_state.board, to_square, move.piece, move.color, to_tier
    )
    if new_piece:
        new_piece.id = moved_piece.id


def _handle_arata_move(
    current_state: ParsedFEN, previous_state: ParsedFEN, move: Move, mode_code: int
) -> None:
    prev_hand_piece = next(
        (hp for hp in previous_state.hand if hp.type == move.piece and hp.color == move.color), None
    )
    if not prev_hand_piece or not prev_hand_piece.id:
        return

    to_rank, to_file, to_tier = map(int, move.to.split("-"))
    to_square = f"{to_rank}-{to_file}"
    placed_piece = _find_piece_at_square(
        current_state.board, to_square, move.piece, move.color, to_tier
    )
    if placed_piece:
        placed_piece.id = prev_hand_piece.id

    current_hand_piece = next(
        (hp for hp in current_state.hand if hp.type == move.piece and hp.color == move.color), None
    )
    if current_hand_piece:
        current_hand_piece.id = _find_next_available_id(
            current_state, mode_code, move.color, move.piece
        )


def _assign_ids_with_position_matching(
    current_state: ParsedFEN, previous_state: ParsedFEN
) -> ParsedFEN:
    result = _clone_state(current_state)
    mode_code = setup_mode_to_code[current_state.mode]

    hand_ids_used_for_arata: set[str] = set()

    for prev_hand_piece in previous_state.hand:
        current_hand_piece = next(
            (
                hp
                for hp in result.hand
                if hp.type == prev_hand_piece.type and hp.color == prev_hand_piece.color
            ),
            None,
        )
        prev_count = prev_hand_piece.count
        curr_count = current_hand_piece.count if current_hand_piece else 0

        if prev_count > curr_count and prev_hand_piece.id:
            new_board_piece = _find_new_board_piece(
                result.board, previous_state.board, prev_hand_piece.type, prev_hand_piece.color
            )
            if new_board_piece:
                new_board_piece.id = prev_hand_piece.id
                hand_ids_used_for_arata.add(prev_hand_piece.id)

    for rank_idx in range(9):
        for file_idx in range(9):
            current_tower = result.board[rank_idx][file_idx]
            previous_tower = previous_state.board[rank_idx][file_idx]
            for tier_idx in range(min(len(current_tower), len(previous_tower))):
                current_piece = current_tower[tier_idx]
                previous_piece = previous_tower[tier_idx]
                if (
                    current_piece
                    and not current_piece.id
                    and previous_piece
                    and current_piece.type == previous_piece.type
                    and current_piece.color == previous_piece.color
                    and previous_piece.id
                ):
                    current_piece.id = previous_piece.id

    for hp in result.hand:
        if hp.id:
            continue
        prev_hand = next(
            (p for p in previous_state.hand if p.type == hp.type and p.color == hp.color), None
        )
        if prev_hand and prev_hand.id and prev_hand.id not in hand_ids_used_for_arata:
            hp.id = prev_hand.id

    _assign_missing_ids(result, mode_code)
    return result


def _assign_missing_ids(state: ParsedFEN, mode_code: int) -> None:
    used_ids = set(_get_all_used_ids(state))
    groups = _group_pieces_by_type_and_color(state)

    for key, pieces in groups.items():
        color, ptype = key.split("|")
        for piece in pieces:
            if not piece.id:
                number = 1
                candidate = f"{mode_code}-{color}-{ptype}-{number}"
                while candidate in used_ids:
                    number += 1
                    candidate = f"{mode_code}-{color}-{ptype}-{number}"
                piece.id = candidate
                used_ids.add(candidate)


def _find_next_available_id(
    state: ParsedFEN, mode_code: int, color: Color, ptype: PieceType
) -> str:
    used_ids = set(_get_all_used_ids(state))
    number = 1
    candidate = f"{mode_code}-{color}-{ptype}-{number}"
    while candidate in used_ids:
        number += 1
        candidate = f"{mode_code}-{color}-{ptype}-{number}"
    return candidate


def _get_all_used_ids(state: ParsedFEN) -> list[str]:
    ids: list[str] = []
    for rank in state.board:
        for file in rank:
            for piece in file:
                if piece and piece.id:
                    ids.append(piece.id)
    for hp in state.hand:
        if hp.id:
            ids.append(hp.id)
    return ids


def _group_pieces_by_type_and_color(state: ParsedFEN) -> dict[str, list[object]]:
    groups: dict[str, list[object]] = {}
    for rank in state.board:
        for file in rank:
            for piece in file:
                if piece:
                    key = f"{piece.color}|{piece.type}"
                    groups.setdefault(key, []).append(piece)
    for hp in state.hand:
        key = f"{hp.color}|{hp.type}"
        groups.setdefault(key, []).append(hp)
    return groups


def _sort_pieces_canonically(pieces: list[object]) -> list[object]:
    def sort_key(p: object) -> tuple[int, int, int, bool]:
        if isinstance(p, HandPiece):
            return (10, 10, 10, True)
        assert isinstance(p, Piece)
        rank, file = map(int, p.square.split("-"))
        return (rank, file, p.tier, False)

    return sorted(pieces, key=sort_key)


def _find_piece_at_square(
    board: Board,
    square: str,
    ptype: PieceType | None = None,
    color: Color | None = None,
    tier: int | None = None,
) -> Piece | None:
    parts = square.split("-")
    rank, file = int(parts[0]), int(parts[1])
    if rank < 1 or rank > 9 or file < 1 or file > 9:
        return None
    tower = board[rank - 1][9 - file]
    for p in tower:
        if (
            p
            and (ptype is None or p.type == ptype)
            and (color is None or p.color == color)
            and (tier is None or p.tier == tier)
        ):
            return p
    return None


def _find_new_board_piece(
    current_board: Board,
    previous_board: Board,
    ptype: PieceType,
    color: Color,
) -> Piece | None:
    for rank_idx in range(9):
        for file_idx in range(9):
            current_tower = current_board[rank_idx][file_idx]
            previous_tower = previous_board[rank_idx][file_idx]

            for tier_idx, current_piece in enumerate(current_tower):
                if (
                    current_piece
                    and current_piece.type == ptype
                    and current_piece.color == color
                    and not current_piece.id
                ):
                    previous_piece = (
                        previous_tower[tier_idx] if tier_idx < len(previous_tower) else None
                    )
                    if (
                        not previous_piece
                        or previous_piece.type != ptype
                        or previous_piece.color != color
                    ):
                        return current_piece

    return None


def _clone_state(state: ParsedFEN) -> ParsedFEN:
    return ParsedFEN(
        board=[
            [[Piece(**vars(piece)) if piece else None for piece in file] for file in rank]
            for rank in state.board
        ],
        hand=[HandPiece(**vars(hp)) for hp in state.hand],
        turn=state.turn,
        mode=state.mode,
        drafting=dict(state.drafting),
        move_number=state.move_number,
    )


def _clone_state_without_ids(state: ParsedFEN) -> ParsedFEN:
    return ParsedFEN(
        board=[
            [
                [Piece(**{**vars(piece), "id": None}) if piece else None for piece in file]
                for file in rank
            ]
            for rank in state.board
        ],
        hand=[HandPiece(**{**vars(hp), "id": None}) for hp in state.hand],
        turn=state.turn,
        mode=state.mode,
        drafting=dict(state.drafting),
        move_number=state.move_number,
    )


# JS-style alias
assignPieceIdsWithState = assign_piece_ids_with_state
