"""ProxAgent — single object for config, skills, and logging."""

import os
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

from supabase import create_client, Client


class ProxAgent:
    """Supabase-backed agent config and logging.

    One object handles everything: fetch skills, write them to disk,
    and log SDK messages to Supabase.

    Usage:
        agent = ProxAgent("otmc-agent", supabase_url="...", supabase_key="...")
        agent.fetch()
        agent.write_skills(".")
        agent.start_run(run_label="COPS-123")

        async for message in client.receive_response():
            agent.log(message)
    """

    def __init__(self, agent_slug: str, supabase_url: str = None, supabase_key: str = None):
        url = supabase_url or os.environ["SUPABASE_URL"]
        key = supabase_key or os.environ["SUPABASE_SERVICE_ROLE_KEY"]
        self.sb: Client = create_client(url, key)
        self.slug = agent_slug
        self.agent: dict | None = None
        self.skills: list[dict] = []
        self.memories: list[dict] = []

        # Run state (set by start_run / log)
        self._run_label: Optional[str] = None
        self._sandbox_id: Optional[str] = None
        self._session_id: Optional[str] = None
        self._model: Optional[str] = None
        self._sequence: int = 0
        self._run_created: bool = False

    # ── Config ──────────────────────────────────────────────────────

    def fetch(self) -> "ProxAgent":
        """Pull agent row and active skills from Supabase."""
        self.agent = (
            self.sb.table("_proxagent_agents")
            .select("*")
            .eq("slug", self.slug)
            .single()
            .execute()
            .data
        )
        self.skills = (
            self.sb.table("_proxagent_skills")
            .select("*")
            .eq("agent_id", self.agent["id"])
            .eq("is_active", True)
            .execute()
            .data
        )
        self.memories = (
            self.sb.table("_proxagent_memories")
            .select("*")
            .eq("agent_id", self.agent["id"])
            .eq("is_active", True)
            .order("created_at")
            .execute()
            .data
        )
        return self

    def skills_files(self) -> list[dict]:
        """Return skills as a list of {path, content} dicts.

        Each entry:
          path:    "{skill_name}/SKILL.md"
          content: YAML frontmatter + markdown body
        """
        files = []
        for skill in self.skills:
            name = skill["name"]
            description = skill.get("description", "")
            body = skill.get("content", "")
            content = f"---\nname: {name}\ndescription: {description}\n---\n\n{body}"
            files.append({"path": f"{name}/SKILL.md", "content": content})
        return files

    def write_skills(self, base_dir: str = ".") -> int:
        """Write skills to {base_dir}/.claude/skills/{name}/SKILL.md.

        Clears the skills directory first to remove stale skills.
        Returns the number of skills written.
        """
        skills_dir = Path(base_dir) / ".claude" / "skills"
        if skills_dir.exists():
            shutil.rmtree(skills_dir)
        for sf in self.skills_files():
            path = skills_dir / sf["path"]
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(sf["content"])
        return len(self.skills)

    def claude_md_content(self) -> str:
        """Return memories formatted as bullet points for CLAUDE.md."""
        lines = [
            f"- {m['learning'].strip()}"
            for m in self.memories
            if m.get("learning", "").strip()
        ]
        return "\n".join(lines) + "\n" if lines else ""

    def write_claude_md(self, base_dir: str = ".") -> int:
        """Write memories as bullet points to {base_dir}/CLAUDE.md.

        Returns the number of memories written.
        """
        content = self.claude_md_content()
        if content:
            Path(base_dir, "CLAUDE.md").write_text(content)
        return len(self.memories)

    def build_prompt(self, **kwargs) -> str:
        """Return system_prompt with placeholders filled in.

        Example:
            agent.build_prompt(issue_key="COPS-123")
        """
        return self.agent["system_prompt"].format(**kwargs)

    def save_memory(self, learning: str, context: str = "") -> None:
        """Save a new memory/learning to _proxagent_memories."""
        row = {"learning": learning}
        if context:
            row["context"] = context
        if self.agent:
            row["agent_id"] = self.agent["id"]
        self.sb.table("_proxagent_memories").insert(row).execute()

    # ── Logging ─────────────────────────────────────────────────────

    def start_run(self, run_label: str = None, sandbox_id: str = None) -> "ProxAgent":
        """Start a new run. Call this before logging messages.

        The actual run row is created when the first SystemMessage arrives
        (which carries the SDK-generated session_id).
        """
        self._run_label = run_label
        self._sandbox_id = sandbox_id
        self._session_id = None
        self._model = None
        self._sequence = 0
        self._run_created = False
        return self

    def log(self, message: Any) -> None:
        """Log a Claude Agent SDK message."""
        from claude_agent_sdk import (
            AssistantMessage,
            UserMessage,
            SystemMessage,
            ResultMessage,
            TextBlock,
            ThinkingBlock,
            ToolUseBlock,
            ToolResultBlock,
        )

        if isinstance(message, SystemMessage):
            if message.subtype == "init":
                self._session_id = message.data.get("session_id")
                self._model = message.data.get("model")
                self._create_run()
            self._insert_log(
                message_type="SystemMessage",
                subtype=message.subtype,
                data=message.data,
            )

        elif isinstance(message, AssistantMessage):
            for block in message.content:
                base = {
                    "message_type": "AssistantMessage",
                    "model": message.model,
                    "error": message.error,
                    "parent_tool_use_id": message.parent_tool_use_id,
                }

                if isinstance(block, TextBlock):
                    self._insert_log(**base, block_type="TextBlock", text=block.text)

                elif isinstance(block, ThinkingBlock):
                    self._insert_log(
                        **base,
                        block_type="ThinkingBlock",
                        thinking=block.thinking,
                        signature=block.signature,
                    )

                elif isinstance(block, ToolUseBlock):
                    self._insert_log(
                        **base,
                        block_type="ToolUseBlock",
                        tool_use_id=block.id,
                        tool_name=block.name,
                        tool_input=block.input,
                    )

                elif isinstance(block, ToolResultBlock):
                    content = block.content
                    if content is not None:
                        content = str(content)[:50000]
                    self._insert_log(
                        **base,
                        block_type="ToolResultBlock",
                        tool_use_id=block.tool_use_id,
                        tool_result_content=content,
                        tool_result_is_error=block.is_error,
                    )

        elif isinstance(message, UserMessage):
            content = message.content if isinstance(message.content, str) else None
            self._insert_log(
                message_type="UserMessage",
                uuid=message.uuid,
                parent_tool_use_id=message.parent_tool_use_id,
                user_content=content,
                tool_use_result=message.tool_use_result,
            )

        elif isinstance(message, ResultMessage):
            self._finalize_run(message)

    def _create_run(self) -> None:
        """Create run row in _proxagent_runs."""
        if self._run_created or not self._session_id:
            return
        try:
            row = {
                "session_id": self._session_id,
                "run_label": self._run_label,
                "sandbox_id": self._sandbox_id,
                "model": self._model,
                "started_at": datetime.utcnow().isoformat(),
            }
            if self.agent:
                row["agent_id"] = self.agent["id"]
            self.sb.table("_proxagent_runs").insert(row).execute()
            self._run_created = True
        except Exception as e:
            print(f"[WARN] proxagent _create_run failed: {e}")

    def _insert_log(self, **kwargs) -> None:
        """Insert log row into _proxagent_logs."""
        if not self._session_id:
            return
        kwargs["session_id"] = self._session_id
        kwargs["sequence"] = self._sequence
        kwargs["timestamp"] = datetime.utcnow().isoformat()
        self._sequence += 1
        kwargs = {k: v for k, v in kwargs.items() if v is not None}
        try:
            self.sb.table("_proxagent_logs").insert(kwargs).execute()
        except Exception as e:
            print(f"[WARN] proxagent _insert_log failed: {e}")

    def _finalize_run(self, r: Any) -> None:
        """Update run with final ResultMessage data."""
        if not self._session_id:
            return
        usage = r.usage or {}
        try:
            self.sb.table("_proxagent_runs").update({
                "subtype": r.subtype,
                "duration_ms": r.duration_ms,
                "duration_api_ms": r.duration_api_ms,
                "is_error": r.is_error,
                "num_turns": r.num_turns,
                "total_cost_usd": float(r.total_cost_usd) if r.total_cost_usd else None,
                "result": r.result[:50000] if r.result else None,
                "structured_output": r.structured_output,
                "input_tokens": usage.get("input_tokens"),
                "output_tokens": usage.get("output_tokens"),
                "cache_creation_input_tokens": usage.get("cache_creation_input_tokens"),
                "cache_read_input_tokens": usage.get("cache_read_input_tokens"),
                "ended_at": datetime.utcnow().isoformat(),
            }).eq("session_id", self._session_id).execute()
        except Exception as e:
            print(f"[WARN] proxagent _finalize_run failed: {e}")
