from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING

from openpoker.ai.messages import AIMessageBuilder
from openpoker.ai.providers import ProviderClients
from openpoker.types import (
    ActionType,
    AIContext,
    AIDecisionDetails,
    AIProvider,
    AvailableAction,
    Player,
    PlayerType,
    PokerAction,
    TokenUsage,
)

if TYPE_CHECKING:
    from anthropic.types import Message as AnthropicMessage
    from openai.types.chat import ChatCompletion

    from openpoker.core.poker_game import PokerGame

logger = logging.getLogger(__name__)

_MEMORY_TOOL_NAME = "add_to_memory"
_MEMORY_TOOL_DESCRIPTION = (
    "Record observations about specific opponents to exploit their tendencies in future hands"
)
_MEMORY_TOOL_SCHEMA: dict[str, object] = {
    "type": "object",
    "properties": {
        "memory": {
            "type": "string",
            "description": "A brief note about a specific opponent's tendencies",
        }
    },
    "required": ["memory"],
}

MEMORY_TOOL_OPENAI: dict[str, object] = {
    "type": "function",
    "function": {
        "name": _MEMORY_TOOL_NAME,
        "description": _MEMORY_TOOL_DESCRIPTION,
        "parameters": _MEMORY_TOOL_SCHEMA,
    },
}

MEMORY_TOOL_ANTHROPIC: dict[str, object] = {
    "name": _MEMORY_TOOL_NAME,
    "description": _MEMORY_TOOL_DESCRIPTION,
    "input_schema": _MEMORY_TOOL_SCHEMA,
}


class AIAgent:
    def __init__(
        self,
        enable_memory: bool = True,
        enable_monte_carlo: bool = True,
        debug: bool = False,
    ):
        self.providers = ProviderClients.from_environment()
        self.message_builder = AIMessageBuilder(
            enable_memory=enable_memory,
            enable_monte_carlo=enable_monte_carlo,
        )
        self.hand_usage: list[TokenUsage] = []
        self.player_memories: dict[str, list[str]] = {}
        self.max_memories_per_player = 20
        self.enable_memory = enable_memory
        self.enable_monte_carlo = enable_monte_carlo
        self.debug = debug

    async def get_action(
        self,
        player: Player,
        context: AIContext,
        available_actions: list[AvailableAction],
        game: PokerGame,
    ) -> tuple[PokerAction | None, AIDecisionDetails | None]:
        if player.player_type != PlayerType.AI:
            return None, None
        if not player.provider or not player.model:
            return self._get_fallback_action(available_actions), None
        if not self.providers.has_provider(player.provider):
            return self._get_fallback_action(available_actions), None

        if player.provider in {AIProvider.OPENAI, AIProvider.OPENROUTER}:
            return await self._get_openai_family_action(
                player=player,
                context=context,
                available_actions=available_actions,
                game=game,
            )
        if player.provider == AIProvider.ANTHROPIC:
            return await self._get_anthropic_action(
                player=player,
                context=context,
                available_actions=available_actions,
                game=game,
            )
        return self._get_fallback_action(available_actions), None

    async def _get_openai_family_action(
        self,
        player: Player,
        context: AIContext,
        available_actions: list[AvailableAction],
        game: PokerGame,
    ) -> tuple[PokerAction | None, AIDecisionDetails | None]:
        try:
            client = (
                self.providers.async_openrouter_client
                if player.provider == AIProvider.OPENROUTER
                else self.providers.async_openai_client
            )
            if client is None:
                return self._get_fallback_action(available_actions), None

            filtered_tools = self._build_dynamic_tools_openai(available_actions)
            history = self._get_or_create_history(game, player.id)
            messages: list[dict[str, object]] = [
                {
                    "role": "system",
                    "content": self.message_builder.get_system_prompt(
                        player=player,
                        memories=self.player_memories.get(player.id, []),
                    ),
                }
            ]
            messages.extend(history)
            messages.append(
                {
                    "role": "user",
                    "content": self.message_builder.build_context_message(
                        context=context,
                        available_actions=available_actions,
                    ),
                }
            )

            response = await client.chat.completions.create(
                model=player.model,
                messages=messages,
                tools=filtered_tools,
                tool_choice="auto",
                parallel_tool_calls=False,
            )
            usage = self._extract_openai_usage(
                response=response,
                model=player.model,
                provider=player.provider,
            )
            self.hand_usage.append(usage)

            message = response.choices[0].message
            rationale = message.content or ""
            history.append(
                {
                    "role": "assistant",
                    "content": rationale,
                    "tool_calls": [
                        {
                            "id": tool_call.id,
                            "type": "function",
                            "function": {
                                "name": tool_call.function.name,
                                "arguments": tool_call.function.arguments,
                            },
                        }
                        for tool_call in (message.tool_calls or [])
                    ]
                    if message.tool_calls
                    else None,
                }
            )

            chosen_action: PokerAction | None = None
            chosen_tool_name: str | None = None
            chosen_args: dict[str, object] | None = None
            if message.tool_calls:
                for tool_call in message.tool_calls:
                    arguments = (
                        json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
                    )
                    parsed_action = self._parse_action(
                        function_name=tool_call.function.name,
                        arguments=arguments,
                        available_actions=available_actions,
                    )
                    if parsed_action:
                        chosen_action = parsed_action
                        chosen_tool_name = tool_call.function.name
                        chosen_args = arguments
                        history.append(
                            {
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": self._format_action_result(parsed_action),
                            }
                        )
                        break

            if not chosen_action:
                chosen_action = self._get_fallback_action(available_actions)

            details = AIDecisionDetails(
                player_id=player.id,
                player_name=player.name,
                provider=player.provider,
                model=player.model,
                rationale=rationale,
                tool_call=chosen_tool_name,
                arguments=chosen_args,
                token_usage=usage,
            )
            return chosen_action, details
        except Exception as error:
            logger.exception("openai-family decision failed for player=%s: %s", player.id, error)
            return self._get_fallback_action(available_actions), None

    async def _get_anthropic_action(
        self,
        player: Player,
        context: AIContext,
        available_actions: list[AvailableAction],
        game: PokerGame,
    ) -> tuple[PokerAction | None, AIDecisionDetails | None]:
        try:
            client = self.providers.async_anthropic_client
            if client is None:
                return self._get_fallback_action(available_actions), None

            filtered_tools = self._build_dynamic_tools_anthropic(available_actions)
            history = self._get_or_create_history(game, player.id)
            messages_to_send = list(history)
            messages_to_send.append(
                {
                    "role": "user",
                    "content": self.message_builder.build_context_message(
                        context=context,
                        available_actions=available_actions,
                    ),
                }
            )

            response = await client.messages.create(
                model=player.model,
                max_tokens=2048,
                system=self.message_builder.get_system_prompt(
                    player=player,
                    memories=self.player_memories.get(player.id, []),
                ),
                messages=messages_to_send,
                tools=filtered_tools,
            )
            usage = self._extract_anthropic_usage(response=response, model=player.model)
            self.hand_usage.append(usage)

            rationale_parts: list[str] = []
            content_blocks: list[dict[str, object]] = []
            chosen_action: PokerAction | None = None
            chosen_tool_name: str | None = None
            chosen_args: dict[str, object] | None = None
            chosen_tool_id: str | None = None

            for block in response.content:
                if block.type == "text":
                    rationale_parts.append(block.text)
                    content_blocks.append({"type": "text", "text": block.text})
                elif block.type == "tool_use":
                    arguments = block.input
                    parsed_action = self._parse_action(
                        function_name=block.name,
                        arguments=arguments,
                        available_actions=available_actions,
                    )
                    if parsed_action and not chosen_action:
                        chosen_action = parsed_action
                        chosen_tool_name = block.name
                        chosen_args = arguments
                        chosen_tool_id = block.id
                    content_blocks.append(
                        {
                            "type": "tool_use",
                            "id": block.id,
                            "name": block.name,
                            "input": block.input,
                        }
                    )

            rationale = " ".join(rationale_parts).strip()
            history.append({"role": "assistant", "content": content_blocks})

            if chosen_action and chosen_tool_id:
                history.append(
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "tool_result",
                                "tool_use_id": chosen_tool_id,
                                "content": self._format_action_result(chosen_action),
                            }
                        ],
                    }
                )
            if not chosen_action:
                chosen_action = self._get_fallback_action(available_actions)

            details = AIDecisionDetails(
                player_id=player.id,
                player_name=player.name,
                provider=AIProvider.ANTHROPIC,
                model=player.model,
                rationale=rationale,
                tool_call=chosen_tool_name,
                arguments=chosen_args,
                token_usage=usage,
            )
            return chosen_action, details
        except Exception as error:
            logger.exception("anthropic decision failed for player=%s: %s", player.id, error)
            return self._get_fallback_action(available_actions), None

    def _extract_openai_usage(
        self,
        response: ChatCompletion,
        model: str,
        provider: AIProvider,
    ) -> TokenUsage:
        usage_obj = getattr(response, "usage", None)
        prompt_tokens = int(getattr(usage_obj, "prompt_tokens", 0) or 0)
        completion_tokens = int(getattr(usage_obj, "completion_tokens", 0) or 0)
        total_tokens = int(getattr(usage_obj, "total_tokens", prompt_tokens + completion_tokens) or 0)

        prompt_details = getattr(usage_obj, "prompt_tokens_details", None)
        completion_details = getattr(usage_obj, "completion_tokens_details", None)
        cached_tokens = 0
        reasoning_tokens = None
        accepted_prediction_tokens = None
        rejected_prediction_tokens = None

        raw_usage: dict[str, object] = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
        }

        if prompt_details is not None:
            if hasattr(prompt_details, "cached_tokens"):
                cached_tokens = int(getattr(prompt_details, "cached_tokens", 0) or 0)
            raw_usage["prompt_tokens_details"] = self._to_plain_dict(prompt_details)
        if completion_details is not None:
            if hasattr(completion_details, "reasoning_tokens"):
                reasoning_tokens = int(getattr(completion_details, "reasoning_tokens", 0) or 0)
            if hasattr(completion_details, "accepted_prediction_tokens"):
                accepted_prediction_tokens = int(
                    getattr(completion_details, "accepted_prediction_tokens", 0) or 0
                )
            if hasattr(completion_details, "rejected_prediction_tokens"):
                rejected_prediction_tokens = int(
                    getattr(completion_details, "rejected_prediction_tokens", 0) or 0
                )
            raw_usage["completion_tokens_details"] = self._to_plain_dict(completion_details)

        return TokenUsage(
            provider=provider,
            model=model,
            input_tokens=prompt_tokens,
            output_tokens=completion_tokens,
            total_tokens=total_tokens,
            cache_read_input_tokens=cached_tokens if cached_tokens > 0 else None,
            reasoning_tokens=reasoning_tokens,
            accepted_prediction_tokens=accepted_prediction_tokens,
            rejected_prediction_tokens=rejected_prediction_tokens,
            raw_usage=raw_usage,
        )

    def _extract_anthropic_usage(self, response: AnthropicMessage, model: str) -> TokenUsage:
        usage_obj = getattr(response, "usage", None)
        input_tokens = int(getattr(usage_obj, "input_tokens", 0) or 0)
        output_tokens = int(getattr(usage_obj, "output_tokens", 0) or 0)
        cache_creation_input_tokens = int(getattr(usage_obj, "cache_creation_input_tokens", 0) or 0)
        cache_read_input_tokens = int(getattr(usage_obj, "cache_read_input_tokens", 0) or 0)
        total_tokens = input_tokens + output_tokens
        raw_usage = {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "cache_creation_input_tokens": cache_creation_input_tokens,
            "cache_read_input_tokens": cache_read_input_tokens,
        }

        return TokenUsage(
            provider=AIProvider.ANTHROPIC,
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=total_tokens,
            cache_creation_input_tokens=(
                cache_creation_input_tokens if cache_creation_input_tokens > 0 else None
            ),
            cache_read_input_tokens=cache_read_input_tokens if cache_read_input_tokens > 0 else None,
            raw_usage=raw_usage,
        )

    def _build_dynamic_tools_openai(
        self,
        available_actions: list[AvailableAction],
    ) -> list[dict[str, object]]:
        tools: list[dict[str, object]] = []
        for action in available_actions:
            spec = self._action_tool_spec(action)
            if spec is None:
                continue
            name, description, properties, required = spec
            tools.append(
                {
                    "type": "function",
                    "function": {
                        "name": name,
                        "description": description,
                        "parameters": {
                            "type": "object",
                            "properties": properties,
                            "required": required,
                        },
                    },
                }
            )
        return tools

    def _build_dynamic_tools_anthropic(
        self,
        available_actions: list[AvailableAction],
    ) -> list[dict[str, object]]:
        tools: list[dict[str, object]] = []
        for action in available_actions:
            spec = self._action_tool_spec(action)
            if spec is None:
                continue
            name, description, properties, required = spec
            tools.append(
                {
                    "name": name,
                    "description": description,
                    "input_schema": {
                        "type": "object",
                        "properties": properties,
                        "required": required,
                    },
                }
            )
        return tools

    def _action_tool_spec(
        self,
        action: AvailableAction,
    ) -> tuple[str, str, dict[str, object], list[str]] | None:
        if action.action_type == ActionType.FOLD:
            return ("fold", "Fold your hand and forfeit the pot", {}, [])
        if action.action_type == ActionType.CHECK:
            return ("check", "Check when no bet is required", {}, [])
        if action.action_type == ActionType.CALL:
            return ("call", f"Call the current bet of {action.min_amount} chips", {}, [])
        if action.action_type == ActionType.BET:
            return (
                "bet",
                f"Make a bet between {action.min_amount} and {action.max_amount} chips.",
                {"amount": {"type": "integer"}},
                ["amount"],
            )
        if action.action_type == ActionType.RAISE:
            return (
                "raise_bet",
                f"Raise to between {action.min_amount} and {action.max_amount} chips.",
                {"amount": {"type": "integer"}},
                ["amount"],
            )
        return None

    def _parse_action(
        self,
        function_name: str,
        arguments: dict[str, object],
        available_actions: list[AvailableAction],
    ) -> PokerAction | None:
        if function_name == "fold":
            return PokerAction(action_type=ActionType.FOLD)
        if function_name == "check":
            return PokerAction(action_type=ActionType.CHECK)
        if function_name == "call":
            return PokerAction(action_type=ActionType.CALL)
        if function_name == "bet":
            amount = self._parse_int(arguments.get("amount"))
            if amount is not None:
                return PokerAction(
                    action_type=ActionType.BET,
                    amount=self._clamp_amount(amount, available_actions, ActionType.BET),
                )
        if function_name == "raise_bet":
            amount = self._parse_int(arguments.get("amount"))
            if amount is not None:
                return PokerAction(
                    action_type=ActionType.RAISE,
                    amount=self._clamp_amount(amount, available_actions, ActionType.RAISE),
                )
        return None

    def _clamp_amount(
        self,
        amount: int,
        available_actions: list[AvailableAction],
        action_type: ActionType,
    ) -> int:
        clamped = amount
        for action in available_actions:
            if (
                action.action_type == action_type
                and action.min_amount is not None
                and action.max_amount is not None
            ):
                if clamped < action.min_amount:
                    clamped = action.min_amount
                elif clamped > action.max_amount:
                    clamped = action.max_amount
                break
        return clamped

    def _format_action_result(self, action: PokerAction) -> str:
        if action.action_type == ActionType.FOLD:
            return "You folded your hand."
        if action.action_type == ActionType.CHECK:
            return "You checked."
        if action.action_type == ActionType.CALL:
            return "You called."
        if action.action_type == ActionType.BET:
            return f"You bet ${action.amount}."
        if action.action_type == ActionType.RAISE:
            return f"You raised to ${action.amount}."
        return "Action executed."

    def _get_fallback_action(self, available_actions: list[AvailableAction]) -> PokerAction:
        for action in available_actions:
            if action.action_type == ActionType.CHECK:
                return PokerAction(action_type=ActionType.CHECK)
            if action.action_type == ActionType.CALL:
                return PokerAction(action_type=ActionType.CALL)
        return PokerAction(action_type=ActionType.FOLD)

    async def reflect_on_hand(self, player: Player, hand_summary: str) -> None:
        if not self.enable_memory:
            return
        if player.player_type != PlayerType.AI:
            return
        if not player.provider or not player.model:
            return

        if player.id not in self.player_memories:
            self.player_memories[player.id] = []
        current_memories = self.player_memories[player.id]
        memories_text = ""
        if current_memories:
            memories_text = (
                "Your current memories:\n"
                + "\n".join([f"- {memory}" for memory in current_memories])
                + "\n\n"
            )
        reflection_prompt = f"""{memories_text}Here is what happened in the hand that just finished:

{hand_summary}

Only record opponent-specific observations. Do not record generic strategy advice.
You can call add_to_memory multiple times (or not at all)."""

        try:
            if player.provider in {AIProvider.OPENAI, AIProvider.OPENROUTER}:
                client = (
                    self.providers.async_openrouter_client
                    if player.provider == AIProvider.OPENROUTER
                    else self.providers.async_openai_client
                )
                await self._reflect_openai_family(
                    player=player,
                    reflection_prompt=reflection_prompt,
                    client=client,
                )
            elif player.provider == AIProvider.ANTHROPIC:
                await self._reflect_anthropic(player=player, reflection_prompt=reflection_prompt)
        except Exception as error:
            logger.exception("reflection failed for player=%s: %s", player.id, error)

    async def _reflect_openai_family(
        self,
        player: Player,
        reflection_prompt: str,
        client: object | None,
    ) -> None:
        if client is None:
            return
        response = await client.chat.completions.create(
            model=player.model or "",
            messages=[
                {
                    "role": "system",
                    "content": f"You are {player.name}, reflecting on a poker hand you just played.",
                },
                {"role": "user", "content": reflection_prompt},
            ],
            tools=[MEMORY_TOOL_OPENAI],
            tool_choice="auto",
        )
        message = response.choices[0].message
        if message.tool_calls:
            for tool_call in message.tool_calls:
                if tool_call.function.name == "add_to_memory":
                    arguments = (
                        json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
                    )
                    memory = str(arguments.get("memory", "")).strip()
                    if memory:
                        self._add_memory(player.id, memory)

    async def _reflect_anthropic(self, player: Player, reflection_prompt: str) -> None:
        client = self.providers.async_anthropic_client
        if client is None:
            return
        response = await client.messages.create(
            model=player.model or "",
            max_tokens=1024,
            system=f"You are {player.name}, reflecting on a poker hand you just played.",
            messages=[{"role": "user", "content": reflection_prompt}],
            tools=[MEMORY_TOOL_ANTHROPIC],
        )
        for block in response.content:
            if block.type == "tool_use" and block.name == "add_to_memory":
                memory = str(block.input.get("memory", "")).strip()
                if memory:
                    self._add_memory(player.id, memory)

    def _add_memory(self, player_id: str, memory: str) -> None:
        if player_id not in self.player_memories:
            self.player_memories[player_id] = []
        self.player_memories[player_id].append(memory)
        if len(self.player_memories[player_id]) > self.max_memories_per_player:
            self.player_memories[player_id].pop(0)

    def get_hand_usage_summary(self) -> dict[str, object]:
        if not self.hand_usage:
            return {
                "total_input_tokens": 0,
                "total_output_tokens": 0,
                "total_tokens": 0,
                "calls": [],
                "by_model": {},
                "by_provider": {},
            }

        total_input = sum(call.input_tokens for call in self.hand_usage)
        total_output = sum(call.output_tokens for call in self.hand_usage)
        total_tokens = sum(call.total_tokens for call in self.hand_usage)

        by_model: dict[str, dict[str, int]] = {}
        by_provider: dict[str, dict[str, int]] = {}
        for call in self.hand_usage:
            if call.model not in by_model:
                by_model[call.model] = {
                    "calls": 0,
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "total_tokens": 0,
                }
            by_model[call.model]["calls"] += 1
            by_model[call.model]["input_tokens"] += call.input_tokens
            by_model[call.model]["output_tokens"] += call.output_tokens
            by_model[call.model]["total_tokens"] += call.total_tokens

            provider_key = call.provider.value
            if provider_key not in by_provider:
                by_provider[provider_key] = {
                    "calls": 0,
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "total_tokens": 0,
                }
            by_provider[provider_key]["calls"] += 1
            by_provider[provider_key]["input_tokens"] += call.input_tokens
            by_provider[provider_key]["output_tokens"] += call.output_tokens
            by_provider[provider_key]["total_tokens"] += call.total_tokens

        return {
            "total_input_tokens": total_input,
            "total_output_tokens": total_output,
            "total_tokens": total_tokens,
            "calls": [call.model_dump() for call in self.hand_usage],
            "by_model": by_model,
            "by_provider": by_provider,
        }

    def reset_hand_usage(self) -> None:
        self.hand_usage = []

    def _parse_int(self, value: object) -> int | None:
        if isinstance(value, int):
            return value
        if isinstance(value, float):
            return int(value)
        if isinstance(value, str):
            stripped = value.strip().replace("$", "").replace(",", "")
            try:
                return int(float(stripped))
            except (ValueError, OverflowError):
                return None
        return None

    def _to_plain_dict(self, value: object) -> dict[str, object]:
        if hasattr(value, "model_dump"):
            dumped = value.model_dump()
            if isinstance(dumped, dict):
                return dumped
        if hasattr(value, "__dict__"):
            raw = value.__dict__
            if isinstance(raw, dict):
                return {str(key): raw[key] for key in raw}
        return {"value": str(value)}

    def _get_or_create_history(
        self,
        game: PokerGame,
        player_id: str,
    ) -> list[dict[str, object]]:
        if player_id not in game.conversation_history:
            game.conversation_history[player_id] = []
        return game.conversation_history[player_id]
