from __future__ import annotations

import math

from .fen import encode_fen, parse_fen
from .utils import (
    ARATA,
    BETRAY,
    SQUARES,
    TAKE,
    TSUKE,
    Board,
    Color,
    HandPiece,
    Move,
    MoveType,
    Piece,
    PieceType,
    convert,
    get,
    get_top,
    put,
    remove,
    remove_top,
    update_hand,
)
from .utils import piece as piece_type

# Directions are ordered to mirror the JS implementation
dirs: list[tuple[int, int]] = [
    (-1, 1),
    (-1, 0),
    (-1, -1),
    (0, 1),
    (0, -1),
    (1, 1),
    (1, 0),
    (1, -1),
]

# Movement probes per piece type
piece_probes = {
    piece_type["marshal"]: [1, 1, 1, 1, 1, 1, 1, 1],
    piece_type["general"]: [1, math.inf, 1, math.inf, math.inf, 1, math.inf, 1],
    piece_type["lieutenant_general"]: [math.inf, 1, math.inf, 1, 1, math.inf, 1, math.inf],
    piece_type["major_general"]: [1, 1, 1, 1, 1, 0, 1, 0],
    piece_type["warrior"]: [1, 1, 1, 0, 0, 0, 1, 0],
    piece_type["lancer"]: [1, (1, 2), 1, 0, 0, 0, 1, 0],
    piece_type["rider"]: [0, (1, 2), 0, 1, 1, 0, (1, 2), 0],
    piece_type["spy"]: [(1, 2), 0, (1, 2), 0, 0, (1, 2), 0, (1, 2)],
    piece_type["fortress"]: [0, 1, 0, 1, 1, 1, 0, 1],
    piece_type["soldier"]: [0, 1, 0, 0, 0, 0, 1, 0],
    piece_type["cannon"]: [0, 3, 0, 1, 1, 0, 1, 0],
    piece_type["archer"]: [2, 2, 2, 0, 0, 0, 1, 0],
    piece_type["musketeer"]: [0, 2, 0, 0, 0, 1, 0, 1],
    piece_type["tactician"]: [1, 0, 1, 0, 0, 0, 1, 0],
}


def _get_available_squares(
    direction: tuple[int, int],
    start: tuple[int, int],
    origin: tuple[int, int],
    length: float,
    board: Board,
) -> list[str]:
    y, x = start
    py, px = origin
    dy, dx = direction
    origin_piece = get_top(f"{py}-{px}", board)
    if origin_piece is None:
        return []

    if origin_piece.type == piece_type["archer"] and abs(dy) == 1 and abs(dx) == 1:
        wing = get_top(f"{py + dy}-{px + dx}", board)
        if wing and wing.tier > origin_piece.tier:
            return []

    available_squares: list[str] = []

    reverse_y, reverse_x = y, x
    while reverse_x != px or reverse_y != py:
        reverse_x -= dx
        reverse_y -= dy

        piece = get_top(f"{reverse_y}-{reverse_x}", board)
        if piece and piece.tier > origin_piece.tier:
            return []

        side = -1 if origin_piece.color == "b" else 1
        below = get_top(f"{reverse_y + side}-{reverse_x}", board)
        if below and below.square == f"{py}-{px}":
            break

    forward_y, forward_x = y, x
    step = 0
    while step < length:
        if forward_x < 1 or forward_x > 9 or forward_y < 1 or forward_y > 9:
            break

        piece = get_top(f"{forward_y}-{forward_x}", board)
        if piece and piece.tier > origin_piece.tier:
            break

        available_squares.append(f"{forward_y}-{forward_x}")

        leap_pieces = {
            piece_type["cannon"],
            piece_type["musketeer"],
            piece_type["archer"],
        }
        if piece:
            if origin_piece.type not in leap_pieces:
                break

            moving_forward = dy < 0 if origin_piece.color == "w" else dy > 0
            if not moving_forward:
                break

        forward_x += dx
        forward_y += dy
        step += 1

    return available_squares


def generate_moves_for_square(square: str, fen: str) -> list[Move]:
    state = parse_fen(fen)
    piece = get_top(square, state.board)
    if not piece or state.turn != piece.color:
        return []

    py, px = map(int, square.split("-"))
    probes = piece_probes[piece.type]
    squares: list[str] = []

    for idx, probe in enumerate(probes):
        pval = probe if isinstance(probe, (int, float)) else probe[0]
        pcarry = 1 if isinstance(probe, (int, float)) else probe[1]
        if pval < 1:
            continue

        dy, dx = dirs[idx]
        if piece.color == "b":
            dy *= -1
            dx *= -1

        target_y, target_x = (py + dy, px + dx) if pval is math.inf else (py + pval * dy, px + dx)
        length = math.inf if pval is math.inf else piece.tier + pcarry - 1
        squares.extend(
            _get_available_squares((dy, dx), (target_y, target_x), (py, px), length, state.board)
        )

    moves: list[Move] = []
    max_tier = 3 if state.mode == "advanced" else 2
    marshal_can_stack = state.mode in {"advanced", "intermediate"}

    for destination in squares:
        top_piece = get_top(destination, state.board)
        tower = get(destination, state.board)

        if not top_piece or not tower:
            moves.append(_create_move(piece, f"{destination}-1", fen, "route"))
            continue

        if top_piece.tier < max_tier and top_piece.type != piece_type["marshal"]:
            if piece.type != piece_type["marshal"] or marshal_can_stack:
                moves.append(
                    _create_move(piece, f"{destination}-{top_piece.tier + 1}", fen, "tsuke")
                )

            if piece.type == piece_type["tactician"] and any(
                tower_piece.color != piece.color for tower_piece in tower
            ):
                moves.extend(
                    _create_move(piece, f"{destination}-{top_piece.tier + 1}", fen, "betray", combo)
                    for combo in _get_betrayal_combos(tower, state.hand, piece.color)
                )

        if top_piece.color != piece.color:
            new_tier = (
                len([tower_piece for tower_piece in tower if tower_piece.color == piece.color]) + 1
            )
            captured = [tower_piece for tower_piece in tower if tower_piece.color != piece.color]
            moves.append(_create_move(piece, f"{destination}-{new_tier}", fen, "capture", captured))

    return moves


def generate_arata(hand_piece: HandPiece, fen: str) -> list[Move]:
    state = parse_fen(fen)
    if not hand_piece or state.turn != hand_piece.color:
        return []

    is_marshal_placed = not any(
        p for p in state.hand if p.type == piece_type["marshal"] and p.color == hand_piece.color
    )
    if not is_marshal_placed and hand_piece.type != piece_type["marshal"]:
        return []

    is_draft = state.drafting["b"] or state.drafting["w"]
    max_tier = 3 if state.mode == "advanced" else 2
    ranks: list[int] = []
    maybe: list[int] = []

    if is_draft:
        ranks = [7, 8, 9] if hand_piece.color == "w" else [1, 2, 3]
    else:
        start, end, step = (1, 9, 1) if hand_piece.color == "b" else (9, 1, -1)
        rank = start
        while (rank <= end) if hand_piece.color == "b" else (rank >= end):
            tops = [get_top(f"{rank}-{index + 1}", state.board) for index in range(9)]
            if any(piece and piece.color == hand_piece.color for piece in tops):
                ranks.extend(maybe)
                ranks.append(rank)
                maybe = []
            else:
                maybe.append(rank)
            rank += step

    squares = [f"{rank}-{file}" for rank in ranks for file in range(1, 10)]
    moves: list[Move] = []

    for destination in squares:
        top_piece = get_top(destination, state.board)
        arata_piece = Piece(
            type=hand_piece.type,
            color=hand_piece.color,
            square=destination,
            tier=0,
        )

        if not top_piece or (
            top_piece.color == arata_piece.color
            and top_piece.tier < max_tier
            and top_piece.type != piece_type["marshal"]
        ):
            tier = (top_piece.tier if top_piece else 0) + 1
            tower = get(destination, state.board)
            betrayal_combos = (
                _get_betrayal_combos(tower, state.hand, hand_piece.color)
                if hand_piece.type == piece_type["tactician"] and tower
                else []
            )
            player_hand_count = sum(hp.count for hp in state.hand if hp.color == hand_piece.color)
            is_last_piece = player_hand_count == 1

            if state.drafting[hand_piece.color]:
                if not is_last_piece:
                    moves.append(_create_move(arata_piece, f"{destination}-{tier}", fen, "arata"))
                    moves.extend(
                        _create_move(arata_piece, f"{destination}-{tier}", fen, "arata", combo)
                        for combo in betrayal_combos
                    )

                moves.append(
                    _create_move(arata_piece, f"{destination}-{tier}", fen, "arata", [], True)
                )
                moves.extend(
                    _create_move(arata_piece, f"{destination}-{tier}", fen, "arata", combo, True)
                    for combo in betrayal_combos
                )
            else:
                moves.append(_create_move(arata_piece, f"{destination}-{tier}", fen, "arata"))
                moves.extend(
                    _create_move(arata_piece, f"{destination}-{tier}", fen, "arata", combo)
                    for combo in betrayal_combos
                )

    return moves


def _create_move(
    piece: Piece,
    to: str,
    fen: str,
    move_type: MoveType,
    captured: list[Piece] | None = None,
    draft_finished: bool | None = None,
) -> Move:
    from_str = f"({piece.square}-{piece.tier})" if piece.tier != 0 else ""
    to_tier = int(to.split("-")[-1])
    arata = ARATA if move_type == "arata" else ""
    capture = TAKE if move_type == "capture" else ""
    tsuke = TSUKE if move_type == "tsuke" or (to_tier != 1 and to_tier - piece.tier > 0) else ""
    betray = (
        f"{BETRAY}{''.join(p.type for p in captured or [])}"
        if move_type == "betray" or (move_type == "arata" and captured and len(captured) > 0)
        else ""
    )
    draft_done = "\u7d42" if draft_finished else ""

    move = Move(
        piece=piece.type,
        color=piece.color,
        from_=from_str[1:-1] if from_str else "",
        to=to,
        type=move_type,
        san=f"{arata}{piece.type}{from_str}{capture}({to}){betray or tsuke}{draft_done}",
        before=fen,
        after="",
        draft_finished=draft_finished,
        captured=captured,
    )

    after_state, game_over_san = _make_move(move, move.before)
    move.after = after_state
    move.san += game_over_san
    return move


def _make_move(move: Move, fen: str) -> tuple[str, str]:
    state = parse_fen(fen)
    rank, file, tier = move.to.split("-")
    to_piece = Piece(
        type=move.piece,
        color=move.color,
        square=f"{rank}-{file}",
        tier=int(tier),
    )

    if move.type in {"route", "tsuke"}:
        if move.from_:
            remove_top(move.from_, state.board)
    elif move.type == "capture":
        if move.from_:
            remove_top(move.from_, state.board)
        remove(f"{rank}-{file}", move.captured or [], state.board)
    elif move.type == "betray":
        if move.from_:
            remove_top(move.from_, state.board)
        convert(f"{rank}-{file}", move.captured or [], state.board)
        update_hand(move.captured or [], state.hand, True)
    elif move.type == "arata":
        update_hand([to_piece], state.hand)
        if move.captured:
            convert(f"{rank}-{file}", move.captured, state.board)
            update_hand(move.captured, state.hand, True)
        if move.draft_finished:
            state.drafting[move.color] = False
            if move.color == "b":
                state.drafting["w"] = False

    put(to_piece, state.board)

    if state.turn == "b" and (state.drafting["w"] or not state.drafting["b"]):
        state.move_number += 1
    if (
        state.turn == "w"
        and not state.drafting["b"]
        and not state.drafting["w"]
        and move.draft_finished
    ):
        state.move_number += 1

    if state.drafting["w"] == state.drafting["b"]:
        if not move.draft_finished or state.turn != "w":
            state.turn = "w" if state.turn == "b" else "b"
    else:
        if not state.drafting["b"] and state.turn == "b":
            state.turn = "w"
        elif not state.drafting["w"] and state.turn == "w":
            state.turn = "b"

    game_over_san = ""
    if move.captured and any(piece.type == piece_type["marshal"] for piece in move.captured):
        game_over_san = "#"

    return encode_fen(state), game_over_san


def _generate_combinations(items: list[Piece]) -> list[list[Piece]]:
    result: list[list[Piece]] = []

    def helper(start: int, current: list[Piece]) -> None:
        if current:
            result.append(list(current))

        for index in range(start, len(items)):
            current.append(items[index])
            helper(index + 1, current)
            current.pop()

    helper(0, [])
    return result


def _get_betrayal_combos(
    tower: list[Piece], hand: list[HandPiece], color: Color
) -> list[list[Piece]]:
    enemies = [piece for piece in tower if piece.color != color]
    if not enemies:
        return []

    player_hand = [hand_piece for hand_piece in hand if hand_piece.color == color]
    enemy_count_map: dict[PieceType, int] = {}
    for enemy in enemies:
        enemy_count_map[enemy.type] = enemy_count_map.get(enemy.type, 0) + 1

    betrayal_options = [
        enemy
        for enemy_type, count in enemy_count_map.items()
        if any(
            hand_piece.type == enemy_type and hand_piece.count >= count
            for hand_piece in player_hand
        )
        for enemy in enemies
        if enemy.type == enemy_type
    ]

    return _generate_combinations(betrayal_options)


def _get_attacked_squares(square: str, board: Board) -> list[str]:
    piece = get_top(square, board)
    if piece is None:
        return []

    py, px = map(int, square.split("-"))
    probes = piece_probes[piece.type]
    attacked_squares: list[str] = []

    for idx, probe in enumerate(probes):
        pval = probe if isinstance(probe, (int, float)) else probe[0]
        pcarry = 1 if isinstance(probe, (int, float)) else probe[1]
        if pval < 1:
            continue

        dy, dx = dirs[idx]
        if piece.color == "b":
            dy *= -1
            dx *= -1

        target_y, target_x = (py + dy, px + dx) if pval is math.inf else (py + pval * dy, px + dx)
        length = math.inf if pval is math.inf else piece.tier + pcarry - 1
        attacked_squares.extend(
            _get_available_squares((dy, dx), (target_y, target_x), (py, px), length, board)
        )

    return attacked_squares


def is_square_attacked(square: str, by_color: Color, board: Board) -> bool:
    target_square = "-".join(square.split("-")[:2])

    for source_square in SQUARES:
        piece = get_top(source_square, board)
        if not piece or piece.color != by_color:
            continue

        attacked_squares = _get_attacked_squares(source_square, board)
        if any(attacked.startswith(target_square) for attacked in attacked_squares):
            return True

    return False


def in_check(color: Color, fen: str) -> bool:
    state = parse_fen(fen)
    opposite_color: Color = "b" if color == "w" else "w"

    for square in SQUARES:
        piece = get_top(square, state.board)
        if piece and piece.type == piece_type["marshal"] and piece.color == color:
            return is_square_attacked(square, opposite_color, state.board)

    return False


def would_be_in_check_after_move(move: Move) -> bool:
    return in_check(move.color, move.after)


# Re-export JS-style names
generateMovesForSquare = generate_moves_for_square
generateArata = generate_arata
isSquareAttacked = is_square_attacked
inCheck = in_check
wouldBeInCheckAfterMove = would_be_in_check_after_move
