"""Per-thread cache for Arcade SelectTools schemas.

Caches the full JSON tool schemas returned by Arcade_SelectTools so the agent
can look up parameter details without re-fetching or bloating context.
Storage: ~/.cadecoder/tool_schemas/{thread_id}.json
"""

from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

log = logging.getLogger("cadecoder")


@dataclass
class CachedToolSchema:
    """One cached tool schema entry."""

    name: str
    description: str
    parameters: dict[str, Any]
    query_id: str
    cached_at: str  # ISO 8601


class ToolSchemaCache:
    """Per-thread cache for Arcade tool schemas.

    Storage: ``~/.cadecoder/tool_schemas/{thread_id}.json``

    Each file is a JSON object accumulating tools across multiple
    ``Arcade_SelectTools`` calls within the same thread.
    """

    def __init__(self, storage_dir: Path | None = None) -> None:
        if storage_dir is None:
            storage_dir = Path.home() / ".cadecoder" / "tool_schemas"
        self.storage_dir = storage_dir
        try:
            self.storage_dir.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            log.warning(f"Cannot create tool schema cache dir: {e}")
        self._disk_ok = self.storage_dir.exists()

    def _thread_path(self, thread_id: str) -> Path:
        """Path to the thread's schema cache file, stored per-thread."""
        safe_id = thread_id.replace("/", "_").replace("..", "_")
        thread_dir = self.storage_dir / safe_id
        thread_dir.mkdir(parents=True, exist_ok=True)
        return thread_dir / "tool_schemas.json"

    # ------------------------------------------------------------------
    # Caching
    # ------------------------------------------------------------------

    def cache_select_tools_result(self, thread_id: str, content: str) -> int:
        """Parse a SelectTools JSON result and cache all tool schemas.

        Returns the number of tools cached.
        """
        try:
            data = json.loads(content)
        except (json.JSONDecodeError, TypeError):
            log.warning("Cannot cache SelectTools result: invalid JSON")
            return 0

        tools, query_id = self._parse_select_tools_data(data)
        if not tools:
            return 0

        # Load existing cache and merge
        cache_data = self._load_thread_cache(thread_id)
        now = datetime.now(UTC).isoformat()
        cache_data["updated_at"] = now

        cached_count = 0
        for tool_obj in tools:
            name = _extract_tool_name(tool_obj)
            if name == "unknown":
                continue
            description = _extract_tool_description(tool_obj)
            parameters = _extract_tool_parameters(tool_obj)

            cache_data["tools"][name] = asdict(
                CachedToolSchema(
                    name=name,
                    description=description,
                    parameters=parameters,
                    query_id=query_id,
                    cached_at=now,
                )
            )
            cached_count += 1

        self._save_thread_cache(thread_id, cache_data)
        return cached_count

    # ------------------------------------------------------------------
    # Retrieval
    # ------------------------------------------------------------------

    def get_tool_schema(self, thread_id: str, tool_name: str) -> CachedToolSchema | None:
        """Get the full schema for a specific tool from cache."""
        cache_data = self._load_thread_cache(thread_id)
        tool_data = cache_data["tools"].get(tool_name)
        if tool_data is None:
            # Case-insensitive fallback
            name_lower = tool_name.lower()
            for key, val in cache_data["tools"].items():
                if key.lower() == name_lower:
                    tool_data = val
                    break
        if tool_data is None:
            return None
        return CachedToolSchema(**tool_data)

    def list_tools(self, thread_id: str) -> list[dict[str, str]]:
        """List all cached tools for a thread (name + description + query_id)."""
        cache_data = self._load_thread_cache(thread_id)
        result = []
        for tool_data in cache_data["tools"].values():
            result.append(
                {
                    "name": tool_data["name"],
                    "description": tool_data.get("description", ""),
                    "query_id": tool_data.get("query_id", ""),
                }
            )
        return result

    # ------------------------------------------------------------------
    # Cache clearing
    # ------------------------------------------------------------------

    def clear_all(self) -> int:
        """Remove all cached tool schemas across all threads.

        Called when an MCP server is disabled to prevent stale schema references.
        Returns the number of cache files removed.
        """
        if not self._disk_ok:
            return 0
        removed = 0
        for path in self.storage_dir.glob("*/tool_schemas.json"):
            try:
                path.unlink()
                removed += 1
            except OSError as e:
                log.warning(f"Failed to remove schema cache {path.name}: {e}")
            # Also remove stale lock files
            lock_path = path.with_suffix(".json.lock")
            if lock_path.exists():
                try:
                    lock_path.unlink()
                except OSError:
                    pass
        if removed:
            log.info(f"Cleared {removed} tool schema cache file(s)")
        return removed

    # ------------------------------------------------------------------
    # Persistence helpers
    # ------------------------------------------------------------------

    def _load_thread_cache(self, thread_id: str) -> dict[str, Any]:
        """Load a thread's cache file. Returns empty structure if missing."""
        path = self._thread_path(thread_id)
        if not self._disk_ok or not path.exists():
            return {"updated_at": "", "tools": {}}
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError) as e:
            log.warning(f"Corrupt schema cache for thread {thread_id}: {e}")
            return {"updated_at": "", "tools": {}}

    def _save_thread_cache(self, thread_id: str, data: dict[str, Any]) -> None:
        """Write the thread's cache file."""
        if not self._disk_ok:
            return
        path = self._thread_path(thread_id)
        try:
            from filelock import FileLock

            lock = FileLock(str(path) + ".lock", timeout=10)
            with lock:
                path.write_text(
                    json.dumps(data, indent=2, ensure_ascii=False),
                    encoding="utf-8",
                )
        except Exception as e:
            log.warning(f"Failed to save schema cache: {e}")

    # ------------------------------------------------------------------
    # SelectTools JSON parsing
    # ------------------------------------------------------------------

    @staticmethod
    def _parse_select_tools_data(
        data: Any,
    ) -> tuple[list[dict[str, Any]], str]:
        """Extract tools list and query_id from various SelectTools formats.

        Returns (tools_list, query_id).
        """
        tools: list[dict[str, Any]] = []
        query_id = ""

        if isinstance(data, dict):
            query_id = data.get("query_id") or data.get("queryId") or ""
            tool_list = data.get("tools") or data.get("results") or data.get("definitions")
            if isinstance(tool_list, list):
                tools = tool_list
            elif "function" in data or "name" in data:
                tools = [data]
        elif isinstance(data, list):
            tools = data

        return tools, str(query_id)


# ---------------------------------------------------------------------------
# Field extraction helpers (shared with tool_result_store.py patterns)
# ---------------------------------------------------------------------------


def _strip_version_suffix(name: str) -> str:
    """Remove Arcade toolkit version suffix (e.g. ``@3_1_2``) from tool names."""
    at_idx = name.find("@")
    return name[:at_idx] if at_idx > 0 else name


def _extract_tool_name(tool_obj: dict[str, Any]) -> str:
    """Extract tool name from various schema formats."""
    func = tool_obj.get("function", {})
    if isinstance(func, dict) and func.get("name"):
        return _strip_version_suffix(func["name"])
    return _strip_version_suffix(tool_obj.get("name", "unknown"))


def _extract_tool_description(tool_obj: dict[str, Any]) -> str:
    """Extract tool description from various schema formats."""
    func = tool_obj.get("function", {})
    if isinstance(func, dict) and func.get("description"):
        return func["description"]
    return tool_obj.get("description", "")


def _extract_tool_parameters(tool_obj: dict[str, Any]) -> dict[str, Any]:
    """Extract tool parameters schema from various formats."""
    func = tool_obj.get("function", {})
    if isinstance(func, dict) and func.get("parameters"):
        return func["parameters"]
    return tool_obj.get("parameters", {})
