"""
Phase 80: Dynamic Tool Loading — Runtime tool registration and execution.

Agents can register, discover, and execute tools at runtime without restarting
the MCP server. Supports importing tool modules, schema export, and usage stats.

Storage:
  state/broker/dynamic_tools/
    registry.json   — tool definitions
    stats.jsonl     — per-call execution log
"""

import hashlib
import importlib
import inspect
import json
import logging
import time
import traceback
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

logger = logging.getLogger(__name__)

# ── Paths ──────────────────────────────────────────────────────────────────

TOOLS_DIR = BROKER_DIR / "dynamic_tools"
REGISTRY_FILE = TOOLS_DIR / "registry.json"
STATS_FILE = TOOLS_DIR / "stats.jsonl"

for _d in (TOOLS_DIR,):
    _d.mkdir(parents=True, exist_ok=True)

# ── In-memory state ────────────────────────────────────────────────────────

_handlers: Dict[str, Callable] = {}
_call_counts: Dict[str, int] = {}
_error_counts: Dict[str, int] = {}


def reset() -> None:
    """Reset in-memory state. Used by tests."""
    global _handlers, _call_counts, _error_counts
    _handlers = {}
    _call_counts = {}
    _error_counts = {}


# ── Internal helpers ───────────────────────────────────────────────────────

def _load_registry() -> dict:
    if REGISTRY_FILE.exists():
        try:
            return json.loads(REGISTRY_FILE.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            return {"tools": {}}
    return {"tools": {}}


def _save_registry(reg: dict) -> None:
    _write_json_sync(REGISTRY_FILE, reg)


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


# ── Public API ─────────────────────────────────────────────────────────────

def register_tool(
    name: str,
    description: str,
    handler_fn: Callable,
    parameters: Optional[dict] = None,
    category: str = "custom",
) -> dict:
    """Register a runtime tool with a callable handler."""
    if not name or not callable(handler_fn):
        raise ValueError("name and callable handler_fn are required")

    reg = _load_registry()
    tool_def = {
        "name": name,
        "description": description,
        "category": category,
        "parameters": parameters or {},
        "registered_at": _now(),
        "active": True,
    }
    reg["tools"][name] = tool_def
    _save_registry(reg)
    _handlers[name] = handler_fn
    _call_counts.setdefault(name, 0)
    _error_counts.setdefault(name, 0)
    logger.info("Registered tool: %s (category=%s)", name, category)
    return tool_def


def unregister_tool(name: str) -> bool:
    """Remove a registered tool."""
    reg = _load_registry()
    if name not in reg["tools"]:
        return False
    del reg["tools"][name]
    _save_registry(reg)
    _handlers.pop(name, None)
    return True


def get_tool(name: str) -> Optional[dict]:
    """Get a tool definition by name."""
    reg = _load_registry()
    return reg["tools"].get(name)


def list_tools(category: Optional[str] = None) -> List[dict]:
    """List all registered tools, optionally filtered by category."""
    reg = _load_registry()
    tools = list(reg["tools"].values())
    if category:
        tools = [t for t in tools if t.get("category") == category]
    return tools


def execute_tool(name: str, arguments: Optional[dict] = None) -> dict:
    """Execute a registered tool with error isolation."""
    if name not in _handlers:
        return {"ok": False, "error": f"Tool '{name}' not found or not loaded"}

    start = time.time()
    _call_counts[name] = _call_counts.get(name, 0) + 1
    try:
        result = _handlers[name](arguments or {})
        elapsed = time.time() - start
        log_entry = {
            "tool": name,
            "ts": _now(),
            "elapsed_ms": round(elapsed * 1000, 2),
            "ok": True,
        }
        _append_jsonl_sync(STATS_FILE, log_entry)
        return {"ok": True, "result": result, "elapsed_ms": log_entry["elapsed_ms"]}
    except Exception as exc:
        elapsed = time.time() - start
        _error_counts[name] = _error_counts.get(name, 0) + 1
        log_entry = {
            "tool": name,
            "ts": _now(),
            "elapsed_ms": round(elapsed * 1000, 2),
            "ok": False,
            "error": str(exc),
        }
        _append_jsonl_sync(STATS_FILE, log_entry)
        return {"ok": False, "error": str(exc), "elapsed_ms": log_entry["elapsed_ms"]}


# Decorator / marker attribute for auto-import
TOOL_MARKER = "_brynq_tool"


def tool(name: Optional[str] = None, description: str = "", parameters: Optional[dict] = None, category: str = "custom"):
    """Decorator to mark a function as a dynamic tool."""
    def decorator(fn):
        fn._brynq_tool = {
            "name": name or fn.__name__,
            "description": description or (fn.__doc__ or "").strip(),
            "parameters": parameters,
            "category": category,
        }
        return fn
    return decorator


def import_tools_from_module(module_path: str) -> List[str]:
    """Import a module and auto-register functions decorated with @tool."""
    mod = importlib.import_module(module_path)
    registered = []
    for attr_name in dir(mod):
        obj = getattr(mod, attr_name)
        if callable(obj) and hasattr(obj, TOOL_MARKER):
            meta = getattr(obj, TOOL_MARKER)
            register_tool(
                name=meta["name"],
                description=meta["description"],
                handler_fn=obj,
                parameters=meta.get("parameters"),
                category=meta.get("category", "custom"),
            )
            registered.append(meta["name"])
    return registered


def export_tool_schema(name: Optional[str] = None) -> dict:
    """Export OpenAPI-style JSON schema for one or all tools."""
    reg = _load_registry()
    if name:
        t = reg["tools"].get(name)
        if not t:
            return {"error": f"Tool '{name}' not found"}
        return _tool_to_schema(t)
    schemas = {}
    for tname, tdef in reg["tools"].items():
        schemas[tname] = _tool_to_schema(tdef)
    return schemas


def _tool_to_schema(t: dict) -> dict:
    return {
        "name": t["name"],
        "description": t.get("description", ""),
        "parameters": {
            "type": "object",
            "properties": t.get("parameters", {}),
        },
        "category": t.get("category", "custom"),
    }


def get_tool_stats() -> dict:
    """Return call counts and error counts per tool."""
    stats = {}
    for name in set(list(_call_counts.keys()) + list(_error_counts.keys())):
        stats[name] = {
            "calls": _call_counts.get(name, 0),
            "errors": _error_counts.get(name, 0),
        }
    return stats
