import logging
import re
from typing import Any

from xent.common.configuration_types import (
    GameMapRoundResult,
    PlayerName,
    is_omniscient_player_name,
)
from xent.common.constants import ZERO_SUM_PLAYER_PAIRS
from xent.common.errors import XentGameError, XentInternalError, XentSyntaxError
from xent.common.token_xent_list import TokenXentList, ValidatedBool
from xent.common.x_flag import XFlag
from xent.common.x_list import XList
from xent.common.x_string import XString
from xent.common.xent_event import (
    ElicitRequestEvent,
    ElicitResponseEvent,
    FailedEnsureEvent,
    RevealEvent,
    RewardEvent,
    TokenUsage,
    XentEvent,
    deserialize_event,
    serialize_event,
)
from xent.runtime.players.base_player import XGP
from xent.runtime.players.players import deserialize_player

# This flag enforces that ensure arguments are ValidatedBools generated by comparisons of xents.
ENFORCE_XENT_ENSURE = False

MAX_ENSURE_FAILURES = 10


class XentRuntime:
    def __init__(
        self,
        player: XGP,
        npcs: list[XGP],
        locals: dict[str, Any],
        globals: dict[str, Any],
        store_full_interactions: bool = False,
    ):
        self.local_vars: dict[str, Any] = locals
        self.player = player
        self.player.reset_score()
        self.npcs = npcs
        self.reset_npcs()
        self.score = 0.0
        self.token_usage: TokenUsage = {"input_tokens": 0, "output_tokens": 0}
        self.globals = globals
        self.beacons: dict[str, XFlag] = {}
        self.history: list[XentEvent] = []
        # Map from line number to the number of times the replay has been called since reset
        self.replay_counters: dict[int, int] = {}
        self.store_full_interactions = store_full_interactions
        self.last_elicit_player: XGP | None = None
        self.revealed_registers: dict[PlayerName, list[str]] = {}

    def serialize(self) -> dict[str, Any]:
        locals: dict[str, Any] = {}
        for name, local in self.local_vars.items():
            if isinstance(local, XString | XList):
                locals[name] = local.serialize()

        beacons: dict[str, Any] = {}
        for name, beacon in self.beacons.items():
            beacons[name] = beacon.serialize()

        last_elicit_player_id = (
            self.last_elicit_player.id if self.last_elicit_player is not None else None
        )
        serialized = {
            "player": self.player.serialize(),
            "npcs": [npc.serialize() for npc in self.npcs],
            "locals": locals,
            "beacons": beacons,
            "score": self.score,
            "token_usage": self.token_usage,
            "history": [serialize_event(event) for event in self.history],
            "replay_counters": self.replay_counters,
            "store_full_interactions": self.store_full_interactions,
            "last_elicit_player_id": last_elicit_player_id,
            "revealed_registers": self.revealed_registers,
        }
        return serialized

    @classmethod
    def deserialize(
        cls,
        data: dict[str, Any],
        globals: dict[str, Any],
    ) -> "XentRuntime":
        def _deserialize_local(val: dict[str, Any]):
            t = val.get("type")
            if t == "XString":
                return XString.deserialize(val)
            if t == "XList":
                return XList.deserialize(val)
            return val

        self = cls.__new__(cls)

        self.player = deserialize_player(data["player"])
        self.npcs = [deserialize_player(x) for x in data["npcs"]]

        self.local_vars = {
            name: _deserialize_local(val) for name, val in data["locals"].items()
        }
        self.local_vars[self.player.name] = self.player
        for npc in self.npcs:
            self.local_vars[npc.name] = npc

        self.globals = globals

        self.score = float(data["score"])
        self.token_usage = {
            "input_tokens": int(data["token_usage"]["input_tokens"]),
            "output_tokens": int(data["token_usage"]["output_tokens"]),
        }
        self.history = [deserialize_event(ev) for ev in data["history"]]

        # Note: JSON keys might be strings; ensure int keys for replay counters
        self.replay_counters = {
            int(k): int(v) for k, v in data["replay_counters"].items()
        }

        self.store_full_interactions = bool(data.get("store_full_interactions", False))

        self.beacons = {
            name: XFlag.deserialize(b) for name, b in data["beacons"].items()
        }

        self.last_elicit_player = None
        last_id = data.get("last_elicit_player_id")
        if last_id is not None:
            if self.player.id == last_id:
                self.last_elicit_player = self.player
            else:
                for npc in self.npcs:
                    if npc.id == last_id:
                        self.last_elicit_player = npc
                        break

        self.revealed_registers = data["revealed_registers"]

        return self

    def reset_npcs(self):
        for npc in self.npcs:
            npc.reset_score()

    def add_token_usage(self, token_usage: TokenUsage) -> None:
        self.token_usage["input_tokens"] += token_usage["input_tokens"]
        self.token_usage["output_tokens"] += token_usage["output_tokens"]

    def instruction_names(self) -> set[str]:
        return {"assign", "elicit", "reveal", "reward", "ensure", "beacon", "replay"}

    def _reset_register_states(self):
        for var_name, var in self.local_vars.items():
            if isinstance(var, XString):
                self.local_vars[var_name] = XString(
                    "", static=var.static, public=var.public, name=var.name
                )
            elif isinstance(var, XList):
                self.local_vars[var_name] = XList(
                    [], static=var.static, public=var.public, name=var.name
                )

    def get_results_and_reset(self) -> GameMapRoundResult:
        score = self.score
        self.score = 0.0
        token_usage = self.token_usage
        self.token_usage = {"input_tokens": 0, "output_tokens": 0}

        game_result = GameMapRoundResult(
            score=score, history=self.history, token_usage=token_usage
        )
        self.history = []
        self.replay_counters = {}
        self.beacons = {}
        self._reset_register_states()
        self.reset_npcs()
        self.revealed_registers = {}
        return game_result

    async def execute(
        self,
        instruction_name: str,
        args: list[Any],
        kwargs: dict[str, Any],
        line: str,
        line_num: int,
    ) -> None | XFlag:
        try:
            if instruction_name == "assign":
                return self.assign(args, kwargs)
            elif instruction_name == "elicit":
                return await self.elicit(args, kwargs, line_num, line)
            elif instruction_name == "reveal":
                return await self.reveal(args, kwargs, line_num, line)
            elif instruction_name == "reward":
                return await self.reward(args, kwargs, line_num, line)
            elif instruction_name == "ensure":
                return await self.ensure(args, kwargs, line_num, line)
            elif instruction_name == "beacon":
                return self.beacon(args, kwargs, line_num)
            elif instruction_name == "replay":
                return self.replay(args, kwargs, line_num)
            else:
                raise XentSyntaxError(f"Unknown instruction: {instruction_name}")
        except XentSyntaxError as e:
            logging.exception(f"Syntax error in instruction: {e}")
            raise XentSyntaxError(
                f"{e}. Line: {line}, Line number: {line_num}"
            ) from None
        except XentInternalError as e:
            logging.exception(f"Internal error in instruction: {e}")
            raise XentInternalError(
                f"{e}. Line: {line}, Line number: {line_num}"
            ) from e

    def _first_n_tokens(self, text: str, n: int) -> str | XString:
        return self.globals["first_n_tokens"](text, n)

    def assert_no_args(self, args: list[Any], instruction_name: str) -> None:
        if len(args) > 0:
            raise XentSyntaxError(
                f"Positional arguments are not allowed for {instruction_name}"
            )

    def assert_no_kwargs(self, kwargs: dict[str, Any], instruction_name: str) -> None:
        if len(kwargs) > 0:
            raise XentSyntaxError(
                f"Keyword arguments are not allowed for {instruction_name}"
            )

    async def send_event(self, player: XGP, event: XentEvent) -> None:
        await player.post(event)
        self.history.append(event)

    def _validate_assign_register(self, var_name: str) -> XString | XList:
        cur = self.local_vars.get(var_name)
        if cur is None:
            raise XentSyntaxError(f"Register {var_name} does not exist")
        if not isinstance(cur, XString) and not isinstance(cur, XList):
            raise XentSyntaxError(
                f"{var_name} is not a valid assign target. Assign requires a String or List register"
            )
        if cur.static:
            raise XentSyntaxError(
                f"Register {var_name} is static: {cur}. Assign requires a non-static register"
            )

        return cur

    def _validate_assign_arg(self, var_name: str, var_value: Any) -> XString | XList:
        if isinstance(var_value, str):
            var_value = XString(var_value)

        if isinstance(var_value, XString | XList):
            return var_value
        raise XentSyntaxError(
            f"Value provided to assign is not a String or List: {var_value}. Assign requires a String or List value"
        )

    def assign(self, args: list[Any], kwargs: dict[str, Any]) -> None:
        self.assert_no_args(args, "assign")
        for var_name, var_value in kwargs.items():
            logging.info(f"Assigning {var_value} to {var_name}")
            cur = self._validate_assign_register(var_name)
            validated_value = self._validate_assign_arg(var_name, var_value)
            if isinstance(cur, XString) and isinstance(validated_value, XString):
                cur.primary_string = validated_value.primary_string
                cur.prefix = validated_value.prefix
            elif isinstance(cur, XList) and isinstance(validated_value, XList):
                cur.items = validated_value.items
            else:
                target_type = "String" if isinstance(cur, XString) else "List"
                value_type = "String" if target_type == "List" else "String"
                raise XentSyntaxError(
                    f"Attempted to assign incompatible value to register. Register type: {target_type}, value type: {value_type}"
                )
        return None

    def _validate_elicit_args(
        self, args: list[Any], line_num: int
    ) -> tuple[int, int, XGP, int]:
        if len(args) < 2:
            raise XentSyntaxError("Elicit requires at least two positional arguments")
        if len(args) < 2:
            raise XentSyntaxError("Elicit requires at least two positional arguments")

        var_arg_start = 1
        player = args[0]
        if not isinstance(player, XGP):
            var_arg_start = 0
            player = self.player
        if not isinstance(player, XGP):
            raise XentInternalError("No player identified for elicit")

        max_len = args[-1]
        if not isinstance(max_len, int):
            raise XentSyntaxError(
                f"Elicit requires an integer for the length, not {type(max_len)}"
            )
        return (
            var_arg_start,
            len(args) - 1,  # Exclude the last argument which is max_len
            player,
            max_len,
        )

    def _validate_elicit_arg(self, var: Any) -> None:
        if not isinstance(var, XString):
            raise XentSyntaxError("Elicit requires a String register target")
        if var.name is None:
            raise XentSyntaxError("Elicit requires a String register target")
        if var.static:
            raise XentSyntaxError(
                "Elicit target is static. Elicit requires a non-static register target"
            )

    def _gather_register_states(
        self, player_name: PlayerName
    ) -> dict[str, XString | XList]:
        register_states = {}
        for var_name, var in self.local_vars.items():
            if not isinstance(var, XString | XList):
                continue

            if (
                var.public
                or is_omniscient_player_name(player_name)
                or var_name in self.revealed_registers.get(player_name, [])
            ):
                register_states[var_name] = var

        logging.debug(
            f"Gathered register states for player {player_name}: {register_states}"
        )
        return register_states

    async def elicit(
        self, args: list[Any], kwargs: dict[str, Any], line_num: int, line: str
    ) -> None:
        self.assert_no_kwargs(kwargs, "elicit")
        var_arg_start, var_arg_end, player, max_len = self._validate_elicit_args(
            args, line_num
        )
        for i in range(var_arg_start, var_arg_end):
            var = args[i]
            self._validate_elicit_arg(var)
            register_states = self._gather_register_states(player.name)
            request_event = ElicitRequestEvent(
                type="elicit_request",
                line=line,
                line_num=line_num,
                player=player.name,
                var_name=var.name,
                max_len=max_len,
                registers=register_states,
            )
            await self.send_event(player, request_event)

            move_result = await player.make_move(var.name, register_states)
            player_move = move_result.response
            token_usage = move_result.token_usage
            prompts = move_result.prompts
            full_response = move_result.full_response

            print(f"Player {player.name} move: {player_move}")
            trimmed_move = self._first_n_tokens(player_move, max_len)
            if player_move != trimmed_move:
                logging.warning(
                    f"Player {player.name} move trimmed from {player_move} to {trimmed_move}"
                )

            response_event = ElicitResponseEvent(
                type="elicit_response",
                line=line,
                line_num=line_num,
                player=player.name,
                response=str(trimmed_move),
                token_usage=token_usage,
            )
            if self.store_full_interactions:
                response_event["prompts"] = prompts
                response_event["full_response"] = full_response

            await self.send_event(player, response_event)
            if player.name == self.player.name:
                self.add_token_usage(token_usage)

            logging.info(f'Setting {var.name} to "{trimmed_move}"')
            var.primary_string = trimmed_move
            var.prefix = ""

        self.beacons["previous_elicit"] = XFlag("previous_elicit", line_num)
        self.last_elicit_player = player
        return None

    async def reveal(
        self, args: list[Any], kwargs: dict[str, Any], line_num: int, line: str
    ) -> None:
        self.assert_no_kwargs(kwargs, "reveal")
        if len(args) == 0:
            raise XentSyntaxError(
                "Reveal requires at least one positional argument, got none"
            )

        arg_start = 1
        player = args[0]
        if not isinstance(player, XGP):
            player = self.player
            arg_start = 0
        if not isinstance(player, XGP):
            raise XentInternalError("No player identified for reward")

        var_names = extract_reveal_parameters(line)[arg_start:]
        rest_of_args = args[arg_start:]
        reveal_event = RevealEvent(
            type="reveal",
            line=line,
            line_num=line_num,
            player=player.name,
            values={name: rest_of_args[i] for i, name in enumerate(var_names)},
        )
        await self.send_event(player, reveal_event)
        if player.name not in self.revealed_registers:
            self.revealed_registers[player.name] = []

        for var_name in var_names:
            self.revealed_registers[player.name].append(var_name)

        logging.info(f"Revealed {reveal_event} to player {player.name}")
        return None

    async def reward(
        self, args: list[Any], kwargs: dict[str, Any], line_num: int, line: str
    ) -> None:
        self.assert_no_kwargs(kwargs, "reward")
        if len(args) == 0:
            raise XentSyntaxError(
                "Reward requires one or two positional arguments, got none"
            )
        if len(args) > 2:
            raise XentSyntaxError(
                f"Reward requires one or two positional arguments, got {len(args)}"
            )

        player = args[0]
        if isinstance(player, XGP):
            score = args[1]
        else:
            score = args[0]
            player = self.player
            if not isinstance(player, XGP):
                raise XentInternalError("No player identified for reward")

        if not isinstance(score, TokenXentList):
            raise XentSyntaxError(
                "Reward amount must generated by linearly combining xents"
            )

        player.add_score(score.total_xent())
        if player.name == self.player.name:
            self.score += score.total_xent()
        reward_event = RewardEvent(
            type="reward",
            line=line,
            line_num=line_num,
            player=player.name,
            value=score,
        )
        await self.send_event(player, reward_event)
        logging.info(
            f"Rewarded player {player.name} with {score.total_xent()}. Reward event: {reward_event}"
        )

        # Handle zero-sum player pairs
        for pair in ZERO_SUM_PLAYER_PAIRS:
            if player.name in pair:
                other_name = pair[1] if player.name == pair[0] else pair[0]

                other_player = self.local_vars.get(other_name)

                if isinstance(other_player, XGP):
                    neg_score = -1 * score
                    other_player.add_score(neg_score.total_xent())
                    if other_player.name == self.player.name:
                        self.score += neg_score.total_xent()

                    reward_event = RewardEvent(
                        type="reward",
                        line=line,
                        line_num=line_num,
                        player=other_player.name,
                        value=neg_score,
                    )
                    await self.send_event(other_player, reward_event)
                    logging.info(
                        f"Rewarded player {other_player.name} with {neg_score.total_xent()}. Reward event: {reward_event}"
                    )
        return None

    def _validate_ensure_args(self, args: list[Any]) -> None:
        if len(args) == 0:
            raise XentSyntaxError(
                "Ensure requires at least one positional argument, got none"
            )

        for i, arg in enumerate(args):
            if ENFORCE_XENT_ENSURE:
                if not isinstance(arg, ValidatedBool):
                    raise XentSyntaxError(
                        f"Argument {i} is not the result of a xent computation. Ensure requires a boolean generated by comparisons of xents"
                    )
            else:
                if not isinstance(arg, bool) and not isinstance(arg, ValidatedBool):
                    raise XentSyntaxError(
                        f"Argument {i} is not a boolean. Ensure requires boolean arguments"
                    )

    def _validate_elicit_occurred(self) -> XGP:
        if self.last_elicit_player is None:
            raise XentSyntaxError(
                "No player identified for previous elicit. Cannot ensure before elicit"
            )
        return self.last_elicit_player

    async def ensure(
        self, args: list[Any], kwargs: dict[str, Any], line_num: int, line: str
    ) -> None | XFlag:
        self.assert_no_kwargs(kwargs, "ensure")
        self._validate_ensure_args(args)
        last_elicit_player = self._validate_elicit_occurred()

        if all(args):
            self.replay_counters[line_num] = 0
            return None

        fail_count = self.replay_counters.get(line_num, 0)
        if fail_count >= MAX_ENSURE_FAILURES:
            raise XentGameError(
                f"Ensure limit exceeded for line {line_num}. Limit: {MAX_ENSURE_FAILURES}"
            )
        previous_elicit = self.beacons["previous_elicit"]

        ensure_event = FailedEnsureEvent(
            type="failed_ensure",
            line=line,
            line_num=line_num,
            player=last_elicit_player.name,
            ensure_results=args,
            beacon=previous_elicit.name,
        )
        await self.send_event(last_elicit_player, ensure_event)
        self.replay_counters[line_num] = fail_count + 1
        return previous_elicit

    def beacon(self, args: list[Any], kwargs: dict[str, Any], line_num: int) -> None:
        self.assert_no_kwargs(kwargs, "beacon")
        if len(args) != 1:
            raise XentSyntaxError(
                f"Beacon requires exactly one positional argument, got {len(args)}"
            )

        flag = args[0]
        if not isinstance(flag, XFlag):
            raise XentSyntaxError("Beacon argument must be either `flag_1` or `flag_2`")

        flag.line_num = line_num
        self.beacons[flag.name] = flag
        return None

    def replay(
        self, args: list[Any], kwargs: dict[str, Any], line_num: int
    ) -> XFlag | None:
        self.assert_no_kwargs(kwargs, "replay")

        if len(args) != 2:
            raise XentSyntaxError(
                f"Replay requires exactly two positional arguments, got {len(args)}"
            )

        flag = args[0]
        if not isinstance(flag, XFlag):
            raise XentSyntaxError("Replay argument must be an XFlag")

        if flag.name not in self.beacons:
            raise XentInternalError(f"Replay beacon {flag.name} not found")

        iterations = args[1]
        if not isinstance(iterations, int):
            raise XentSyntaxError("Replay iterations must be an integer")

        replay_count = self.replay_counters.get(line_num, 0)
        if replay_count >= iterations:
            logging.info(
                f"Replay limit exceeded for {flag.name}. Limit: {iterations}, Current: {replay_count}. Continuing execution"
            )
            self.replay_counters[line_num] = 0
            return None

        self.replay_counters[line_num] = replay_count + 1
        return flag


def extract_reveal_parameters(input_str: str) -> list[str]:
    match = re.search(r"\(([^)]+)\)", input_str)
    if not match:
        return []
    params = [param.strip() for param in match.group(1).split(",")]
    return params
