"""Tool loading -- dynamic tools, workspace tools, tool handlers."""

from __future__ import annotations

import importlib.util
import json
import logging
import os
from collections.abc import Awaitable, Callable
from pathlib import Path
from typing import TYPE_CHECKING, Any

from core.pathways.models.tool import ToolOutput
from core.tools import get_tool_registry
from infrastructure.tools.script_tool_loader import load_tools_from_directory

if TYPE_CHECKING:
    from core.pathways.context import Context

logger = logging.getLogger(__name__)

ToolHandler = Callable[["dict[str, Any]", "Context"], Awaitable[ToolOutput | dict[str, Any]]]


async def load_registered_tools(persist_dir: str) -> int:
    """Load dynamically registered tools from previous sessions."""
    workspace_root = os.getenv("ALBUS_WORKSPACE", os.getcwd())
    potential_paths = [
        Path(persist_dir) / "tools" / "manifest.json",
        Path(workspace_root) / ".albus" / "tools" / "manifest.json",
        Path(".albus") / "tools" / "manifest.json",
    ]

    manifest_path = None
    for p in potential_paths:
        if p.exists():
            manifest_path = p
            break

    if not manifest_path:
        logger.debug("No tool manifest found, skipping tool restoration")
        return 0

    try:
        with open(manifest_path) as f:
            manifest = json.load(f)

        tools_data = manifest.get("tools", manifest)
        loaded = 0

        for tool_name, tool_meta in tools_data.items():
            try:
                script_path = tool_meta.get("path")
                if not script_path:
                    continue

                script_path_obj = Path(script_path)
                full_path = None
                for base in [
                    Path(workspace_root),
                    Path.cwd(),
                    manifest_path.parent.parent.parent,
                ]:
                    candidate = base / script_path_obj
                    if candidate.exists():
                        full_path = candidate
                        break

                if not full_path or not full_path.exists():
                    logger.warning("Tool script not found: %s", script_path)
                    continue

                handler = await create_tool_handler(full_path, tool_meta)
                if handler:
                    registry = get_tool_registry()
                    registry.register(
                        name=tool_name,
                        handler=handler,
                        description=tool_meta.get(
                            "description", f"Custom tool: {tool_name}"
                        ),
                        parameters=tool_meta.get("parameters", {}),
                        requires_privileged=False,
                    )

                    loaded += 1
                    logger.debug("Restored tool: %s", tool_name)

            except Exception as e:
                logger.warning("Failed to restore tool %s: %s", tool_name, e)

        if loaded > 0:
            logger.info("Restored %d dynamically registered tools", loaded)

        return loaded

    except Exception as e:
        logger.warning("Failed to load tool manifest: %s", e)
        return 0


async def create_tool_handler(
    script_path: Path, meta: dict[str, Any]
) -> ToolHandler | None:
    """Create a tool handler from a script file."""
    try:
        spec = importlib.util.spec_from_file_location("dynamic_tool", script_path)
        if not spec or not spec.loader:
            return None

        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

        if not hasattr(module, "run"):
            logger.warning("Tool script missing run(): %s", script_path)
            return None

        run_func = module.run
        is_async = meta.get("is_async", True)

        async def handler(
            arguments: dict[str, Any], context: Context
        ) -> ToolOutput:
            try:
                result = await run_func(**arguments, context=context) if is_async else run_func(**arguments, context=context)
                if isinstance(result, ToolOutput):
                    return result
                # Wrap raw results for backwards compatibility
                return ToolOutput(success=True, data=result)
            except Exception as e:
                logger.debug("Dynamic tool handler failed: %s", e, exc_info=True)
                return ToolOutput(success=False, error=str(e))

        return handler

    except Exception as e:
        logger.warning("Failed to create handler for %s: %s", script_path, e)
        return None


async def load_workspace_tools(tools_dir: str | Path | None = None) -> int:
    """Auto-discover and register tools from a workspace tools/ directory.

    Each ``.py`` file with a ``run()`` function becomes a tool.
    Tool name is derived from the file path: ``custom.{filename}``.
    """
    if tools_dir is None:
        tools_dir = Path.cwd() / "tools"
    else:
        tools_dir = Path(os.path.expanduser(os.path.expandvars(str(tools_dir))))

    if not tools_dir.exists() or not tools_dir.is_dir():
        return 0

    try:
        results = load_tools_from_directory(tools_dir)
        registry = get_tool_registry()
        loaded = 0

        for tool_id, schema, handler in results:
            if "." not in tool_id:
                tool_id = f"custom.{tool_id}"

            registry.register(
                name=tool_id,
                handler=handler,
                description=getattr(schema, "description", f"Custom tool: {tool_id}"),
                parameters=getattr(schema, "input_schema", {}),
                requires_privileged=False,
            )
            loaded += 1
            logger.debug("Loaded workspace tool: %s", tool_id)

        if loaded:
            logger.info("Loaded %d workspace tools from %s", loaded, tools_dir)
        return loaded

    except ImportError:
        logger.debug("script_tool_loader not available, skipping workspace tool discovery")
        return 0
    except Exception as e:
        logger.warning("Failed to load workspace tools from %s: %s", tools_dir, e)
        return 0
