"""OmniLink Tool Runner — base class for local tool controllers.

Provides a reusable game/task loop that handles:
- Agent profile setup
- Tool-call kickoff (one API call)
- Memory persistence and UI command polling
- Periodic agent reviews
- Final analysis

Subclasses supply game-specific behaviour by overriding hooks.

Quick start
-----------
::

    from omnilink.tool_runner import ToolRunner

    class BreakoutRunner(ToolRunner):
        agent_name = "breakout-agent"
        display_name = "Breakout"

        def get_state(self):
            return breakout_api.get_state()

        def execute_action(self, state):
            action = breakout_engine.decide_action(state)
            breakout_api.send_action(action)

        def state_summary(self, state):
            return breakout_engine.state_summary(state)

        def is_game_over(self, state):
            return state.get("game_state") == "GAMEOVER"

        def log_events(self, state):
            ...

    if __name__ == "__main__":
        BreakoutRunner().run()
"""

from __future__ import annotations

import os
import re
import time
from typing import Any

from omnilink.client import OmniLinkClient


class ToolRunner:
    """Base class for OmniLink local tool controllers.

    Override the hook methods to adapt this runner to any game or task.
    Call ``run()`` to start the tool-calling loop.
    """

    # ── Required overrides ────────────────────────────────────────────
    agent_name: str = "tool-agent"
    display_name: str = "Tool"

    # ── Configuration (override as needed) ────────────────────────────
    base_url: str = "https://www.omnilink-agents.com"
    omni_key: str = os.environ.get("OMNI_KEY", "")
    engine: str = "g3-engine"
    poll_interval: float = 0.0
    memory_every: int = 60
    ask_every: int = 2400
    tool_name: str = "make_move"
    tool_description: str = "Execute the next action."
    commands: str = "stop_game, pause_game, resume_game"

    # ── Hooks (override in subclass) ──────────────────────────────────

    def get_state(self) -> dict[str, Any]:
        """Fetch the current state from the target system.

        Must return a dict. Called every tick.
        """
        raise NotImplementedError

    def execute_action(self, state: dict[str, Any]) -> None:
        """Decide and send the next action based on ``state``.

        Called every tick while the game is active and not paused.
        """
        raise NotImplementedError

    def state_summary(self, state: dict[str, Any]) -> str:
        """Return a concise text summary of the current state.

        Used for memory persistence and final reporting.
        """
        raise NotImplementedError

    def is_game_over(self, state: dict[str, Any]) -> bool:
        """Return True if the task/game has ended."""
        raise NotImplementedError

    def log_events(self, state: dict[str, Any]) -> None:
        """Print noteworthy events (score changes, lives lost, etc.).

        Called every tick. Override to add game-specific logging.
        """
        pass

    def on_start(self) -> None:
        """Called after the agent kicks off, before the main loop.

        Override to send a RESUME/START command to the game server.
        """
        pass

    def game_over_message(self, state: dict[str, Any]) -> str:
        """Return the game-over banner text."""
        return f"GAME OVER"

    # ── System instructions (override for custom prompts) ─────────────

    def get_system_instruction(self) -> dict[str, Any]:
        """System instruction for the initial tool-call kickoff."""
        return {
            "mainTask": f"{self.display_name} coordinator. Call {self.tool_name} now.",
            "availableTools": self.tool_name,
            "availableToolDetails": [
                {"name": self.tool_name, "description": self.tool_description},
            ],
            "availableCommands": self.commands,
            "allowToolUse": True,
        }

    def get_review_instruction(self) -> dict[str, Any]:
        """System instruction for periodic agent reviews."""
        return {
            "mainTask": (
                f"Monitor {self.display_name}. GAMEOVER or 0 lives: stop_game. "
                f"Otherwise: {self.tool_name}."
            ),
            "availableTools": self.tool_name,
            "availableToolDetails": [
                {"name": self.tool_name, "description": "Continue."},
            ],
            "availableCommands": self.commands,
            "allowToolUse": True,
        }

    def get_profile_settings(self) -> dict[str, Any]:
        """Settings dict for the agent profile."""
        return {
            "agentName": self.agent_name,
            "mainTask": (
                f"{self.display_name} coordinator. Use {self.tool_name} to play. "
                f"Check memory for state."
            ),
            "availableTools": self.tool_name,
            "availableToolDetails": [
                {"name": self.tool_name, "description": self.tool_description},
            ],
            "availableCommands": self.commands,
            "allowToolUse": True,
        }

    # ── Internal helpers ──────────────────────────────────────────────

    _profile_id: str | None = None
    _conversation_history: list[dict[str, Any]] = []

    def _memory_agent_key(self) -> str:
        """Return the key used for short-term memory storage.

        The UI resolves memory by profile UUID, so prefer that over the
        human-readable agent name.
        """
        return self._profile_id or self.agent_name

    def _append_turn(self, prompt: str, response: str) -> None:
        """Append a user/model turn to the conversation history."""
        ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        self._conversation_history.append(
            {"role": "user", "parts": [{"text": prompt}], "timestamp": ts},
        )
        self._conversation_history.append(
            {"role": "model", "parts": [{"text": response}], "timestamp": ts},
        )

    def _ensure_profile(self, client: OmniLinkClient) -> None:
        try:
            profiles = client.list_profiles()
            if profiles:
                pid = profiles[0].get("id")
                self._profile_id = pid
                client.update_profile(pid, name=self.agent_name, settings=self.get_profile_settings())
                print(f"  Profile updated to '{self.agent_name}' (id={pid}).")
            else:
                result = client.create_profile(self.agent_name, settings=self.get_profile_settings())
                self._profile_id = result.get("id")
                print(f"  Profile '{self.agent_name}' created (id={self._profile_id}).")
        except Exception as e:
            err = str(e).encode("ascii", errors="ignore").decode("ascii")
            print(f"  Could not set profile (continuing anyway): {err}")

    def _save_memory(self, client: OmniLinkClient, summary: str) -> None:
        self._append_turn(
            "Report current state.",
            f"Command: none\nResponse: {summary}",
        )
        try:
            client.set_memory(self._memory_agent_key(), self._conversation_history)
        except Exception:
            pass

    _max_none_retries: int = 2

    def _ask_for_tool_command(self, client: OmniLinkClient, prompt: str) -> str | None:
        """Ask the agent for the next command via a tool call.

        Sends *prompt* (typically the current state + result of the last
        command) and parses the ``Command:`` line from the agent's response.
        Returns the command string, or ``None`` only if the agent explicitly
        said ``stop_game``.  If the agent says ``Command: none``, retries
        with a stronger nudge up to ``_max_none_retries`` times.
        """
        suffix = (
            "\n\nCRITICAL: You are in tool-execution mode. "
            "You MUST reply with a valid robot command in the Command: field. "
            "Do NOT use Command: none unless the operator explicitly said stop. "
            "Always issue the next command to continue the operator's task."
        )
        current_prompt = prompt
        for attempt in range(1 + self._max_none_retries):
            try:
                result = client.chat(
                    current_prompt,
                    agent_name=self.agent_name,
                    engine=self.engine,
                    system_instruction=self.get_system_instruction(),
                    system_instruction_suffix=suffix,
                    temperature=0.3,
                )
                raw = result.get("text", "")
                raw_clean = raw.encode("ascii", errors="ignore").decode("ascii")

                # Check for explicit stop.
                if re.search(r"(?i)Command:\s*stop_game", raw_clean):
                    return None

                # Check for none — retry with a nudge.
                if re.search(r"(?i)Command:\s*none\b", raw_clean):
                    if attempt < self._max_none_retries:
                        print(f"  << Agent replied Command: none — retrying ({attempt + 1}/{self._max_none_retries})...")
                        current_prompt = (
                            f"{prompt}\n\n"
                            f"You replied Command: none but the operator's request is NOT fully completed yet. "
                            f"Review the request again and send the next robot_command. "
                            f"Only use stop_game when every single part is done."
                        )
                        continue
                    return None

                # Extract the command value.
                m = re.search(r"(?i)^Command:\s*(.+)", raw_clean, re.MULTILINE)
                if m:
                    cmd = m.group(1).strip()
                    # Strip trailing Tool: lines that may follow.
                    cmd = re.split(r"\n", cmd)[0].strip()
                    if cmd:
                        return cmd
                return None
            except Exception as e:
                err = str(e).encode("ascii", errors="ignore").decode("ascii")
                print(f"  [WARN] Agent tool call failed: {err}")
                return None
        return None

    def _ask_agent(self, client: OmniLinkClient, question: str) -> str:
        try:
            result = client.chat(
                question,
                agent_name=self.agent_name,
                engine=self.engine,
                temperature=0.3,
            )
            raw = result.get("text", "")
            m = re.search(r"(?i)^Response:\s*(.+)", raw, re.MULTILINE | re.DOTALL)
            text = m.group(1).strip() if m else raw.strip()
            return text.encode("ascii", errors="ignore").decode("ascii")
        except Exception as e:
            return f"(could not reach agent: {e})"

    def _check_memory_command(self, client: OmniLinkClient) -> str | None:
        try:
            memory = client.get_memory(self._memory_agent_key())
            if not memory:
                return None
            for entry in memory:
                for part in entry.get("parts", []):
                    text = part.get("text", "")
                    if re.search(r"(?i)Command:\s*stop_game", text):
                        return "stop_game"
                    if re.search(r"(?i)Command:\s*pause_game", text):
                        return "pause_game"
                    if re.search(r"(?i)Command:\s*resume_game", text):
                        return "resume_game"
            return None
        except Exception:
            return None

    def _agent_should_stop(self, client: OmniLinkClient) -> tuple[bool, str]:
        try:
            result = client.chat(
                f"Review state. {self.tool_name} or stop_game.",
                agent_name=self.agent_name,
                engine=self.engine,
                system_instruction=self.get_review_instruction(),
                temperature=0.0,
            )
            raw = result.get("text", "")
            raw_clean = raw.encode("ascii", errors="ignore").decode("ascii")

            stop_called = bool(re.search(r"(?im)^\s*Command:\s*stop_game", raw_clean))
            m = re.search(
                r"(?i)^Response:\s*(.+?)(?=\n(?:Tool|Stop-Tool|Command):|\Z)",
                raw_clean, re.MULTILINE | re.DOTALL,
            )
            reason = m.group(1).strip() if m else raw_clean.strip()[:200]
            return stop_called, reason
        except Exception as e:
            return False, f"(could not reach agent: {e})"

    # ── Main entry point ──────────────────────────────────────────────

    def run(self) -> None:
        """Execute the full tool-calling loop."""
        self._conversation_history = []
        client = OmniLinkClient(omni_key=self.omni_key, base_url=self.base_url, timeout=60)
        self._client = client

        self._ensure_profile(client)

        # Clear stale memory from previous sessions so old stop_game
        # commands don't immediately halt this run.
        try:
            client.set_memory(self._memory_agent_key(), [])
        except Exception:
            pass

        # Kick off with one API call.
        print("  Requesting tool call from OmniLink agent...")
        try:
            result = client.chat(
                f"Call {self.tool_name}.",
                agent_name=self.agent_name,
                engine=self.engine,
                system_instruction=self.get_system_instruction(),
                temperature=0.0,
            )
            raw = result.get("text", "").encode("ascii", errors="ignore").decode("ascii")
            tool_called = bool(re.search(rf"(?i)^Tool:\s*{re.escape(self.tool_name)}", raw, re.MULTILINE))
            if tool_called:
                print(f"  Agent called Tool: {self.tool_name}")
            else:
                print(f"  Agent responded (proceeding anyway): {raw.strip()[:80]}")
        except Exception as e:
            err = str(e).encode("ascii", errors="ignore").decode("ascii")
            print(f"  API call failed (proceeding with local engine): {err}")

        print()
        print("=" * 60)
        print(f"  OmniLink {self.display_name} — AI Controller via Tool Call")
        print("=" * 60)
        print()

        self.on_start()

        last_memory_time = time.time()
        last_ask_time = time.time()
        paused = False

        while True:
            try:
                state = self.get_state()
            except Exception as e:
                print(f"  [ERROR] Cannot reach server: {e}")
                time.sleep(1)
                continue

            if self.is_game_over(state):
                print()
                print(f"  {self.game_over_message(state)}")
                break

            if not paused:
                self.execute_action(state)
                self.log_events(state)

            now = time.time()

            # Check for UI commands before overwriting memory.
            check_interval = 3 if paused else self.memory_every
            if now - last_memory_time >= check_interval:
                cmd = self._check_memory_command(client)
                if cmd == "stop_game":
                    print()
                    print("  >> Stop requested from OmniLink UI.")
                    break
                elif cmd == "pause_game" and not paused:
                    paused = True
                    print()
                    print("  >> Pause requested from OmniLink UI. Game paused.")
                elif cmd == "resume_game" and paused:
                    paused = False
                    print()
                    print("  >> Resume requested from OmniLink UI. Game resumed.")

                if not paused:
                    summary = self.state_summary(state)
                    self._save_memory(client, summary)
                last_memory_time = now

            # Periodic agent review.
            if now - last_ask_time >= self.ask_every:
                print()
                print(f"  >> Agent reviewing {self.display_name}...")
                stop, reason = self._agent_should_stop(client)
                if stop:
                    print(f"  << Agent called stop_game: {reason}")
                    print()
                    break
                else:
                    print(f"  << Agent called {self.tool_name}: {reason}")
                    print()
                last_ask_time = now

            time.sleep(self.poll_interval)

        # Final summary.
        print()
        print("=" * 60)
        try:
            final_state = self.get_state()
            final_summary = self.state_summary(final_state)
        except Exception:
            final_summary = "(could not fetch final state)"

        print(f"  {final_summary}")
        print("=" * 60)

        try:
            self._append_turn(
                "Session ended.",
                f"Command: stop_game\nResponse: {final_summary}",
            )
            client.set_memory(self._memory_agent_key(), self._conversation_history)
            print("  Game saved to OmniLink memory.")
        except Exception:
            pass

        print()
        print("  >> Asking agent for final analysis...")
        answer = self._ask_agent(client, "Summarize the game.")
        print(f"  << Agent: {answer}")
        print()
