"""DynamicToolRegistry — Gateway Mode for Cross-LLM Compatibility.

Reduces visible tool count to ~18 "gateway" tools + discover_tools() + run_tool(),
while keeping all 257 tools fully executable behind the scenes. Activated by setting
DATABRIDGE_TOOL_MODE=dynamic (default is "full" for backward compatibility).
"""

import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

ALWAYS_VISIBLE: Set[str] = {
    # Core data ops
    "load_csv", "load_json", "profile_data", "find_files",
    # Hierarchy (core product)
    "create_hierarchy", "import_flexible_hierarchy", "export_hierarchy_csv",
    # Reconciliation
    "fuzzy_match_columns", "compare_hashes",
    # AI / workflow
    "submit_orchestrated_task", "get_smart_recommendations",
    # Gateway tools (registered by server.py)
    "discover_tools", "run_tool", "search_tools", "list_domains",
    # System
    "get_license_status", "get_application_documentation",
}

SLIM_VISIBLE: Set[str] = {
    # 12 slim workflow tools
    "onboard_data", "reconcile", "search_knowledge", "review_knowledge",
    "discover_model", "generate_dbt", "assess_quality", "resolve_entities",
    "extract_documents", "run_workflow", "generate_sql", "agent_communicate",
    # System / gateway
    "discover_tools", "run_tool", "get_license_status",
}

PREFIX_DOMAIN_MAP: Dict[str, str] = {
    "blce_": "blce",
    "model_": "blce",
    "catalog_": "catalog",
    "cortex_": "cortex",
    "analyst_": "cortex",
    "wright_": "wright",
    "obs_": "observe",
    "version_": "versioning",
    "rag_": "graphrag",
    "ai_": "ai_discovery",
    "dbt_": "dbt",
    "git_": "git",
    "github_": "git",
    "diff_": "diff",
    "hierarchy_": "hierarchy",
    "smart_": "recommendations",
    "plan_": "planner",
    "suggest_": "planner",
    "track_": "lineage",
    "analyze_change_": "lineage",
    "snowflake_": "snowflake",
}

DOMAIN_DESCRIPTIONS: Dict[str, str] = {
    "blce": "Business Logic Comprehension Engine — SQL parsing, normalization, cross-referencing, DDL execution",
    "catalog": "Data Catalog — metadata scanning, search, lineage, glossary",
    "cortex": "Cortex AI — text generation, reasoning, semantic models, natural language SQL",
    "wright": "Wright Pipeline — Snowflake Dynamic Table pipeline generation from hierarchies",
    "observe": "Data Observability — metrics, alerts, asset health monitoring",
    "versioning": "Data Versioning — snapshots, diffs, rollbacks",
    "graphrag": "GraphRAG Engine — vector search, entity extraction, context retrieval",
    "ai_discovery": "AI Relationship Discovery — schema analysis, naming patterns, FK detection",
    "dbt": "dbt Integration — project creation, model generation",
    "git": "Git & CI/CD — commits, PRs, deployment automation",
    "diff": "Diff Utilities — text, dict, list comparison and explanation",
    "hierarchy": "Hierarchy Builder — multi-level hierarchy management, graph bridge",
    "recommendations": "Smart Recommendations — AI-powered data import and workflow suggestions",
    "planner": "Planner Agent — workflow planning, agent suggestions",
    "lineage": "Lineage & Impact — column lineage tracking, change impact analysis",
    "reconciliation": "Data Reconciliation — fuzzy matching, hash comparison, conflict detection",
    "quality": "Data Quality — expectation suites, validation rules",
    "core": "Core Tools — file discovery, data loading, profiling",
    "templates": "Templates, Skills & Knowledge Base",
    "orchestrator": "AI Orchestrator — task submission, agent registration",
    "datashield": "DataShield — data masking and privacy",
    "unified_agent": "Unified AI Agent — book/librarian checkout and sync",
    "snowflake": "Snowflake-Specific Tooling — lineage extraction, dependency evidence, DDL/sample bundle exports",
}


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class ToolMetadata:
    """Metadata for a single tool."""
    name: str
    domain: str
    description: str
    tier: str = "CE"
    plugin: str = "core"


# ---------------------------------------------------------------------------
# DynamicToolRegistry
# ---------------------------------------------------------------------------

class DynamicToolRegistry:
    """Registry that hides tools from FastMCP's visible list while keeping them executable.

    Modes:
        "full"    — backward-compatible; all tools stay visible in FastMCP (default)
        "dynamic" — only ALWAYS_VISIBLE tools remain; rest are hidden but callable via run_tool()
    """

    def __init__(self, mcp: Any, mode: str = "full"):
        self._mcp = mcp
        self.mode = mode.lower().strip()
        self._hidden_tools: Dict[str, Any] = {}          # name → FunctionTool (popped from FastMCP)
        self._metadata: Dict[str, ToolMetadata] = {}      # name → ToolMetadata
        self._domain_map: Dict[str, List[str]] = {}       # domain → [tool_names]
        self._current_plugin: Tuple[str, str] = ("core", "CE")
        self._seen_tools: Set[str] = set()                # track tools already captured

    # ------------------------------------------------------------------
    # Plugin lifecycle hooks
    # ------------------------------------------------------------------

    def set_current_plugin(self, name: str, tier: str) -> None:
        """Set the current plugin context (called before each plugin loads)."""
        self._current_plugin = (name, tier)

    def capture_tool_metadata(self) -> None:
        """Capture metadata for any new tools registered since last call.

        Called after each plugin's register function completes. Compares the
        current FastMCP tool list against _seen_tools to find newly added ones.
        """
        current_tools = set(self._mcp._tool_manager._tools.keys())
        new_tools = current_tools - self._seen_tools

        plugin_name, tier = self._current_plugin

        for tool_name in new_tools:
            tool_obj = self._mcp._tool_manager._tools[tool_name]
            domain = self._infer_domain(tool_name, plugin_name)
            description = getattr(tool_obj, "description", "") or ""

            meta = ToolMetadata(
                name=tool_name,
                domain=domain,
                description=description,
                tier=tier,
                plugin=plugin_name,
            )
            self._metadata[tool_name] = meta

            # Update domain map
            if domain not in self._domain_map:
                self._domain_map[domain] = []
            self._domain_map[domain].append(tool_name)

        self._seen_tools = current_tools.copy()

    # ------------------------------------------------------------------
    # Finalize — hide tools if in dynamic mode
    # ------------------------------------------------------------------

    def finalize(self) -> None:
        """Called after ALL plugins are loaded.

        Modes:
            "full"    — all tools stay visible (default)
            "dynamic" — only ALWAYS_VISIBLE tools remain visible
            "slim"    — only SLIM_VISIBLE tools remain visible (~15 tools)
        """
        if self.mode == "full":
            return

        if self.mode == "slim":
            visible_set = SLIM_VISIBLE
        else:
            visible_set = ALWAYS_VISIBLE

        tools_dict = self._mcp._tool_manager._tools
        to_hide = [name for name in list(tools_dict.keys()) if name not in visible_set]

        for name in to_hide:
            self._hidden_tools[name] = tools_dict.pop(name)

        visible_count = len(tools_dict)
        hidden_count = len(self._hidden_tools)
        print(f"[Gateway] {self.mode} mode: {visible_count} visible, {hidden_count} hidden across {len(self._domain_map)} domains")

    # ------------------------------------------------------------------
    # Query methods
    # ------------------------------------------------------------------

    def get_tools_by_domain(self, domain: str) -> List[Dict[str, str]]:
        """Return tool info for a given domain."""
        domain = domain.lower().strip()
        tool_names = self._domain_map.get(domain, [])
        results = []
        for name in sorted(tool_names):
            meta = self._metadata.get(name)
            if meta:
                results.append({
                    "name": meta.name,
                    "description": meta.description,
                    "domain": meta.domain,
                    "tier": meta.tier,
                })
        return results

    def get_all_domains(self) -> List[Dict[str, Any]]:
        """Return all domains with tool counts and descriptions."""
        domains = []
        for domain, tool_names in sorted(self._domain_map.items()):
            domains.append({
                "domain": domain,
                "tool_count": len(tool_names),
                "description": self._get_domain_description(domain),
            })
        return domains

    def search_tools(self, query: str) -> List[Dict[str, str]]:
        """Fuzzy search across all tool names and descriptions."""
        query_lower = query.lower().strip()
        if not query_lower:
            return []

        # Split query into tokens for multi-word matching
        tokens = query_lower.split()
        results = []

        for name, meta in self._metadata.items():
            searchable = f"{name} {meta.description} {meta.domain}".lower()
            # Score: count how many tokens match
            score = sum(1 for token in tokens if token in searchable)
            if score > 0:
                # Bonus for exact name match
                if query_lower in name.lower():
                    score += 5
                results.append((score, {
                    "name": meta.name,
                    "description": meta.description,
                    "domain": meta.domain,
                    "tier": meta.tier,
                }))

        # Sort by score descending, then by name
        results.sort(key=lambda x: (-x[0], x[1]["name"]))
        return [r[1] for r in results[:30]]

    # ------------------------------------------------------------------
    # Tool execution
    # ------------------------------------------------------------------

    async def execute_tool(self, name: str, arguments: dict) -> Any:
        """Execute a tool by name, whether hidden or visible.

        Args:
            name: The tool name to execute.
            arguments: Dict of arguments to pass to the tool.

        Returns:
            The tool's result content as a string.

        Raises:
            ValueError: If the tool is not found.
        """
        # 1. Check hidden tools first (dynamic mode)
        if name in self._hidden_tools:
            tool = self._hidden_tools[name]
            result = await tool.run(arguments)
            # Extract text from ToolResult content list
            return self._extract_result_text(result)

        # 2. Check visible tools (still in FastMCP)
        if name in self._mcp._tool_manager._tools:
            tool = self._mcp._tool_manager._tools[name]
            result = await tool.run(arguments)
            return self._extract_result_text(result)

        raise ValueError(
            f"Tool '{name}' not found. Use discover_tools() to see available tools."
        )

    # ------------------------------------------------------------------
    # Metadata access for UI
    # ------------------------------------------------------------------

    def get_tool_metadata(self, name: str) -> Optional[ToolMetadata]:
        """Get metadata for a specific tool."""
        return self._metadata.get(name)

    def get_all_metadata(self) -> Dict[str, ToolMetadata]:
        """Get metadata for all tools."""
        return dict(self._metadata)

    @property
    def hidden_count(self) -> int:
        return len(self._hidden_tools)

    @property
    def visible_count(self) -> int:
        return len(self._mcp._tool_manager._tools)

    @property
    def total_count(self) -> int:
        return len(self._metadata)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _infer_domain(tool_name: str, plugin_name: str) -> str:
        """Infer domain from tool name prefix, falling back to plugin name."""
        for prefix, domain in PREFIX_DOMAIN_MAP.items():
            if tool_name.startswith(prefix):
                return domain
        # Fallback to plugin name
        return plugin_name.lower().replace(" ", "_")

    @staticmethod
    def _get_domain_description(domain: str) -> str:
        """Return a stable, human-friendly domain description."""
        explicit = DOMAIN_DESCRIPTIONS.get(domain)
        if explicit and explicit.strip():
            return explicit.strip()
        label = domain.replace("_", " ").strip().title()
        if not label:
            label = "General"
        return f"{label} Tooling - domain-specific operations and workflows"

    @staticmethod
    def _extract_result_text(result: Any) -> str:
        """Extract text content from a FastMCP ToolResult."""
        if result is None:
            return ""
        content = getattr(result, "content", None)
        if content and isinstance(content, list):
            texts = []
            for item in content:
                text = getattr(item, "text", None)
                if text:
                    texts.append(text)
            return "\n".join(texts) if texts else str(result)
        return str(result)
