"""Base Agent and Playback classes for ARC-AGI-3."""

from __future__ import annotations

import json
import logging
import os
import time
from abc import ABC, abstractmethod
from copy import deepcopy
from typing import TYPE_CHECKING, Any

import requests
import requests.cookies
from pydantic import ValidationError
from requests.cookies import RequestsCookieJar

from ._recorder import Recorder
from ._structs import FrameData, GameAction, GameState, Scorecard
from ._tracing import trace_agent_session

if TYPE_CHECKING:
    from requests import Response

logger = logging.getLogger(__name__)


class Agent(ABC):
    """Abstract base class for an agent that plays one ARC-AGI-3 game."""

    MAX_ACTIONS: int = 80  # Avoid looping forever if agent doesn't exit
    ROOT_URL: str

    action_counter: int = 0

    timer: float = 0
    agent_name: str
    card_id: str
    game_id: str
    guid: str
    frames: list[FrameData]

    recorder: Recorder
    headers: dict[str, str]
    _session: requests.Session

    # AgentOps tracing attributes
    trace: Any = None
    tags: list[str]

    def __init__(
        self,
        card_id: str,
        game_id: str,
        agent_name: str,
        ROOT_URL: str,
        record: bool,
        tags: list[str] | None = None,
        cookies: RequestsCookieJar | None = None,
    ) -> None:
        """Initialize an agent.

        Args:
            card_id: The scorecard ID
            game_id: The game ID to play
            agent_name: Name of this agent
            ROOT_URL: Base URL for the API
            record: Whether to record the session
            tags: Optional tags for tracing
            cookies: Optional cookies for the session
        """
        self.ROOT_URL = ROOT_URL
        self.card_id = card_id
        self.game_id = game_id
        self.guid = ""
        self.agent_name = agent_name
        self.tags = tags or []
        self.frames = [FrameData(score=0)]
        self._cleanup = True
        if record:
            self.start_recording()
        self.headers = {
            "X-API-Key": os.getenv("ARC_API_KEY", ""),
            "Accept": "application/json",
        }
        # Reuse session
        self._session = requests.Session()
        self._session.cookies = deepcopy(cookies) if cookies else RequestsCookieJar()
        self._session.headers.update(self.headers)

    @trace_agent_session
    def main(self) -> None:
        """The main agent loop. Play the game until finished, then exit."""
        self.timer = time.time()
        while (
            not self.is_done(self.frames, self.frames[-1])
            and self.action_counter <= self.MAX_ACTIONS
        ):
            action = self.choose_action(self.frames, self.frames[-1])
            if frame := self.take_action(action):
                self.append_frame(frame)
                logger.info(
                    f"{self.game_id} - {action.name}: count {self.action_counter}, "
                    f"score {frame.score}, avg fps {self.fps})"
                )
            self.action_counter += 1

        self.cleanup()

    @property
    def state(self) -> GameState:
        """Current game state."""
        return self.frames[-1].state

    @property
    def score(self) -> int:
        """Current game score."""
        return self.frames[-1].score

    @property
    def seconds(self) -> float:
        """Elapsed time in seconds."""
        return (time.time() - self.timer) * 100 // 1 / 100

    @property
    def fps(self) -> float:
        """Average actions per second."""
        if self.action_counter == 0:
            return 0.0
        elapsed_time = max(self.seconds, 0.1)
        return round(self.action_counter / elapsed_time, 2)

    @property
    def is_playback(self) -> bool:
        """Check if this is a playback agent."""
        return type(self) is Playback

    @property
    def name(self) -> str:
        """Generate standardized agent name."""
        n = self.__class__.__name__.lower()
        return f"{self.game_id}.{n}"

    def start_recording(self) -> None:
        """Initialize recording session."""
        filename = self.agent_name if self.is_playback else None
        self.recorder = Recorder(prefix=self.name, filename=filename)
        logger.info(f"created new recording for {self.name} into {self.recorder.filename}")

    def append_frame(self, frame: FrameData) -> None:
        """Add frame data and record if applicable."""
        self.frames.append(frame)
        if frame.guid:
            self.guid = frame.guid
        if hasattr(self, "recorder") and not self.is_playback:
            self.recorder.record(json.loads(frame.model_dump_json()))

    def do_action_request(self, action: GameAction) -> Response:
        """Submit an action via HTTP POST to the API."""
        data = action.action_data.model_dump()
        if action == GameAction.RESET:
            data["card_id"] = self.card_id
        if self.guid:
            data["guid"] = self.guid
        if action.reasoning:
            data["reasoning"] = action.reasoning
        if self.game_id:
            data["game_id"] = self.game_id

        json_str = json.dumps(data)
        r = self._session.post(
            f"{self.ROOT_URL}/api/cmd/{action.name}",
            json=json.loads(json_str),
            headers=self.headers,
        )
        if "error" in r.json():
            logger.warning(f"Exception during action request: {r.json()}")
        return r

    def take_action(self, action: GameAction) -> FrameData | None:
        """Submit the action and get the next frame."""
        frame_data = self.do_action_request(action).json()
        try:
            frame = FrameData.model_validate(frame_data)
        except ValidationError as e:
            logger.warning(f"Incoming frame data did not validate: {e}")
            return None
        return frame

    def get_scorecard(self) -> Scorecard:
        """Get the scorecard for this agent's game."""
        r = self._session.get(
            f"{self.ROOT_URL}/api/scorecard/{self.card_id}/{self.game_id}",
            timeout=1,
            headers=self.headers,
        )
        response_data = r.json()
        if "error" in response_data:
            logger.warning(f"Exception during scorecard request: {response_data}")
        return Scorecard.model_validate(response_data)

    def cleanup(self, scorecard: Scorecard | None = None) -> None:
        """Called after main loop is finished."""
        if self._cleanup:
            self._cleanup = False  # Only cleanup once per agent
            if hasattr(self, "recorder") and not self.is_playback:
                if scorecard:
                    self.recorder.record(scorecard.get(self.game_id))
                else:
                    scorecard_obj = self.get_scorecard()
                    self.recorder.record(scorecard_obj.get(self.game_id))
                logger.info(
                    f"recording for {self.name} is available in {self.recorder.filename}"
                )
            if self.action_counter >= self.MAX_ACTIONS:
                logger.info(
                    f"Exiting: agent reached MAX_ACTIONS of {self.MAX_ACTIONS}, "
                    f"took {self.seconds} seconds ({self.fps} average fps)"
                )
            else:
                logger.info(
                    f"Finishing: agent took {self.action_counter} actions, "
                    f"took {self.seconds} seconds ({self.fps} average fps)"
                )
            if hasattr(self, "_session"):
                self._session.close()

    @abstractmethod
    def is_done(self, frames: list[FrameData], latest_frame: FrameData) -> bool:
        """Decide if the agent is done playing.

        Args:
            frames: All frames collected so far
            latest_frame: The most recent frame

        Returns:
            True if the agent should stop playing
        """
        raise NotImplementedError

    @abstractmethod
    def choose_action(
        self,
        frames: list[FrameData],
        latest_frame: FrameData,
    ) -> GameAction:
        """Choose which action the Agent should take.

        Args:
            frames: All frames collected so far
            latest_frame: The most recent frame

        Returns:
            The action to take, with any arguments filled in
        """
        raise NotImplementedError


class Playback(Agent):
    """An agent that plays back from a recorded session from another agent."""

    MAX_ACTIONS = 1000000
    PLAYBACK_FPS = 5

    recorded_actions: list[dict[str, Any]]

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.recorder = Recorder(
            prefix=Recorder.get_prefix(self.agent_name),
            guid=Recorder.get_guid(self.agent_name),
        )
        self.recorded_actions = []
        if self.agent_name in Recorder.list():
            try:
                self.recorded_actions = self.filter_actions()
                logger.info(
                    f"Loaded {len(self.recorded_actions)} actions from {self.agent_name}"
                )
            except Exception as e:
                logger.error(f"Failed to load recording {self.agent_name}: {e}")
                self.recorded_actions = []
        else:
            logger.warning(
                f"Recording {self.agent_name} not found in available recordings"
            )

    def filter_actions(self) -> list[dict[str, Any]]:
        """Filter recorded events to only include action frames."""
        return [
            a for a in self.recorder.get() if "data" in a and "action_input" in a["data"]
        ]

    def is_done(self, frames: list[FrameData], latest_frame: FrameData) -> bool:
        """Check if all recorded actions have been played."""
        return bool(self.action_counter >= len(self.recorded_actions))

    def choose_action(
        self,
        frames: list[FrameData],
        latest_frame: FrameData,
    ) -> GameAction:
        """Get the next action from the recording."""
        loop_start_time = time.time()

        if self.action_counter >= len(self.recorded_actions):
            logger.warning(
                f"No more recorded actions available "
                f"(counter: {self.action_counter}, total: {len(self.recorded_actions)})"
            )
            return GameAction.RESET

        recorded_data = self.recorded_actions[self.action_counter]["data"]
        action_input = recorded_data["action_input"]

        action = GameAction.from_id(action_input["id"])
        data = action_input["data"].copy()
        data["game_id"] = self.game_id
        action.set_data(data)
        if "reasoning" in action_input and action_input["reasoning"] is not None:
            action.reasoning = action_input["reasoning"]

        logger.debug(f"Playback action {self.action_counter}: {action.name} with data {data}")

        # Rate limit playback
        target_frame_time = 1.0 / getattr(self, "PLAYBACK_FPS", 5)
        elapsed_time = time.time() - loop_start_time
        sleep_time = max(0, target_frame_time - elapsed_time)
        if sleep_time > 0:
            time.sleep(sleep_time)

        return action

    def append_frame(self, frame: FrameData) -> None:
        """Override to not double record during playback."""
        self.frames.append(frame)
        if frame.guid:
            self.guid = frame.guid


__all__ = ["Agent", "Playback"]
