"""Script Tool Loader - Load tools from Python scripts.

Tools are simple Python scripts with a `run()` function.
Schema is auto-extracted from type hints and docstrings.

Tool file format:
    '''Tool description here.'''

    async def run(query: str, max_results: int = 5, context=None) -> ToolOutput:
        '''
        Args:
            query: Search query string
            max_results: Maximum results to return
        '''
        # implementation
        return ToolOutput(success=True, data={"results": [...]})

Tool discovery:
    tools/
      web/
        search.py      → tool_id: "web.search"
        fetch.py       → tool_id: "web.fetch"
      workspace/
        read_file.py   → tool_id: "workspace.read_file"
"""

import importlib.util
import inspect
import logging
import sys
import types
from collections.abc import Awaitable, Callable
from pathlib import Path
from typing import Any, get_type_hints

from core.pathways.context import Context
from core.pathways.models.tool import ToolArgs, ToolOutput
from infrastructure.tools import ToolSchema

logger = logging.getLogger(__name__)

ToolHandler = Callable[[ToolArgs, Context], Awaitable[ToolOutput]]


def load_tools_from_directory(
    tools_dir: Path | str,
) -> list[tuple[str, ToolSchema, ToolHandler]]:
    """Load all tools from a directory.

    Internal helper modules (``_*.py``) are loaded with scoped
    ``sys.modules`` entries so that identically-named helpers in
    different tool packages (e.g. ``memory/tools/_store.py`` vs
    ``knowledge/tools/_store.py``) never collide.

    Args:
        tools_dir: Path to tools directory

    Returns
    -------
        List of (tool_id, schema, handler) tuples
    """
    tools_dir = Path(tools_dir)
    if not tools_dir.exists():
        logger.warning("Tools directory does not exist: %s", tools_dir)
        return []

    tools_parent = str(tools_dir)

    # ------------------------------------------------------------------
    # 1. Pre-load _*.py helper modules with directory-scoped names so
    #    they don't collide across tool packages.
    # ------------------------------------------------------------------
    scope = tools_dir.parent.name  # e.g. "memory", "knowledge"
    helpers: dict[str, Any] = {}  # bare_name -> module

    for helper_file in sorted(tools_dir.glob("_*.py")):
        bare_name = helper_file.stem  # "_store"
        scoped_name = f"_toolpkg_{scope}.{bare_name}"

        if tools_parent not in sys.path:
            sys.path.insert(0, tools_parent)

        try:
            spec = importlib.util.spec_from_file_location(scoped_name, helper_file)
            if spec and spec.loader:
                mod = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(mod)
                sys.modules[scoped_name] = mod
                helpers[bare_name] = mod
        except Exception as e:
            logger.warning("Failed to load helper %s: %s", helper_file, e)

        if tools_parent in sys.path:
            sys.path.remove(tools_parent)

    # ------------------------------------------------------------------
    # 2. Temporarily install helpers under their bare names while
    #    loading the tool modules, then restore the originals.
    # ------------------------------------------------------------------
    saved: dict[str, Any] = {}
    for bare_name, mod in helpers.items():
        saved[bare_name] = sys.modules.get(bare_name)
        sys.modules[bare_name] = mod

    try:
        tools: list[tuple[str, ToolSchema, ToolHandler]] = []
        for tool_file in sorted(tools_dir.rglob("*.py")):
            if tool_file.name.startswith("_"):
                continue

            try:
                result = load_tool_from_file(tool_file, tools_dir)
                if result:
                    tools.append(result)
            except Exception as e:
                logger.warning("Failed to load tool from %s: %s", tool_file, e)
    finally:
        # Restore bare-name entries so later packages get a clean slate
        for bare_name, original in saved.items():
            if original is None:
                sys.modules.pop(bare_name, None)
            else:
                sys.modules[bare_name] = original

    return tools


def load_tool_from_file(
    tool_file: Path,
    base_dir: Path | None = None,
) -> tuple[str, ToolSchema, ToolHandler] | None:
    """Load a single tool from a Python file.

    Args:
        tool_file: Path to the tool file
        base_dir: Base directory for computing tool_id

    Returns
    -------
        (tool_id, schema, handler) or None if no run() function
    """
    # Compute tool_id from path
    rel_path = tool_file.relative_to(base_dir) if base_dir else Path(tool_file.name)

    tool_id = str(rel_path.with_suffix("")).replace("/", ".").replace("\\", ".")

    # Load module
    # Add tool directory to sys.path so sibling imports work
    tools_parent = str(tool_file.parent)
    added_to_path = False
    if tools_parent not in sys.path:
        sys.path.insert(0, tools_parent)
        added_to_path = True

    try:
        spec = importlib.util.spec_from_file_location(f"tool_{tool_id}", tool_file)
        if not spec or not spec.loader:
            return None

        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
    finally:
        # Clean up sys.path
        if added_to_path and tools_parent in sys.path:
            sys.path.remove(tools_parent)

    # Get run function
    run_fn = getattr(module, "run", None)
    if not run_fn or not callable(run_fn):
        logger.debug("No run() function in %s", tool_file)
        return None

    # Extract schema
    schema = extract_schema(tool_id, run_fn, module)

    # Create handler
    handler = create_handler(run_fn)

    logger.debug("Loaded tool: %s from %s", tool_id, tool_file)
    return tool_id, schema, handler


def extract_schema(
    tool_id: str,
    run_fn: Callable,
    module: Any,
) -> ToolSchema:
    """Extract ToolSchema from function signature and docstrings.

    Args:
        tool_id: Tool identifier
        run_fn: The run() function
        module: The module containing the function

    Returns
    -------
        ToolSchema with extracted metadata
    """
    # Get description from module docstring or function docstring
    description = module.__doc__ or run_fn.__doc__ or f"Tool: {tool_id}"
    description = description.strip().split("\n")[0]  # First line only

    # Extract parameter info from signature
    sig = inspect.signature(run_fn)

    # Try to get type hints
    try:
        hints = get_type_hints(run_fn)
    except Exception:
        hints = {}

    # Parse docstring for parameter descriptions
    param_docs = _parse_docstring_params(run_fn.__doc__ or "")

    # Build JSON schema for parameters
    properties: dict[str, Any] = {}
    required: list[str] = []

    for name, param in sig.parameters.items():
        # Skip context parameter if present
        if name in ("context", "ctx"):
            continue

        # Get full JSON schema for the type (handles arrays with items, etc.)
        hint = hints.get(name)
        prop = _python_type_to_json_schema(hint)

        # Description from docstring
        if name in param_docs:
            prop["description"] = param_docs[name]

        # Default value
        if param.default is not inspect.Parameter.empty:
            prop["default"] = param.default
        else:
            required.append(name)

        properties[name] = prop

    input_schema = {
        "type": "object",
        "properties": properties,
        "required": required,
    }

    # Generic output contract for all tools.
    # Tool-specific payload lives under ToolOutput.data.
    output_schema = ToolOutput.model_json_schema()

    return ToolSchema(
        id=tool_id,
        name=tool_id.rsplit(".", maxsplit=1)[-1],
        description=description,
        input_schema=input_schema,
        output_schema=output_schema,
        category="builtin",
        requires_privileged=False,
    )


def create_handler(run_fn: Callable) -> ToolHandler:
    """Create a ToolRegistry-compatible handler from a run() function.

    Wraps run(typed_args) → handler(ToolArgs, context) → ToolOutput
    Pydantic end-to-end: ToolArgs in, ToolOutput out.
    """
    sig = inspect.signature(run_fn)
    is_async = inspect.iscoroutinefunction(run_fn)

    async def handler(args: ToolArgs, context: Context) -> ToolOutput:
        # Extract kwargs from ToolArgs Pydantic model
        # Use model_dump() to get extra fields, then map to function signature
        args_dict = args.model_dump()

        kwargs: dict[str, Any] = {}

        for name, param in sig.parameters.items():
            if name in ("context", "ctx"):
                kwargs[name] = context
            elif name in args_dict:
                kwargs[name] = args_dict[name]
            elif param.default is not inspect.Parameter.empty:
                kwargs[name] = param.default

        # Call the function
        try:
            if is_async:
                result = await run_fn(**kwargs)
            else:
                result = run_fn(**kwargs)

        except Exception as e:
            # Tool runtime failure: return canonical ToolOutput (do not raise)
            return ToolOutput(success=False, error=str(e))

        # STRICT: tools MUST return ToolOutput (no legacy dict/primitive shims).
        if isinstance(result, ToolOutput):
            return result

        return ToolOutput(
            success=False,
            error="Tool must return ToolOutput",
            metadata={"returned_type": type(result).__name__},
        )

    return handler


def _python_type_to_json_schema(hint: Any) -> dict[str, Any]:
    """Convert Python type hint to JSON schema dict.

    Returns a full schema dict (not just type string) to properly handle
    arrays with items, etc. OpenAI requires arrays to have items defined.
    """
    if hint is None:
        return {"type": "string"}

    # Handle Optional, Union, etc.
    origin = getattr(hint, "__origin__", None)
    args = getattr(hint, "__args__", ())

    if origin is not None:
        # List[T] -> array with items
        if origin is list:
            if args:
                # Has type argument like list[str]
                item_schema = _python_type_to_json_schema(args[0])
                return {"type": "array", "items": item_schema}
            # Bare list without type arg
            return {"type": "array", "items": {"type": "string"}}

        # Dict[K, V] -> object
        if origin is dict:
            return {"type": "object"}

        # Union[X, None] = Optional[X] or X | Y (Python 3.10+)
        if origin is types.UnionType or (
            hasattr(origin, "__name__") and origin.__name__ == "Union"
        ):
            # Filter out None and use the first non-None type
            non_none = [a for a in args if a is not type(None)]
            if non_none:
                return _python_type_to_json_schema(non_none[0])
            return {"type": "string"}

        # Handle Python 3.10+ union syntax (X | Y)
        # Check if this is a union type (not Optional)
        if hasattr(hint, "__args__") and len(args) > 1:
            # For unions like list[dict] | dict, prefer the first non-None type
            # or use object as fallback
            non_none_types = [a for a in args if a is not type(None)]
            if non_none_types:
                # If one is a generic (like list), prefer it, otherwise use first
                for t in non_none_types:
                    if hasattr(t, "__origin__") and t.__origin__ is list:
                        return _python_type_to_json_schema(t)
                return _python_type_to_json_schema(non_none_types[0])
            return {"type": "object"}  # Fallback for complex unions

    # Simple types
    if hint is str:
        return {"type": "string"}
    if hint is int:
        return {"type": "integer"}
    if hint is float:
        return {"type": "number"}
    if hint is bool:
        return {"type": "boolean"}
    if hint is list:
        # Bare list without type hints - default to array of strings
        return {"type": "array", "items": {"type": "string"}}
    if hint is dict:
        return {"type": "object"}

    return {"type": "string"}  # Default


def _parse_docstring_params(docstring: str) -> dict[str, str]:
    """Parse parameter descriptions from docstring.

    Supports Google-style docstrings:
        Args:
            query: Search query string
            max_results: Maximum results
    """
    params: dict[str, str] = {}

    if not docstring:
        return params

    lines = docstring.split("\n")
    in_args = False

    for line in lines:
        stripped = line.strip()

        if stripped.lower() in ("args:", "arguments:", "parameters:"):
            in_args = True
            continue

        if in_args:
            # Check if we've left the Args section
            if (
                stripped
                and not stripped.startswith(" ")
                and ":" not in stripped[:20]
                and stripped.lower() in ("returns:", "raises:", "yields:", "examples:")
            ):
                break

            # Parse parameter line: "name: description" or "name (type): description"
            if ":" in stripped:
                parts = stripped.split(":", 1)
                name = parts[0].strip()
                # Remove type annotation if present: "name (type)" → "name"
                if "(" in name:
                    name = name.split("(")[0].strip()
                desc = parts[1].strip() if len(parts) > 1 else ""
                if name and not name.startswith("-"):
                    params[name] = desc

    return params


__all__ = [
    "create_handler",
    "extract_schema",
    "load_tool_from_file",
    "load_tools_from_directory",
]
