"""
AutonomousAgent — Claude Agent SDK wrapper for open-ended autonomous tasks.

Wraps the claude-agent-sdk Python package to provide autonomous agent
capabilities within the DataBridge AI Planner. Streams real-time tool
actions and text back via async iterators for SSE consumption.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, AsyncIterator, Optional

logger = logging.getLogger(__name__)

# Graceful import — feature hidden when SDK not installed
try:
    from claude_agent_sdk import (
        query as claude_query,
        ClaudeAgentOptions,
        ResultMessage,
        AssistantMessage,
        UserMessage,
        SystemMessage,
    )
    from claude_agent_sdk.types import (
        Message,
        TextBlock,
        ToolUseBlock,
        ToolResultBlock,
    )
    AGENT_SDK_AVAILABLE = True
except ImportError:
    AGENT_SDK_AVAILABLE = False
    logger.info("claude-agent-sdk not installed — autonomous mode disabled")


@dataclass
class SessionInfo:
    """Metadata for an autonomous agent session."""
    session_id: str
    created_at: str
    last_activity: str
    sdk_session_id: str | None = None  # Session ID returned by SDK for resume
    prompt_history: list[str] = field(default_factory=list)
    total_turns: int = 0
    total_cost_usd: float = 0.0


class AutonomousAgent:
    """Manages autonomous Claude Agent SDK sessions with DataBridge tools.

    Each session maintains conversation continuity via the SDK's built-in
    session resume mechanism. Tool calls, text, and cost are streamed as
    normalized dicts for SSE consumption.
    """

    def __init__(
        self,
        work_dir: str | Path | None = None,
        max_turns: int = 25,
        max_budget_usd: float = 1.0,
    ):
        self.work_dir = str(work_dir or Path.cwd())
        self.max_turns = max_turns
        self.max_budget_usd = max_budget_usd
        self.sessions: dict[str, SessionInfo] = {}
        self._purpose_text = self._load_purpose()

    @staticmethod
    def is_available() -> bool:
        """Check if the Agent SDK + CLI prerequisites are met."""
        return AGENT_SDK_AVAILABLE

    def _load_purpose(self) -> Optional[str]:
        """Load purpose.md strategic vision document from project root."""
        candidates = [
            Path(self.work_dir) / "purpose.md",
            Path(__file__).parent.parent.parent / "purpose.md",
        ]
        for p in candidates:
            try:
                if p.exists():
                    return p.read_text(encoding="utf-8")
            except Exception:
                continue
        return None

    def _get_purpose_context(self) -> str:
        """Return curated strategic summary from purpose.md."""
        if not self._purpose_text:
            return ""
        return (
            "## DataBridge AI Strategic Context\n"
            "Core strategy: Intelligence Arbitrage — converting consulting expertise "
            "into scalable autonomous agents.\n"
            "Key principles: Automated Trust (detect > diagnose > fix > verify), "
            "GraphRAG Memory, Sprinkler Behavior (drive closure), "
            "CEO/IT Bridge (business semantics), Delivery Compression.\n\n"
        )

    def _get_connections_context(self) -> str:
        """Discover Snowflake connections from config.toml."""
        try:
            import tomllib
        except ImportError:
            try:
                import tomli as tomllib  # type: ignore[no-redef]
            except ImportError:
                return ""
        config_path = Path(os.environ.get("USERPROFILE", str(Path.home()))) / ".snowflake" / "config.toml"
        if not config_path.exists():
            return ""
        try:
            with open(config_path, "rb") as f:
                config = tomllib.load(f)
            conns = config.get("connections", {})
            if not conns:
                return ""
            lines = [f"## Snowflake Connections ({len(conns)} configured)\n"]
            for name, params in list(conns.items())[:10]:
                if isinstance(params, dict):
                    acct = params.get("account", "?")
                    db = params.get("database", "?")
                    wh = params.get("warehouse", "?")
                    lines.append(f"- **{name}**: account={acct}, database={db}, warehouse={wh}")
            return "\n".join(lines) + "\n\n"
        except Exception:
            return ""

    def _get_hierarchy_context(self) -> str:
        """Discover hierarchy projects from data/hierarchy_projects/."""
        projects_dir = Path(self.work_dir) / "data" / "hierarchy_projects"
        if not projects_dir.exists():
            return ""
        try:
            project_files = list(projects_dir.glob("*/project.json"))
            if not project_files:
                return ""
            lines = [f"## Hierarchy Projects ({len(project_files)} found)\n"]
            for pf in project_files[:5]:
                pdata = json.loads(pf.read_text(encoding="utf-8"))
                name = pdata.get("name", pf.parent.name)
                hcount = len(pdata.get("hierarchies", []))
                lines.append(f"- **{name}** (dir: {pf.parent.name}): {hcount} hierarchies")
            return "\n".join(lines) + "\n\n"
        except Exception:
            return ""

    def _get_workflow_context(self) -> str:
        """Discover recent workflow run artifacts."""
        artifacts_dir = Path(self.work_dir) / "data" / "workflow_runs" / "artifacts"
        if not artifacts_dir.exists():
            return ""
        try:
            reports = sorted(artifacts_dir.glob("run_*_report.html"), key=lambda p: p.stat().st_mtime, reverse=True)
            if not reports:
                return ""
            lines = [f"## Recent Workflow Runs ({len(reports)} total)\n"]
            for r in reports[:5]:
                name = r.stem.replace("_report", "")
                lines.append(f"- {name} ({r.name})")
            return "\n".join(lines) + "\n\n"
        except Exception:
            return ""

    def _build_system_prompt(self) -> str:
        """Build the DataBridge-aware system prompt with dynamic context."""
        base = (
            "You are a DataBridge AI autonomous agent. You help users investigate "
            "data quality issues, explore database schemas, profile data, and solve "
            "open-ended data engineering problems.\n\n"
            "## Available DataBridge Context\n"
            "- Data directory: data/\n"
            "- Snowflake config: ~/.snowflake/config.toml\n"
            "- Hierarchy projects: data/hierarchy_projects/\n"
            "- Workflow artifacts: data/workflow_runs/artifacts/\n\n"
            "## Guidelines\n"
            "- Be thorough but concise in your investigations\n"
            "- When examining data, always profile it first before making conclusions\n"
            "- Show your reasoning as you work\n"
            "- Summarize findings at the end\n"
            "- If you find data quality issues, quantify them (row counts, percentages)\n\n"
        )
        # Append dynamic context sections
        base += self._get_purpose_context()
        base += self._get_connections_context()
        base += self._get_hierarchy_context()
        base += self._get_workflow_context()
        return base

    def _build_options(self, session: SessionInfo | None = None) -> "ClaudeAgentOptions":
        """Build ClaudeAgentOptions for an SDK query."""
        # Build clean env — remove CLAUDECODE to allow nested agent sessions
        clean_env = {k: v for k, v in os.environ.items() if k != "CLAUDECODE"}

        kwargs = {
            "system_prompt": self._build_system_prompt(),
            "cwd": self.work_dir,
            "allowed_tools": ["Read", "Write", "Edit", "Bash", "Grep", "Glob"],
            "permission_mode": "acceptEdits",
            "max_turns": self.max_turns,
            "max_budget_usd": self.max_budget_usd,
            "env": clean_env,
        }

        # Resume from existing session if available
        if session and session.sdk_session_id:
            kwargs["resume"] = session.sdk_session_id
            kwargs["continue_conversation"] = True

        return ClaudeAgentOptions(**kwargs)

    async def start_session(
        self,
        prompt: str,
        session_id: str | None = None,
    ) -> AsyncIterator[dict]:
        """Start a new autonomous session, yielding normalized event dicts."""
        if not AGENT_SDK_AVAILABLE:
            yield {"type": "error", "message": "Claude Agent SDK not installed. Run: pip install claude-agent-sdk"}
            return

        sid = session_id or str(uuid.uuid4())[:12]
        now = datetime.now(timezone.utc).isoformat()
        session = SessionInfo(
            session_id=sid,
            created_at=now,
            last_activity=now,
            prompt_history=[prompt],
        )
        self.sessions[sid] = session

        yield {"type": "session_start", "session_id": sid}

        async for event in self._run_query(sid, prompt, session):
            yield event

    async def continue_session(
        self,
        session_id: str,
        prompt: str,
    ) -> AsyncIterator[dict]:
        """Send a follow-up message to an existing session."""
        if session_id not in self.sessions:
            yield {"type": "error", "message": f"Session '{session_id}' not found"}
            return

        session = self.sessions[session_id]
        session.prompt_history.append(prompt)
        session.last_activity = datetime.now(timezone.utc).isoformat()

        async for event in self._run_query(session_id, prompt, session):
            yield event

    async def _run_query(
        self,
        session_id: str,
        prompt: str,
        session: SessionInfo,
    ) -> AsyncIterator[dict]:
        """Execute a query against the Claude Agent SDK and yield events."""
        try:
            # Remove CLAUDECODE from current process env to prevent
            # "cannot launch inside another session" detection by the SDK.
            # _build_options() also strips it from child env dict, but the
            # SDK checks os.environ before spawning the subprocess.
            os.environ.pop("CLAUDECODE", None)

            options = self._build_options(session)

            # claude_query is an async iterator — iterate directly
            async for msg in claude_query(prompt=prompt, options=options):
                # ResultMessage signals completion with usage stats
                if isinstance(msg, ResultMessage):
                    cost = msg.total_cost_usd or 0.0
                    turns = msg.num_turns or 0

                    # Capture SDK session_id for future resume
                    if msg.session_id:
                        session.sdk_session_id = msg.session_id

                    session.total_cost_usd += cost
                    session.total_turns += turns

                    # Skip msg.result — the same text was already streamed
                    # via AssistantMessage content blocks above.

                    yield {
                        "type": "done",
                        "session_id": session_id,
                        "cost_usd": round(cost, 4),
                        "total_cost_usd": round(session.total_cost_usd, 4),
                        "turns": turns,
                        "total_turns": session.total_turns,
                        "is_error": msg.is_error,
                    }
                    continue

                # Process assistant/user/system messages
                async for event in _process_message(msg):
                    yield event

        except Exception as exc:
            logger.exception("Autonomous agent query failed")
            error_msg = str(exc)

            # Detect common failure modes
            if "CLINotFoundError" in type(exc).__name__ or (
                "claude" in error_msg.lower() and "not found" in error_msg.lower()
            ):
                error_msg = (
                    "Claude Code CLI not found. Install it with: "
                    "npm install -g @anthropic-ai/claude-code"
                )
            elif "api" in error_msg.lower() and "key" in error_msg.lower():
                error_msg = "API key not configured. Set ANTHROPIC_API_KEY environment variable."

            yield {"type": "error", "message": error_msg}

    async def close_session(self, session_id: str) -> bool:
        """Clean up and remove a session."""
        if session_id in self.sessions:
            del self.sessions[session_id]
            return True
        return False

    def list_sessions(self) -> list[dict]:
        """Return metadata for all active sessions."""
        return [
            {
                "session_id": s.session_id,
                "created_at": s.created_at,
                "last_activity": s.last_activity,
                "prompts": len(s.prompt_history),
                "total_turns": s.total_turns,
                "total_cost_usd": round(s.total_cost_usd, 4),
            }
            for s in self.sessions.values()
        ]


async def _process_message(msg: Any) -> AsyncIterator[dict]:
    """Convert an SDK message into normalized event dicts.

    The claude-agent-sdk streams various message types:
    - SystemMessage (subtype='init') — session init metadata
    - AssistantMessage — with content list of TextBlock/ToolUseBlock/ThinkingBlock
    - UserMessage — with content list of ToolResultBlock
    - Raw objects like TextBlock, ToolUseBlock, ToolResultBlock, ThinkingBlock
    """
    cls_name = type(msg).__name__
    # SDK message types don't have a .role attribute — infer from class name
    role = getattr(msg, 'role', None)
    if role is None:
        if cls_name in ('AssistantMessage', 'TextBlock', 'ToolUseBlock'):
            role = 'assistant'
        elif cls_name in ('UserMessage', 'ToolResultBlock'):
            role = 'user'
        else:
            role = 'system'
    msg_str = str(msg)

    # ── Filter out noisy system/internal messages ──
    # Check class name AND string repr (SDK may use different class paths)
    if cls_name == 'SystemMessage' or msg_str.startswith('SystemMessage('):
        logger.warning("[_process_message] FILTERED SystemMessage")
        return

    # ── Handle raw block objects (SDK sometimes yields these directly) ──
    if cls_name == 'ThinkingBlock' or msg_str.startswith('ThinkingBlock('):
        return

    if cls_name == 'TextBlock' or msg_str.startswith('TextBlock('):
        text = getattr(msg, 'text', None)
        if text is None and msg_str.startswith('TextBlock('):
            # Extract text from repr: TextBlock(text='...')
            import re
            m = re.search(r"text=['\"](.+?)['\"]", msg_str)
            text = m.group(1) if m else msg_str
        if text:
            yield {"type": "text", "text": text, "role": "assistant"}
        return

    if cls_name == 'ToolUseBlock':
        name = getattr(msg, 'name', 'unknown')
        tool_input = getattr(msg, 'input', {})
        tool_id = getattr(msg, 'id', '')
        yield {
            "type": "tool_use",
            "tool_name": name,
            "tool_input": _truncate_input(tool_input),
            "tool_id": tool_id,
        }
        return

    if cls_name == 'ToolResultBlock':
        tool_id = getattr(msg, 'tool_use_id', '')
        result_content = getattr(msg, 'content', '')
        is_error = getattr(msg, 'is_error', False)
        yield {
            "type": "tool_result",
            "tool_id": tool_id,
            "result": _truncate_text(str(result_content), 2000),
            "is_error": is_error,
        }
        return

    # ── Filter system-role messages (init prompts, config, etc.) ──
    if role == 'system':
        logger.debug("[_process_message] FILTERED system-role message: %s", cls_name)
        return

    # ── Handle messages with content lists (AssistantMessage, UserMessage) ──
    content = getattr(msg, 'content', None)
    if content is None:
        # Unknown message type — skip if it looks internal, else emit
        if cls_name in ('ResultMessage',):
            return  # Handled upstream
        # Skip anything that looks like internal SDK metadata
        if len(msg_str) > 5000:
            logger.debug("[_process_message] FILTERED large unknown message: %s", cls_name)
            return
        yield {"type": "text", "text": str(msg), "role": role}
        return

    if isinstance(content, str):
        yield {"type": "text", "text": content, "role": role}
        return

    if isinstance(content, list):
        for block in content:
            block_cls = type(block).__name__
            block_type = getattr(block, 'type', None) or (block.get('type') if isinstance(block, dict) else None)

            # Skip thinking blocks within content lists
            if block_cls == 'ThinkingBlock' or block_type == 'thinking':
                continue

            if block_cls == 'TextBlock' or block_type == 'text':
                text = getattr(block, 'text', None) or (block.get('text') if isinstance(block, dict) else str(block))
                if text:
                    yield {"type": "text", "text": text, "role": role}

            elif block_cls == 'ToolUseBlock' or block_type == 'tool_use':
                name = getattr(block, 'name', None) or (block.get('name') if isinstance(block, dict) else 'unknown')
                tool_input = getattr(block, 'input', None) or (block.get('input') if isinstance(block, dict) else {})
                tool_id = getattr(block, 'id', None) or (block.get('id') if isinstance(block, dict) else '')
                yield {
                    "type": "tool_use",
                    "tool_name": name,
                    "tool_input": _truncate_input(tool_input),
                    "tool_id": tool_id,
                }

            elif block_cls == 'ToolResultBlock' or block_type == 'tool_result':
                tool_id = getattr(block, 'tool_use_id', None) or (block.get('tool_use_id') if isinstance(block, dict) else '')
                result_content = getattr(block, 'content', None) or (block.get('content') if isinstance(block, dict) else '')
                is_error = getattr(block, 'is_error', False) or (block.get('is_error', False) if isinstance(block, dict) else False)
                yield {
                    "type": "tool_result",
                    "tool_id": tool_id,
                    "result": _truncate_text(str(result_content), 2000),
                    "is_error": is_error,
                }

            else:
                # Unknown block type — skip if empty, else emit
                text = str(block)
                if text and len(text) < 500:
                    yield {"type": "text", "text": text, "role": role}


def _truncate_input(tool_input: Any, max_len: int = 500) -> dict:
    """Truncate tool input values for display."""
    if not isinstance(tool_input, dict):
        return {"value": _truncate_text(str(tool_input), max_len)}
    result = {}
    for k, v in tool_input.items():
        s = str(v)
        result[k] = s if len(s) <= max_len else s[:max_len] + "..."
    return result


def _truncate_text(text: str, max_len: int = 2000) -> str:
    """Truncate text with ellipsis indicator."""
    if len(text) <= max_len:
        return text
    return text[:max_len] + f"\n... ({len(text) - max_len} chars truncated)"
