"""ToolRegistry — tool discovery, schema generation, and execution.

Scans ``{data_dir}/tools/`` for tool group directories (with ``tool.json``)
and provides methods to list, search, and execute tools.

Supports two modes:
- **Group mode**: ``tool.json`` contains a ``tools`` array of sub-tools,
  each with its own handler file.  Tool names use ``{group}.{sub}`` format.
- **Legacy mode**: ``tool.json`` + ``main.py`` defines a single tool
  (backward compatible).
"""

from __future__ import annotations

import importlib.util
import json
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Any, Optional

from oedix.tools.base_tool import BaseTool

logger = logging.getLogger("oedix.tools")


@dataclass
class ToolEntry:
    """Metadata entry for a registered tool."""

    name: str
    title: str = ""
    description: str = ""
    input_schema: dict[str, Any] = field(default_factory=dict)
    version: str = "1.0.0"
    tags: list[str] = field(default_factory=list)
    dir_path: str = ""
    handler_path: str = ""
    source: str = "builtin"  # builtin | mcp
    group: str = ""  # group name (e.g. "storage") if sub-tool
    canvas: Optional[dict[str, Any]] = field(default=None)  # canvas metadata
    storage: Optional[dict[str, Any]] = field(default=None)  # storage integration
    context_text: str = ""  # additional context for LLM
    keywords: list[str] = field(default_factory=list)
    use_cases: list[str] = field(default_factory=list)
    _instance: Optional[BaseTool] = field(default=None, repr=False)

    def to_schema(self) -> dict[str, Any]:
        """Return LLM-compatible tool schema."""
        return {
            "name": self.name,
            "description": self.description,
            "input_schema": self.input_schema or {
                "type": "object",
                "properties": {},
            },
        }

    def to_meta(self) -> dict[str, Any]:
        """Return metadata for search and display."""
        meta: dict[str, Any] = {
            "name": self.name,
            "title": self.title or self.name,
            "description": self.description,
            "version": self.version,
            "tags": self.tags,
            "source": self.source,
        }
        if self.group:
            meta["group"] = self.group
        if self.canvas:
            meta["canvas"] = self.canvas
        if self.storage:
            meta["storage"] = self.storage
        return meta


class ToolRegistry:
    """Discover, manage, and execute tools.

    Usage::

        registry = ToolRegistry("/path/to/data/tools")
        registry.scan()
        schemas = registry.list_schemas()
        result = registry.execute("storage.explorer", action="list", path="/")
    """

    def __init__(self, tools_dir: str) -> None:
        self.tools_dir = tools_dir
        self._entries: dict[str, ToolEntry] = {}
        self._groups: dict[str, dict[str, Any]] = {}
        self._context: dict[str, Any] = {}

    def set_context(self, context: dict[str, Any]) -> None:
        """Set the execution context passed to tools."""
        self._context = context

    def scan(self) -> list[ToolEntry]:
        """Scan the tools directory and register all valid tools.

        Each tool directory is either:
        - **Group**: ``tool.json`` with ``tools`` array → multiple sub-tools
        - **Legacy**: ``tool.json`` + ``main.py`` → single tool
        - **Single file**: ``.py`` with ``name``/``description`` attributes

        Returns:
            List of discovered ToolEntry items.
        """
        self._entries.clear()
        self._groups.clear()
        if not os.path.isdir(self.tools_dir):
            logger.warning("Tools directory not found: %s", self.tools_dir)
            return []

        discovered: list[ToolEntry] = []
        for item_name in sorted(os.listdir(self.tools_dir)):
            item_path = os.path.join(self.tools_dir, item_name)
            if item_name.startswith("_"):
                continue

            entries: list[ToolEntry] = []
            if os.path.isdir(item_path):
                entries = self._scan_dir(item_path, item_name)
            elif item_name.endswith(".py"):
                entry = self._scan_file(item_path, item_name)
                if entry is not None:
                    entries = [entry]

            for entry in entries:
                self._entries[entry.name] = entry
                discovered.append(entry)
                # Register legacy/single-file tools as their own group
                group_key = entry.group or entry.name
                if group_key not in self._groups:
                    self._groups[group_key] = {
                        "name": group_key,
                        "title": entry.title,
                        "description": entry.description,
                        "tool_count": 1,
                        "has_help": False,
                    }

        logger.info(
            "ToolRegistry scanned %d tools", len(discovered)
        )
        return discovered

    def list_schemas(self) -> list[dict[str, Any]]:
        """Return LLM tool schemas for all registered tools."""
        if not self._entries:
            self.scan()
        return [e.to_schema() for e in self._entries.values()]

    def list_meta(self) -> list[dict[str, Any]]:
        """Return metadata list for all registered tools."""
        if not self._entries:
            self.scan()
        return [e.to_meta() for e in self._entries.values()]

    def get(self, name: str) -> Optional[ToolEntry]:
        """Get a tool entry by name."""
        if not self._entries:
            self.scan()
        return self._entries.get(name)

    def get_schema(self, name: str) -> Optional[dict[str, Any]]:
        """Get the schema for a specific tool."""
        entry = self.get(name)
        return entry.to_schema() if entry else None

    def execute(
        self,
        tool_name: str,
        params: Optional[dict[str, Any]] = None,
        context: Optional[dict] = None,
    ) -> dict[str, Any]:
        """Execute a tool by name.

        Args:
            tool_name: Tool name (e.g. ``storage.explorer``).
            params: Parameters to pass to the tool's run() method.
            context: Optional execution context override.

        Returns:
            Tool result dict (``{"text": ..., "data": ..., "canvas": ...}``).

        Raises:
            KeyError: If tool not found.
            Exception: Propagated from tool execution.
        """
        entry = self.get(tool_name)
        if entry is None:
            raise KeyError(f"Tool '{tool_name}' not found in registry")

        instance = self._get_instance(entry, context or self._context)
        # Pass sub-tool action name (e.g. "export" from "platform_data.export")
        run_kwargs = dict(params or {})
        if "." in tool_name:
            run_kwargs.setdefault("_action", tool_name.rsplit(".", 1)[-1])
        result = instance.run(**run_kwargs)

        # Normalize result
        if isinstance(result, str):
            return {"text": result}
        if not isinstance(result, dict):
            return {"text": str(result)}
        if "text" not in result:
            result["text"] = str(result.get("data", ""))

        # Attach canvas metadata from entry if tool returned canvas=True
        if result.get("canvas") is True and entry.canvas:
            result["canvas"] = entry.canvas

        return result

    def names(self) -> list[str]:
        """Return list of registered tool names."""
        if not self._entries:
            self.scan()
        return list(self._entries.keys())

    def list_groups(self) -> list[dict[str, Any]]:
        """Return group-level summaries for hierarchical tool selection.

        Each group entry contains:
        - name: group directory name (e.g. "storage", "pdf")
        - title: human-readable title
        - description: group-level description
        - tool_count: number of sub-tools (excluding help)
        - has_help: whether a {group}.help tool exists

        Legacy (single-tool) entries are treated as their own group.
        Group metadata is populated during scan(), not re-read from disk.
        """
        if not self._entries:
            self.scan()
        return list(self._groups.values())

    def list_group_tools(self, group_name: str) -> list[dict[str, Any]]:
        """Return schemas for non-help tools within a specific group.

        For legacy (single-tool) groups, returns the tool's schema.
        """
        if not self._entries:
            self.scan()

        schemas = []
        for entry in self._entries.values():
            group_key = entry.group or entry.name
            if group_key == group_name and not entry.name.endswith(".help"):
                schemas.append(entry.to_schema())
        return schemas

    def get_help_tool_name(self, group_name: str) -> Optional[str]:
        """Return the help tool name for a group, or None."""
        help_name = f"{group_name}.help"
        if help_name in self._entries:
            return help_name
        return None

    def search_tools(
        self, query: str, top_n: int = 5
    ) -> list[dict[str, Any]]:
        """Search and rank tools by relevance to query.

        Scores each non-help tool against the query terms using
        keyword, tag, description, use_case, and name matching.

        Args:
            query: Search query string (keywords or natural language).
            top_n: Maximum number of results to return.

        Returns:
            List of tool dicts sorted by relevance score (descending).
            Each dict contains the standard schema fields plus:
            score, title, keywords, use_cases, context, group.
        """
        if not self._entries:
            self.scan()

        terms = self._tokenize(query)
        if not terms:
            return [
                self._tool_search_result(e, 0.0)
                for e in self._entries.values()
                if not e.name.endswith(".help")
            ][:top_n]

        scored: list[tuple[ToolEntry, float]] = []
        for entry in self._entries.values():
            if entry.name.endswith(".help"):
                continue
            score = self._score_tool(entry, terms)
            if score > 0:
                scored.append((entry, score))

        scored.sort(key=lambda x: x[1], reverse=True)
        return [
            self._tool_search_result(entry, score)
            for entry, score in scored[:top_n]
        ]

    def _score_tool(
        self, entry: ToolEntry, terms: list[str]
    ) -> float:
        """Score a tool's relevance to search terms."""
        score = 0.0
        desc_lower = entry.description.lower()

        for term in terms:
            t = term.lower()
            # Keyword match (highest weight)
            for kw in entry.keywords:
                if t in kw.lower() or kw.lower() in t:
                    score += 3.0
                    break
            # Tag match
            for tag in entry.tags:
                if t in tag.lower() or tag.lower() in t:
                    score += 2.0
                    break
            # Use case match
            for uc in entry.use_cases:
                if t in uc.lower():
                    score += 2.0
                    break
            # Description match
            if t in desc_lower:
                score += 1.0
            # Name match
            if t in entry.name.lower():
                score += 2.0
        return score

    def _tool_search_result(
        self, entry: ToolEntry, score: float
    ) -> dict[str, Any]:
        """Build a search result dict for a tool."""
        result = entry.to_schema()
        result["score"] = score
        result["title"] = entry.title
        result["keywords"] = entry.keywords
        result["use_cases"] = entry.use_cases
        result["context"] = entry.context_text
        result["group"] = entry.group
        return result

    @staticmethod
    def _tokenize(query: str) -> list[str]:
        """Tokenize a query into meaningful search terms."""
        import re as _re
        tokens = _re.split(r'[\s,.\-_/\\:;!?(){}\[\]"\'\']+', query)
        return [t for t in tokens if len(t) >= 2]

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _scan_dir(
        self, dir_path: str, dir_name: str
    ) -> list[ToolEntry]:
        """Scan a tool directory.

        If ``tool.json`` contains a ``tools`` array, register each sub-tool.
        Otherwise fall back to legacy single-tool mode (``main.py``).
        """
        meta_path = os.path.join(dir_path, "tool.json")

        # Load metadata
        meta: dict[str, Any] = {}
        if os.path.isfile(meta_path):
            try:
                with open(meta_path, "r", encoding="utf-8") as f:
                    meta = json.load(f)
            except (json.JSONDecodeError, OSError):
                logger.warning("Invalid tool.json in %s", dir_path)

        group_name = meta.get("name", dir_name)
        group_version = meta.get("version", "1.0.0")
        group_tags = meta.get("tags", [])
        group_context = meta.get("context", "")
        group_title = meta.get("title", dir_name)
        group_desc = meta.get("description", "")
        group_keywords = meta.get("keywords", [])

        # --- Group mode: tool.json has "tools" array ---
        sub_tools = meta.get("tools")
        if isinstance(sub_tools, list) and sub_tools:
            entries: list[ToolEntry] = []
            for sub in sub_tools:
                sub_name = sub.get("name", "")
                if not sub_name:
                    continue
                handler_file = sub.get("handler", "")
                if not handler_file:
                    continue
                handler_path = os.path.join(dir_path, handler_file)
                if not os.path.isfile(handler_path):
                    logger.warning(
                        "Handler not found for sub-tool '%s': %s",
                        sub_name, handler_path,
                    )
                    continue
                sub_keywords = sub.get("keywords", [])
                merged_keywords = list(set(group_keywords + sub_keywords))
                entries.append(ToolEntry(
                    name=sub_name,
                    title=sub.get("title", sub_name),
                    description=sub.get("description", ""),
                    input_schema=sub.get("input_schema", {}),
                    version=group_version,
                    tags=sub.get("tags", group_tags),
                    dir_path=dir_path,
                    handler_path=handler_path,
                    group=group_name,
                    canvas=sub.get("canvas"),
                    storage=sub.get("storage"),
                    context_text=sub.get("context", group_context),
                    keywords=merged_keywords,
                    use_cases=sub.get("use_cases", []),
                ))
            if entries:
                # Register group metadata
                self._groups[group_name] = {
                    "name": group_name,
                    "title": group_title,
                    "description": group_desc,
                    "tool_count": sum(
                        1 for e in entries if not e.name.endswith(".help")
                    ),
                    "has_help": any(
                        e.name.endswith(".help") for e in entries
                    ),
                }
                return entries

        # --- Legacy mode: single tool with main.py ---
        handler_path = os.path.join(dir_path, "main.py")
        if not os.path.isfile(handler_path):
            handler_path = os.path.join(dir_path, "handler.py")
        if not os.path.isfile(handler_path):
            return []

        return [ToolEntry(
            name=meta.get("name", dir_name),
            title=meta.get("title", dir_name),
            description=meta.get("description", ""),
            input_schema=meta.get("input_schema", {}),
            version=group_version,
            tags=group_tags,
            dir_path=dir_path,
            handler_path=handler_path,
            canvas=meta.get("canvas"),
            storage=meta.get("storage"),
            context_text=group_context,
            keywords=meta.get("keywords", []),
            use_cases=meta.get("use_cases", []),
        )]

    def _scan_file(
        self, file_path: str, filename: str
    ) -> Optional[ToolEntry]:
        """Scan a single-file tool (no tool.json)."""
        if filename == "base_tool.py":
            return None

        cls = self._load_handler(file_path)
        if cls is None:
            return None

        name = getattr(cls, "name", "") or filename.removesuffix(".py")
        return ToolEntry(
            name=name,
            title=name,
            description=getattr(cls, "description", ""),
            input_schema=getattr(cls, "input_schema", {}),
            handler_path=file_path,
        )

    def _get_instance(
        self, entry: ToolEntry, context: dict[str, Any]
    ) -> Any:
        """Get or create a tool instance."""
        cls = self._load_handler(entry.handler_path)
        if cls is None:
            raise RuntimeError(
                f"Failed to load handler for tool '{entry.name}'"
            )
        # BaseTool subclasses accept context in __init__
        if issubclass(cls, BaseTool):
            return cls(context)
        # Standalone tool classes: try (context), then ()
        try:
            return cls(context)
        except TypeError:
            return cls()

    def _load_handler(
        self, handler_path: str
    ) -> Optional[type]:
        """Load a tool handler class from a Python file."""
        if not os.path.isfile(handler_path):
            return None

        module_name = (
            f"oedix_tool_{os.path.basename(handler_path).removesuffix('.py')}"
        )
        try:
            if module_name in sys.modules:
                del sys.modules[module_name]

            spec = importlib.util.spec_from_file_location(
                module_name, handler_path
            )
            if spec is None or spec.loader is None:
                return None

            module = importlib.util.module_from_spec(spec)
            sys.modules[module_name] = module

            handler_dir = os.path.dirname(handler_path)
            pkg_root = os.path.dirname(
                os.path.abspath(__file__)
            )  # oedix/tools/
            added_paths = []
            for p in (handler_dir, pkg_root):
                if p and p not in sys.path:
                    sys.path.insert(0, p)
                    added_paths.append(p)
            try:
                spec.loader.exec_module(module)
            finally:
                for p in added_paths:
                    try:
                        sys.path.remove(p)
                    except ValueError:
                        pass

            # 1) Prefer BaseTool subclass with a name attribute
            for attr_name in dir(module):
                attr = getattr(module, attr_name)
                if (
                    isinstance(attr, type)
                    and issubclass(attr, BaseTool)
                    and attr is not BaseTool
                    and getattr(attr, "name", "")
                ):
                    return attr

            # 2) Fallback: any class with a run() method (duck typing)
            for attr_name in dir(module):
                attr = getattr(module, attr_name)
                if (
                    isinstance(attr, type)
                    and attr.__module__ == module_name
                    and hasattr(attr, "run")
                    and callable(getattr(attr, "run", None))
                ):
                    return attr

            return None

        except Exception:
            logger.exception(
                "Failed to load tool handler: %s", handler_path
            )
            return None
