"""MCP tool registration for the Admin plugin (Enterprise tier).

Registers 5 tools: 3 Vanna RAG + 2 PydanticAI agents.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Dict

logger = logging.getLogger(__name__)


def register_admin_tools(mcp, settings) -> Dict[str, Any]:
    """Register Vanna RAG and PydanticAI agent tools.

    Args:
        mcp: FastMCP server instance.
        settings: Application settings.

    Returns:
        Dict with registration summary.
    """
    # Import trust helpers from the main server module
    try:
        from src.server import _coerce_trust_attestation_str, _enforce_trust_for_ai_tool
    except ImportError:
        def _coerce_trust_attestation_str(explicit_attestation="", **kw):
            return explicit_attestation or ""
        def _enforce_trust_for_ai_tool(*, tool_name, lane, trust_attestation=""):
            return None

    # Import log_action helper
    try:
        from src.server import log_action
    except ImportError:
        def log_action(actor, action, detail):
            logger.info("[%s] %s: %s", actor, action, detail)

    # Vanna client (optional)
    try:
        from src.vanna_client import get_vanna_client, VANNA_AVAILABLE
        _vanna = get_vanna_client() if VANNA_AVAILABLE else None
    except ImportError:
        VANNA_AVAILABLE = False
        _vanna = None

    # PydanticAI agents (optional)
    try:
        from src.agents.pydantic_agents import (
            PYDANTIC_AI_AVAILABLE,
            run_planner_agent as _run_planner,
            run_sql_agent as _run_sql,
        )
    except ImportError:
        PYDANTIC_AI_AVAILABLE = False
        _run_planner = None
        _run_sql = None

    _vanna_mark = "YES" if VANNA_AVAILABLE else "NO"
    _pai_mark = "YES" if PYDANTIC_AI_AVAILABLE else "NO"
    print(f"[Admin] Vanna={_vanna_mark}  PydanticAI={_pai_mark}")

    # ==================================================================
    # VANNA RAG TOOLS (1-3)
    # ==================================================================

    @mcp.tool()
    def vanna_train(
        source_type: str,
        content: str,
        question: str = "",
        sf_account: str = "",
        sf_database: str = "",
        sf_schema: str = "",
    ) -> str:
        """Train the Vanna RAG model on DDL, documentation, Q&A pairs, or auto-extract from Snowflake.

        source_type: One of 'ddl', 'documentation', 'qa', or 'snowflake'.
        content: DDL string, documentation text, or SQL (for qa). Ignored for snowflake.
        question: The original question (required only for source_type='qa').
        sf_account: Snowflake account (required only for source_type='snowflake').
        sf_database: Snowflake database (required only for source_type='snowflake').
        sf_schema: Snowflake schema (required only for source_type='snowflake').
        """
        if not VANNA_AVAILABLE or _vanna is None:
            return json.dumps({"error": "Vanna not installed. Run: pip install 'vanna[chromadb,anthropic]'"})

        if source_type == "ddl":
            result = _vanna.train_from_ddl(content)
        elif source_type == "documentation":
            result = _vanna.train_from_documentation(content)
        elif source_type == "qa":
            if not question:
                return json.dumps({"error": "question is required for source_type='qa'"})
            result = _vanna.train_from_query_history(content, question)
        elif source_type == "snowflake":
            if not sf_account or not sf_database or not sf_schema:
                return json.dumps({"error": "sf_account, sf_database, sf_schema required for snowflake"})
            result = _vanna.train_from_snowflake(sf_account, sf_database, sf_schema)
        else:
            return json.dumps({"error": f"Unknown source_type: {source_type}. Use ddl/documentation/qa/snowflake."})

        log_action("AI_AGENT", "vanna_train", f"Trained Vanna ({source_type})")
        return json.dumps(result, indent=2, default=str)

    @mcp.tool()
    def vanna_ask(
        question: str,
    ) -> str:
        """Generate SQL from a natural-language question using Vanna RAG.

        Requires prior training via vanna_train().

        question: Natural language question (e.g. "total revenue by month").
        """
        if not VANNA_AVAILABLE or _vanna is None:
            return json.dumps({"error": "Vanna not installed. Run: pip install 'vanna[chromadb,anthropic]'"})

        if not _vanna.is_trained():
            return json.dumps({"error": "Vanna not trained yet. Use vanna_train() first."})

        result = _vanna.ask(question)
        if result is None:
            return json.dumps({"error": "Could not generate SQL for this question."})

        log_action("AI_AGENT", "vanna_ask", f"Generated SQL for: {question[:80]}")
        return json.dumps(result, indent=2, default=str)

    @mcp.tool()
    def vanna_status() -> str:
        """Get Vanna RAG training data stats and readiness status."""
        if not VANNA_AVAILABLE or _vanna is None:
            return json.dumps({"available": False, "reason": "Vanna not installed"})

        return json.dumps(_vanna.get_training_status(), indent=2, default=str)

    # ==================================================================
    # PYDANTIC AI AGENT TOOLS (4-5)
    # ==================================================================

    @mcp.tool()
    async def ai_plan_workflow(
        request: str,
        context: str = "",
        trust_attestation: str = "",
        _trust_attestation: str = "",
    ) -> str:
        """Plan a workflow using PydanticAI multi-step reasoning agent.

        Enhanced version of plan_workflow that uses iterative tool calls to
        validate the plan before returning it.

        request: Natural language description of the desired workflow.
        context: Optional JSON string with additional context.
        """
        trust_attestation = _coerce_trust_attestation_str(
            explicit_attestation=trust_attestation or _trust_attestation,
        )
        trust_error = _enforce_trust_for_ai_tool(
            tool_name="ai_plan_workflow",
            lane="databridge_ai_masked",
            trust_attestation=trust_attestation,
        )
        if trust_error:
            return json.dumps(trust_error)

        if not PYDANTIC_AI_AVAILABLE or _run_planner is None:
            return json.dumps({"error": "pydantic-ai not installed. Run: pip install pydantic-ai"})

        ctx_dict = None
        if context:
            try:
                ctx_dict = json.loads(context)
            except json.JSONDecodeError:
                ctx_dict = {"raw_context": context}

        # Gather available agents
        try:
            from src.agents.planner_agent import DEFAULT_AGENTS
            agents_dicts = [a.to_dict() for a in DEFAULT_AGENTS]
        except ImportError:
            agents_dicts = []

        plan_output = await _run_planner(request, agents_dicts, ctx_dict)
        if plan_output is None:
            return json.dumps({"error": "PydanticAI planner returned no result — check ANTHROPIC_API_KEY"})

        log_action("AI_AGENT", "ai_plan_workflow", f"PydanticAI plan: {plan_output.name}")
        return json.dumps(plan_output.model_dump(), indent=2, default=str)

    @mcp.tool()
    async def ai_generate_sql(
        question: str,
        tables: str = "",
        trust_attestation: str = "",
        _trust_attestation: str = "",
    ) -> str:
        """Generate SQL using PydanticAI SQL agent (combines Vanna + DuckDB + schema context).

        question: Natural language question about the data.
        tables: Comma-separated list of table names for context (optional).
        """
        trust_attestation = _coerce_trust_attestation_str(
            explicit_attestation=trust_attestation or _trust_attestation,
        )
        trust_error = _enforce_trust_for_ai_tool(
            tool_name="ai_generate_sql",
            lane="databridge_ai_masked",
            trust_attestation=trust_attestation,
        )
        if trust_error:
            return json.dumps(trust_error)

        if not PYDANTIC_AI_AVAILABLE or _run_sql is None:
            return json.dumps({"error": "pydantic-ai not installed. Run: pip install pydantic-ai"})

        table_list = [t.strip() for t in tables.split(",") if t.strip()] if tables else None

        sql_output = await _run_sql(question, table_list)
        if sql_output is None:
            return json.dumps({"error": "PydanticAI SQL agent returned no result"})

        log_action("AI_AGENT", "ai_generate_sql", f"SQL generated for: {question[:80]}")
        return json.dumps(sql_output.model_dump(), indent=2, default=str)

    # ------------------------------------------------------------------
    logger.info("[Plugin] admin registered 5 tools")
    return {"tools_registered": 5, "plugin": "admin"}
