"""
cvc.agent.subagent — Sub-agent architecture (Claude Code-style).

Enables the parent agent to spawn lightweight child agents with:
  - Isolated context windows (no shared history with parent)
  - Restricted tool sets (e.g. read-only for Explore)
  - Configurable max turns and model selection
  - Single-string result returned to the parent

Built-in agents:
  - Explore: Fast read-only codebase exploration (cheapest model, read-only tools)
  - Plan:    Analyze and propose without modifying (current model, read-only tools)

Custom agents can be defined in `.cvc/agents/<name>/agent.md` with YAML frontmatter.
"""

from __future__ import annotations

import logging
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger("cvc.agent.subagent")

# Read-only tools that sub-agents with restricted access can use
READ_ONLY_TOOLS = frozenset({
    "read_file", "grep", "glob", "list_dir",
    "cvc_status", "cvc_log", "cvc_search", "cvc_smart_search", "cvc_diff",
    "cvc_document_search", "cvc_list_documents",
})

# Full coding tools — read + write + execute
CODING_TOOLS = frozenset({
    "read_file", "write_file", "edit_file", "patch_file",
    "bash", "grep", "glob", "list_dir",
    "cvc_status", "cvc_log", "cvc_commit", "cvc_branch",
    "cvc_search", "cvc_smart_search", "cvc_diff",
    "web_search",
})

# Research tools — read + web
RESEARCH_TOOLS = frozenset({
    "read_file", "grep", "glob", "list_dir",
    "bash", "web_search",
    "cvc_status", "cvc_log", "cvc_search", "cvc_smart_search", "cvc_diff",
    "cvc_document_search", "cvc_list_documents",
})

# Orchestration tools — everything including sub-agents
ORCHESTRATION_TOOLS = frozenset({
    "read_file", "write_file", "edit_file", "patch_file",
    "bash", "grep", "glob", "list_dir",
    "web_search", "agent",
    "cvc_status", "cvc_log", "cvc_commit", "cvc_branch",
    "cvc_restore", "cvc_merge",
    "cvc_search", "cvc_smart_search", "cvc_diff",
    "cvc_ingest_document", "cvc_document_search", "cvc_list_documents",
})


@dataclass
class SubAgentConfig:
    """Configuration for a sub-agent."""
    name: str
    description: str = ""
    model: str = ""            # Empty = use parent's model
    tools: frozenset[str] = field(default_factory=lambda: READ_ONLY_TOOLS)
    max_turns: int = 10
    system_prompt: str = ""    # Additional instructions prepended to system prompt
    custom_instructions: str = ""  # From agent.md file


# ── Built-in agent configs ──────────────────────────────────────────────────

BUILTIN_AGENTS: dict[str, SubAgentConfig] = {
    "Explore": SubAgentConfig(
        name="Explore",
        description=(
            "Fast read-only codebase exploration. Uses cheapest model tier "
            "with only read-only tools (read_file, grep, glob, list_dir, "
            "cvc_search). Returns a summary of findings."
        ),
        tools=READ_ONLY_TOOLS,
        max_turns=10,
        system_prompt=(
            "You are a fast codebase explorer. Your job is to search and read files "
            "to answer the user's question. You CANNOT modify files — you only have "
            "read-only access. Be thorough but concise. Return your findings as a "
            "clear summary."
        ),
    ),
    "Plan": SubAgentConfig(
        name="Plan",
        description=(
            "Analyze and produce a plan without modifying anything. Uses the "
            "parent model with read-only tools. Returns an analysis and "
            "step-by-step plan."
        ),
        tools=READ_ONLY_TOOLS,
        max_turns=15,
        system_prompt=(
            "You are a planning agent. Analyze the codebase and produce a detailed "
            "step-by-step implementation plan. You CANNOT modify files — you only "
            "have read-only access. Return a structured plan with file paths, "
            "specific changes needed, and ordering considerations."
        ),
    ),
    "Security": SubAgentConfig(
        name="Security",
        description=(
            "Code security analyst. Scans for vulnerabilities (OWASP Top 10, "
            "CVEs, secrets, insecure patterns), performs dependency audits, "
            "and suggests hardening fixes."
        ),
        tools=RESEARCH_TOOLS,
        max_turns=20,
        system_prompt=(
            "You are an elite application security engineer. Your mission is to "
            "find and fix security vulnerabilities in the codebase.\n\n"
            "What you do:\n"
            "- Scan source code for OWASP Top 10 vulnerabilities (injection, XSS, "
            "broken auth, SSRF, insecure deserialization, etc.)\n"
            "- Detect hardcoded secrets, API keys, tokens, and credentials\n"
            "- Audit dependencies for known CVEs using web_search\n"
            "- Check for insecure cryptographic practices\n"
            "- Review access control and authentication logic\n"
            "- Identify command injection, path traversal, and IDOR risks\n\n"
            "Output a prioritized report: CRITICAL → HIGH → MEDIUM → LOW findings, "
            "each with file path, line number, description, impact, and recommended fix."
        ),
    ),
    "AI": SubAgentConfig(
        name="AI",
        description=(
            "AI/ML engineering specialist. Designs and implements AI pipelines, "
            "LLM integrations, embeddings, RAG systems, model training, and "
            "inference optimization in Python."
        ),
        tools=CODING_TOOLS,
        max_turns=25,
        system_prompt=(
            "You are a world-class AI/ML engineer specializing in Python. You build "
            "production-grade AI systems.\n\n"
            "Your expertise:\n"
            "- LLM integration (OpenAI, Anthropic, Google, local models via Ollama)\n"
            "- RAG pipelines (embeddings, vector stores, retrieval, reranking)\n"
            "- Prompt engineering and chain-of-thought design\n"
            "- Model fine-tuning and evaluation\n"
            "- ML pipelines (PyTorch, scikit-learn, transformers, langchain)\n"
            "- Data preprocessing, feature engineering, tokenization\n"
            "- Inference optimization (quantization, batching, caching)\n"
            "- Agent architectures and tool-use patterns\n\n"
            "Write clean, well-structured Python code. Use type hints. Follow "
            "best practices for error handling, logging, and configuration. "
            "Prefer async patterns for I/O-bound work."
        ),
    ),
    "UI": SubAgentConfig(
        name="UI",
        description=(
            "UI/UX design and frontend engineering specialist. Creates beautiful, "
            "accessible interfaces — web (React, Vue, HTML/CSS/JS), terminal TUI, "
            "and CLI experiences."
        ),
        tools=CODING_TOOLS,
        max_turns=20,
        system_prompt=(
            "You are the world's most advanced UI/UX developer. You create beautiful, "
            "intuitive, and accessible user interfaces.\n\n"
            "Your expertise:\n"
            "- Frontend frameworks: React, Next.js, Vue, Svelte, Angular\n"
            "- Styling: CSS, Tailwind, styled-components, animations, responsive design\n"
            "- Terminal UIs: Rich, Textual, prompt_toolkit, curses, chalk\n"
            "- CLI design: argument parsing, help text, progress bars, color output\n"
            "- Accessibility: WCAG compliance, ARIA, keyboard navigation, screen readers\n"
            "- Design systems: component architecture, theming, dark/light modes\n"
            "- Performance: lazy loading, code splitting, rendering optimization\n\n"
            "Prioritize user experience above all. Every interaction should feel "
            "smooth and intuitive. Use modern patterns, follow platform conventions, "
            "and ensure responsive, accessible designs."
        ),
    ),
    "Data": SubAgentConfig(
        name="Data",
        description=(
            "Data analytics and engineering specialist. Analyzes datasets, builds "
            "ETL pipelines, creates visualizations, writes SQL, and performs "
            "statistical analysis."
        ),
        tools=CODING_TOOLS,
        max_turns=20,
        system_prompt=(
            "You are an expert data analyst and data engineer.\n\n"
            "Your expertise:\n"
            "- Data analysis: pandas, numpy, scipy, statistical modeling\n"
            "- Visualization: matplotlib, seaborn, plotly, D3.js\n"
            "- SQL: complex queries, window functions, CTEs, query optimization\n"
            "- ETL/ELT: data pipelines, transformation, cleaning, validation\n"
            "- Databases: PostgreSQL, SQLite, MongoDB, Redis\n"
            "- Big data: Spark, Dask, Arrow, Parquet, data streaming\n"
            "- Reporting: automated dashboards, summary statistics, insights\n\n"
            "Always validate data quality first. Handle missing values, outliers, "
            "and type mismatches explicitly. Provide clear interpretations alongside "
            "code and visualizations."
        ),
    ),
    "Orchestrator": SubAgentConfig(
        name="Orchestrator",
        description=(
            "Meta-agent that coordinates other specialized agents to tackle "
            "complex multi-domain tasks. Decomposes problems, delegates to "
            "Security/AI/UI/Data/Explore agents, and synthesizes results."
        ),
        tools=ORCHESTRATION_TOOLS,
        max_turns=30,
        system_prompt=(
            "You are a senior engineering orchestrator. You coordinate complex tasks "
            "by delegating to specialized sub-agents and synthesizing their results.\n\n"
            "Available agents you can invoke with the 'agent' tool:\n"
            "- Explore: Fast read-only codebase exploration and search\n"
            "- Plan: Analyze codebase and produce implementation plans\n"
            "- Security: Scan for vulnerabilities and security issues\n"
            "- AI: AI/ML engineering (Python, LLMs, RAG, training)\n"
            "- UI: UI/UX design and frontend development\n"
            "- Data: Data analysis, SQL, ETL, visualization\n\n"
            "Your workflow:\n"
            "1. Analyze the request and identify which domains are involved\n"
            "2. Use Explore/Plan first to understand the codebase\n"
            "3. Delegate implementation to the appropriate specialist agent(s)\n"
            "4. Review and integrate results\n"
            "5. Verify the combined output works correctly\n\n"
            "You can also directly read/write files and run commands when needed. "
            "Think step-by-step about the best decomposition strategy."
        ),
    ),
}


def _get_cheap_model(provider: str, parent_model: str) -> str:
    """Pick the cheapest model for the given provider (for Explore agent)."""
    cheap_models = {
        "anthropic": "claude-haiku-4-5",
        "openai": "gpt-5-mini",
        "google": "gemini-3-flash-preview",
        "ollama": parent_model,  # Use whatever's loaded locally
        "lmstudio": parent_model,
    }
    return cheap_models.get(provider, parent_model)


def load_custom_agents(workspace: Path) -> dict[str, SubAgentConfig]:
    """Load custom agent definitions from .cvc/agents/<name>/agent.md."""
    agents_dir = workspace / ".cvc" / "agents"
    if not agents_dir.is_dir():
        return {}

    custom: dict[str, SubAgentConfig] = {}

    for agent_dir in agents_dir.iterdir():
        if not agent_dir.is_dir():
            continue

        agent_file = agent_dir / "agent.md"
        if not agent_file.exists():
            continue

        try:
            text = agent_file.read_text(encoding="utf-8")
            config = _parse_agent_md(agent_dir.name, text)
            if config:
                custom[config.name] = config
        except Exception as e:
            logger.warning("Failed to load agent %s: %s", agent_dir.name, e)

    return custom


def _parse_agent_md(name: str, text: str) -> SubAgentConfig | None:
    """Parse an agent.md file with optional YAML frontmatter."""
    frontmatter: dict[str, Any] = {}
    body = text

    # Extract YAML frontmatter (between --- markers)
    fm_match = re.match(r"^---\s*\n(.*?)\n---\s*\n", text, re.DOTALL)
    if fm_match:
        body = text[fm_match.end():]
        try:
            import yaml
            frontmatter = yaml.safe_load(fm_match.group(1)) or {}
        except ImportError:
            # Minimal YAML parsing fallback
            for line in fm_match.group(1).splitlines():
                if ":" in line:
                    k, v = line.split(":", 1)
                    k, v = k.strip(), v.strip()
                    if v.startswith("[") and v.endswith("]"):
                        v = [x.strip().strip("'\"") for x in v[1:-1].split(",")]
                    frontmatter[k] = v

    # Build tool set from frontmatter
    tools_list = frontmatter.get("tools", None)
    if isinstance(tools_list, list):
        tools = frozenset(tools_list)
    else:
        tools = READ_ONLY_TOOLS

    return SubAgentConfig(
        name=frontmatter.get("name", name),
        description=frontmatter.get("description", ""),
        model=frontmatter.get("model", ""),
        tools=tools,
        max_turns=int(frontmatter.get("max_turns", 10)),
        system_prompt=frontmatter.get("system_prompt", ""),
        custom_instructions=body.strip(),
    )


def get_available_agents(workspace: Path) -> dict[str, SubAgentConfig]:
    """Return all available agents (built-in + custom)."""
    agents = dict(BUILTIN_AGENTS)
    agents.update(load_custom_agents(workspace))
    return agents


class SubAgent:
    """
    Run an isolated sub-agent session and return its result.

    The sub-agent gets:
    - Its own context window (no shared history)
    - A restricted tool set
    - A configurable turn limit
    - A specialized system prompt
    """

    def __init__(
        self,
        config: SubAgentConfig,
        workspace: Path,
        provider: str,
        api_key: str,
        parent_model: str,
        base_url: str = "",
    ) -> None:
        self.config = config
        self.workspace = workspace
        self.provider = provider
        self.api_key = api_key
        self.parent_model = parent_model
        self.base_url = base_url

    async def run(self, prompt: str) -> str:
        """Execute the sub-agent and return its final response."""
        from cvc.agent.llm import AgentLLM
        from cvc.agent.tools import AGENT_TOOLS, get_tools_for_provider

        # Pick model
        model = self.config.model
        if not model:
            if self.config.name == "Explore":
                model = _get_cheap_model(self.provider, self.parent_model)
            else:
                model = self.parent_model

        # Build LLM client for this sub-agent
        llm = AgentLLM(
            provider=self.provider,
            api_key=self.api_key,
            model=model,
            base_url=self.base_url,
        )

        try:
            # Filter tools to only those allowed by the config
            allowed_tools = [
                t for t in AGENT_TOOLS
                if t["function"]["name"] in self.config.tools
            ]
            provider_tools = get_tools_for_provider(self.provider, allowed_tools)

            # Build system prompt
            sys_parts = []
            if self.config.system_prompt:
                sys_parts.append(self.config.system_prompt)
            if self.config.custom_instructions:
                sys_parts.append(self.config.custom_instructions)
            sys_parts.append(
                f"\nWorkspace: {self.workspace}\n"
                f"You are the '{self.config.name}' sub-agent. Complete the task and "
                "provide a clear, concise summary of your findings."
            )

            messages: list[dict[str, Any]] = [
                {"role": "system", "content": "\n\n".join(sys_parts)},
                {"role": "user", "content": prompt},
            ]

            # Build a minimal executor for the sub-agent
            from cvc.core.models import CVCConfig
            from cvc.core.database import ContextDatabase
            from cvc.operations.engine import CVCEngine
            from cvc.agent.permissions import PermissionEngine

            config = CVCConfig.for_project(
                project_root=self.workspace,
                provider=self.provider,
                model=model,
                mode="cli",
            )
            config.ensure_dirs()
            db = ContextDatabase(config)
            engine = CVCEngine(config, db)

            # Sub-agent permission engine: only allow configured tools
            perm = PermissionEngine()
            from cvc.agent.permissions import PermissionRule
            for tool_name in self.config.tools:
                perm.add_rule(PermissionRule.parse(tool_name, "allow"))

            from cvc.agent.executor import ToolExecutor
            executor = ToolExecutor(self.workspace, engine, permission_engine=perm)

            # Agentic loop (simplified — no streaming, no commit)
            final_text = ""
            for turn in range(self.config.max_turns):
                response_text = ""
                tool_calls = []

                # Collect full response
                async for event in llm.stream_chat(
                    messages=messages,
                    tools=provider_tools,
                ):
                    if event.type == "text_delta":
                        response_text += event.text
                    elif event.type == "tool_call_start":
                        tool_calls.append(event.tool_call)
                    elif event.type == "tool_call_delta":
                        if tool_calls:
                            idx = event.tool_call_index
                            if idx < len(tool_calls):
                                tool_calls[idx].arguments += event.args_delta

                if not tool_calls:
                    final_text = response_text
                    break

                # Execute tool calls
                messages.append({
                    "role": "assistant",
                    "content": response_text,
                    "tool_calls": [tc.__dict__ for tc in tool_calls],
                })

                for tc in tool_calls:
                    try:
                        args = tc.arguments
                        if isinstance(args, str):
                            import json
                            args = json.loads(args)
                        result = executor.execute(tc.name, args)
                    except Exception as e:
                        result = f"Error: {e}"

                    messages.append({
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "content": result,
                    })

            return final_text or "(Sub-agent completed without producing a response)"

        finally:
            await llm.close()
