"""Core agent loop. Single loop, no sub-agents. Trust the model.

Features:
- Retry on API errors with exponential backoff
- Tool execution with error boundaries
- Doom loop detection (same tool call repeated)
- Context compaction when approaching token limits
- Max turns enforcement with graceful degradation
- Per-run stats tracking (turns, tool calls, token estimates)
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
import sys
from typing import Any, TypedDict

if sys.version_info >= (3, 11):
    from typing import NotRequired
else:
    from typing_extensions import NotRequired

from rich.console import Console
from rich.markup import escape as rich_escape
from rich.panel import Panel

from krim.models.base import Model, ModelResponse, ToolCall
from krim.tools import get_tool, tool_schemas
from krim.tools.base import Tool
from krim.compaction import needs_compaction, compact, estimate_message_tokens
from krim.retry import with_retry

console = Console()


# Type definitions for message formats
class ToolUseBlock(TypedDict):
    type: str  # "tool_use"
    id: str
    name: str
    input: dict[str, Any]


class ToolResultBlock(TypedDict):
    type: str  # "tool_result"
    tool_use_id: str
    content: str | list[dict[str, Any]]


class TextBlock(TypedDict):
    type: str  # "text"
    text: str


class ImageBlock(TypedDict):
    type: str  # "image"
    source: dict[str, str]


ContentBlock = ToolUseBlock | ToolResultBlock | TextBlock | ImageBlock


class Message(TypedDict):
    """Message format for conversation history."""
    role: str  # "system" | "user" | "assistant" | "tool"
    content: NotRequired[str | list[ContentBlock]]
    tool_calls: NotRequired[list[dict[str, Any]]]  # OpenAI format
    tool_call_id: NotRequired[str]  # OpenAI tool response
    name: NotRequired[str]  # OpenAI tool name


@dataclass
class RunStats:
    """Stats for a single agent.run() invocation."""
    turns: int = 0
    tool_calls: int = 0
    compactions: int = 0
    tool_call_names: dict[str, int] = field(default_factory=dict)

    def record_tool_call(self, name: str):
        self.tool_calls += 1
        self.tool_call_names[name] = self.tool_call_names.get(name, 0) + 1

MAX_STEPS_PROMPT = (
    "You have reached the maximum number of turns. "
    "You MUST finish now. If there are pending changes, finalize them immediately. "
    "Write any remaining files and verify critical results. "
    "Then provide a brief summary of what was accomplished and what remains."
)

DOOM_RECOVERY_PROMPT = (
    "You are repeating the same action. This approach is not working. "
    "Stop and reconsider: try a completely different strategy, "
    "read error messages carefully, or check your assumptions."
)

LAST_TURN_WARNING = (
    "\n[SYSTEM: You have {remaining} turn(s) remaining. "
    "Prioritize completing the task over exploration. "
    "Finalize any changes and verify results NOW.]"
)

NUDGE_PROMPT = (
    "You provided thoughts but did not call any tools. "
    "If you are still working, please proceed with your next action. "
    "If you have completed the task, you MUST call the `submit` tool to finish."
)

NUDGE_WARNING_PROMPT = (
    "WARNING: You have not used any tools for 3 consecutive turns. "
    "Stop thinking and take ACTION now. Either use tools to continue working, "
    "or call `submit` if you are done."
)


# -- message builders (provider-specific formatting) --

def _build_assistant_msg_claude(response: ModelResponse) -> dict:
    content = []
    if response.text:
        content.append({"type": "text", "text": response.text})
    for tc in response.tool_calls:
        content.append({
            "type": "tool_use",
            "id": tc.id,
            "name": tc.name,
            "input": tc.args,
        })
    return {"role": "assistant", "content": content}


def _build_assistant_msg_openai(response: ModelResponse) -> dict:
    msg: dict = {"role": "assistant"}
    if response.text:
        msg["content"] = response.text
    if response.tool_calls:
        msg["tool_calls"] = [
            {
                "id": tc.id,
                "type": "function",
                "function": {
                    "name": tc.name,
                    "arguments": json.dumps(tc.args),
                },
            }
            for tc in response.tool_calls
        ]
    return msg


def _build_tool_result_openai(tool_call_id: str, name: str, result: str) -> dict:
    return {
        "role": "tool",
        "tool_call_id": tool_call_id,
        "name": name,
        "content": result,
    }


# Image marker from read.py
_IMAGE_MARKER = "<<<IMAGE:"


def _parse_image_result(result: str) -> dict | None:
    """Parse image marker and return Claude API image block, or None if not an image."""
    if not result.startswith(_IMAGE_MARKER):
        return None
    try:
        # Format: <<<IMAGE:media_type:base64_data>>>
        content = result[len(_IMAGE_MARKER):-3]  # Remove marker and closing >>>
        media_type, b64_data = content.split(":", 1)
        return {
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": media_type,
                "data": b64_data,
            }
        }
    except Exception:
        return None


def _build_tool_result_content_claude(result: str) -> list | str:
    """Build tool result content for Claude API. Returns image block list or string."""
    image_block = _parse_image_result(result)
    if image_block:
        return [image_block]
    return result


class Agent:
    def __init__(
        self,
        model: Model,
        provider: str,
        system_prompt: str,
        tools: list[Tool],
        mcp_tools: list[Tool] | None = None,
        max_turns: int = 10,
        verbose: bool = False,
        token_limit: int = 120_000,
        log_dir: str | None = None,
    ):
        self.model = model
        self.provider = provider
        self.max_turns = max_turns
        self.tools = tools
        self.mcp_tools = mcp_tools or []
        self.verbose = verbose
        self.token_limit = token_limit
        self.log_dir = log_dir
        self._system_prompt = system_prompt  # keep for logging
        self.messages: list[Message] = [{"role": "system", "content": system_prompt}]

        # doom loop detection: track recent tool calls + recovery count
        self._recent_calls: list[str] = []
        self._doom_recoveries: int = 0

        # consecutive no-tool turns (for nudge logic)
        self._consecutive_no_tool_turns: int = 0

        # track last meaningful response for final output
        self._last_response: str | None = None

        # cumulative stats
        self.last_stats: RunStats | None = None
        self.total_turns: int = 0
        self.total_tool_calls: int = 0

    # -- tool management --

    def _all_tool_schemas(self) -> list[dict]:
        schemas = tool_schemas(self.tools)
        for mt in self.mcp_tools:
            schemas.append(mt.schema())
        return schemas

    def _execute_tool(self, name: str, args: dict) -> str:
        """Execute a tool with error boundary."""
        # check mcp tools first
        for mt in self.mcp_tools:
            if mt.name == name:
                try:
                    return mt.run(**args)
                except Exception as e:
                    return f"error: tool '{name}' raised: {e}"

        tool = get_tool(self.tools, name)
        if not tool:
            return f"error: unknown tool '{name}'"
        try:
            return tool.run(**args)
        except Exception as e:
            return f"error: tool '{name}' raised: {e}"

    # -- doom loop detection --

    def _check_doom_loop(self, tool_calls: list[ToolCall]) -> bool:
        """Detect if the agent is stuck calling the same tools repeatedly."""
        sig = json.dumps([(tc.name, tc.args) for tc in tool_calls], sort_keys=True)
        self._recent_calls.append(sig)
        if len(self._recent_calls) > 10:
            self._recent_calls = self._recent_calls[-10:]

        # same call 3 times in a row
        if len(self._recent_calls) >= 3:
            last3 = self._recent_calls[-3:]
            if last3[0] == last3[1] == last3[2]:
                return True
        return False

    def _build_turn_hint(self, current_turn: int) -> str | None:
        """Build context hint for the agent based on current state."""
        parts = []
        # doom recovery nudge (1st doom loop was detected this turn)
        if self._doom_recoveries == 1 and len(self._recent_calls) == 0:
            parts.append(DOOM_RECOVERY_PROMPT)
        # last turn warning (only when max_turns > 5 to avoid noise in short runs)
        remaining = self.max_turns - current_turn
        if self.max_turns > 5 and remaining <= 2 and remaining > 0:
            parts.append(LAST_TURN_WARNING.format(remaining=remaining))
        return "\n".join(parts) if parts else None

    # -- display --

    def _print_tool_call(self, name: str, args: dict):
        icon = {"bash": ">>", "read": "<<", "write": "=>", "edit": "<>", "grep": "??", "glob": "**"}.get(name, "::")
        summary = f"[bold cyan]{icon} {name}[/]"
        if name == "bash":
            cmd = args.get("command", "")
            # show first line of multi-line commands
            first_line = cmd.split("\n")[0]
            if len(first_line) > 100:
                first_line = first_line[:97] + "..."
            summary += f"  `{first_line}`"
        elif name in ("read", "write", "edit"):
            path = args.get("path", "")
            summary += f"  {path}"
        else:
            summary += f"  {json.dumps(args, ensure_ascii=False)[:80]}"
        console.print(summary)

    def _print_tool_result(self, result: str):
        # Handle image results - don't print base64
        if result.startswith(_IMAGE_MARKER):
            try:
                content = result[len(_IMAGE_MARKER):-3]
                media_type = content.split(":", 1)[0]
                console.print(Panel(f"[image: {media_type}]", border_style="dim", expand=False, padding=(0, 1)))
                return
            except Exception:
                pass

        lines = result.splitlines()
        preview = "\n".join(lines[:20])
        # Escape Rich markup characters to prevent parsing errors
        preview = rich_escape(preview)
        if len(lines) > 20:
            preview += f"\n[dim]... ({len(lines) - 20} more lines)[/]"
        console.print(Panel(preview, border_style="dim", expand=False, padding=(0, 1)))

    # -- core loop --

    def run(self, user_input: str) -> RunStats:
        """Execute a single user request through the agent loop."""
        # Initialize submit tool state for this run (instance-level, thread-safe)
        from krim.tools.submit import SubmitTool
        for tool in self.tools:
            if isinstance(tool, SubmitTool):
                tool.reset(user_input)
                break

        self.messages.append({"role": "user", "content": user_input})
        self._recent_calls.clear()
        self._doom_recoveries = 0
        self._consecutive_no_tool_turns = 0
        self._last_response = None
        stats = RunStats()

        # cache tool schemas (deterministic order for prompt cache)
        cached_schemas = self._all_tool_schemas()

        # wrap model.chat with retry
        chat_with_retry = with_retry(self.model.chat)

        # stream callback (only print in verbose mode)
        def stream_cb(text: str):
            if self.verbose:
                console.print(text, end="", highlight=False)

        turn = 0
        while turn < self.max_turns:
            turn += 1
            stats.turns = turn

            if self.verbose:
                tokens = estimate_message_tokens(self.messages)
                console.print(f"\n[dim]--- turn {turn}/{self.max_turns}  ~{tokens:,} tokens ---[/]")

            # check for compaction
            if needs_compaction(self.messages, max_tokens=self.token_limit):
                if self.verbose:
                    console.print("[dim]compacting conversation...[/]")
                self.messages = compact(self.messages, max_tokens=self.token_limit)
                stats.compactions += 1

            # call model with retry
            try:
                response = chat_with_retry(
                    messages=self.messages,
                    tools=cached_schemas,
                    stream_callback=stream_cb,
                )
            except Exception as e:
                console.print(f"\n[red]model error: {e}[/]")
                break

            if response.text:
                if self.verbose:
                    console.print()
                # capture longest response (real answers are longer than meta-commentary)
                if self._last_response is None or len(response.text) > len(self._last_response):
                    self._last_response = response.text

            # no tool calls -> nudge to continue or use submit
            if not response.tool_calls:
                self._consecutive_no_tool_turns += 1

                if response.text:
                    self.messages.append({"role": "assistant", "content": response.text})

                # after 3 consecutive no-tool turns, give stronger warning
                if self._consecutive_no_tool_turns >= 3:
                    if self.verbose:
                        console.print("[yellow]no tools for 3 turns, warning agent[/]")
                    self.messages.append({"role": "user", "content": NUDGE_WARNING_PROMPT})
                else:
                    if self.verbose:
                        console.print("[dim]no tool calls, nudging agent to continue[/]")
                    self.messages.append({"role": "user", "content": NUDGE_PROMPT})

                continue

            # doom loop detection with recovery
            if self._check_doom_loop(response.tool_calls):
                self._doom_recoveries += 1
                if self._doom_recoveries >= 2:
                    # 2nd doom loop: hard stop
                    if self.verbose:
                        console.print("[yellow]doom loop detected (2nd time), forcing stop[/]")
                    self.messages.append({"role": "user", "content": MAX_STEPS_PROMPT})
                    try:
                        final = chat_with_retry(
                            messages=self.messages,
                            tools=[],  # no tools, force text response
                            stream_callback=stream_cb,
                        )
                        if final.text:
                            if self.verbose:
                                console.print()
                            self._last_response = final.text
                            self.messages.append({"role": "assistant", "content": final.text})
                    except Exception as e:
                        if self.verbose:
                            console.print(f"[dim]final summary failed: {e}[/]")
                    break
                else:
                    # 1st doom loop: warn but keep tools available
                    if self.verbose:
                        console.print("[yellow]doom loop detected, nudging recovery[/]")
                    self._recent_calls.clear()  # reset so it gets a fresh chance

            # reset consecutive no-tool counter since we have tool calls
            self._consecutive_no_tool_turns = 0

            # check for submit tool first
            from krim.tools.submit import VERIFICATION_MARKER, SUBMIT_MARKER
            submitted = False
            submit_verification = False

            for tc in response.tool_calls:
                if tc.name == "submit":
                    # Execute submit tool to check if verification needed
                    result = self._execute_tool(tc.name, tc.args)
                    stats.record_tool_call("submit")

                    if VERIFICATION_MARKER in result:
                        # First submit: verification required
                        if self.verbose:
                            console.print("[yellow]Verification required - reviewing task[/]")
                        submit_verification = True
                        if self.provider in ("claude", "vertex_ai"):
                            self.messages.append(_build_assistant_msg_claude(response))
                            self.messages.append({"role": "user", "content": [
                                {"type": "tool_result", "tool_use_id": tc.id, "content": result}
                            ]})
                        else:
                            self.messages.append(_build_assistant_msg_openai(response))
                            self.messages.append(_build_tool_result_openai(tc.id, tc.name, result))
                    else:
                        # Second submit: actually done - extract summary as final response
                        summary = result.replace(SUBMIT_MARKER, "", 1) if result.startswith(SUBMIT_MARKER) else result
                        self._last_response = summary
                        if self.verbose:
                            console.print(f"[bold green]{summary}[/]")
                        if self.provider in ("claude", "vertex_ai"):
                            self.messages.append(_build_assistant_msg_claude(response))
                            self.messages.append({"role": "user", "content": [
                                {"type": "tool_result", "tool_use_id": tc.id, "content": result}
                            ]})
                        else:
                            self.messages.append(_build_assistant_msg_openai(response))
                            self.messages.append(_build_tool_result_openai(tc.id, tc.name, result))
                        submitted = True
                    break

            if submitted:
                break

            if submit_verification:
                continue  # Skip normal tool processing, go to next turn

            # add assistant message with tool calls
            if self.provider in ("claude", "vertex_ai"):
                self.messages.append(_build_assistant_msg_claude(response))
            else:
                self.messages.append(_build_assistant_msg_openai(response))

            # execute each tool call
            tool_results = []
            for tc in response.tool_calls:
                if self.verbose:
                    self._print_tool_call(tc.name, tc.args)
                result = self._execute_tool(tc.name, tc.args)
                if self.verbose:
                    self._print_tool_result(result)
                tool_results.append((tc, result))
                stats.record_tool_call(tc.name)

            # add tool results to messages
            if self.provider in ("claude", "vertex_ai"):
                content = []
                for tc, result in tool_results:
                    content.append({
                        "type": "tool_result",
                        "tool_use_id": tc.id,
                        "content": _build_tool_result_content_claude(result),
                    })
                # inject recovery/warning hints as text block in same message
                hint = self._build_turn_hint(turn)
                if hint:
                    content.append({"type": "text", "text": hint})
                self.messages.append({"role": "user", "content": content})
            else:
                for tc, result in tool_results:
                    self.messages.append(_build_tool_result_openai(tc.id, tc.name, result))
                # inject recovery/warning hints as separate user message
                hint = self._build_turn_hint(turn)
                if hint:
                    self.messages.append({"role": "user", "content": hint})

        else:
            # while-else: loop condition became false (not break) = all turns used with tool calls still pending
            if self.verbose:
                console.print(f"\n[yellow]reached max turns ({self.max_turns})[/]")
            # inject max_steps prompt for graceful summary
            self.messages.append({"role": "user", "content": MAX_STEPS_PROMPT})
            try:
                final = chat_with_retry(
                    messages=self.messages,
                    tools=[],
                    stream_callback=stream_cb,
                )
                if final.text:
                    if self.verbose:
                        console.print()
                    self._last_response = final.text
                    self.messages.append({"role": "assistant", "content": final.text})
            except Exception as e:
                if self.verbose:
                    console.print(f"[dim]final summary failed: {e}[/]")

        # print run stats (verbose only)
        if self.verbose:
            self._print_stats(stats)
        self.last_stats = stats
        self.total_turns += stats.turns
        self.total_tool_calls += stats.tool_calls

        # save logs if log_dir specified
        self._save_logs(stats)

        # print final response (non-verbose mode)
        if not self.verbose and self._last_response:
            console.print(self._last_response)

        return stats

    def _print_stats(self, stats: RunStats):
        parts = [f"turns: {stats.turns}"]
        if stats.tool_calls:
            tools_summary = ", ".join(f"{n}:{c}" for n, c in sorted(stats.tool_call_names.items()))
            parts.append(f"tool calls: {stats.tool_calls} ({tools_summary})")
        if stats.compactions:
            parts.append(f"compactions: {stats.compactions}")
        tokens = estimate_message_tokens(self.messages)
        parts.append(f"~{tokens:,} tokens")
        console.print(f"\n[dim]{' | '.join(parts)}[/]")

    def token_count(self) -> int:
        return estimate_message_tokens(self.messages)

    def force_compact(self):
        """Manually trigger compaction."""
        before = estimate_message_tokens(self.messages)
        self.messages = compact(self.messages)
        after = estimate_message_tokens(self.messages)
        console.print(f"[dim]compacted: ~{before:,} → ~{after:,} tokens[/]")

    def _save_logs(self, stats: RunStats):
        """Save conversation logs to log_dir (Terminus-compatible format)."""
        if not self.log_dir:
            return

        try:
            log_path = Path(self.log_dir)
            log_path.mkdir(parents=True, exist_ok=True)

            # prompt.txt - system prompt
            (log_path / "prompt.txt").write_text(self._system_prompt)

            # messages.json - full conversation history
            (log_path / "messages.json").write_text(
                json.dumps(self.messages, indent=2, ensure_ascii=False)
            )

            # meta.json - run metadata
            meta = {
                "provider": self.provider,
                "max_turns": self.max_turns,
                "turns": stats.turns,
                "tool_calls": stats.tool_calls,
                "tool_call_names": stats.tool_call_names,
                "compactions": stats.compactions,
                "tokens": estimate_message_tokens(self.messages),
            }
            (log_path / "meta.json").write_text(
                json.dumps(meta, indent=2, ensure_ascii=False)
            )

            if self.verbose:
                console.print(f"[dim]logs saved to {log_path}[/]")

        except Exception as e:
            console.print(f"[yellow]failed to save logs: {e}[/]")
