"""Tool utilities for agents.

Provides reusable helpers for tool argument parsing, skill loading,
agent delegation, input filtering, tool-call batch execution,
and tool handler auto-discovery.
"""

from __future__ import annotations

import asyncio
import json
from collections.abc import Awaitable, Callable
from pathlib import Path
from typing import TYPE_CHECKING, Any

from fluxibly.agent.skills import SkillMetadata, read_skill_body
from fluxibly.schema.tools import ToolCall, ToolResult

if TYPE_CHECKING:
    from fluxibly.agent.base import AgentConfig, BaseAgent
    from fluxibly.tools.executor import ToolService


# ══════════════════════════════════════════════════════════════
# Argument Parsing
# ══════════════════════════════════════════════════════════════


def parse_tool_args(arguments: dict[str, Any] | str | None) -> dict[str, Any]:
    """Parse tool call arguments to a dict."""
    if arguments is None:
        return {}
    if isinstance(arguments, dict):
        return arguments
    try:
        return json.loads(arguments)
    except (json.JSONDecodeError, TypeError):
        return {}


# ══════════════════════════════════════════════════════════════
# Skill Handling
# ══════════════════════════════════════════════════════════════


def handle_load_skill(
    tool_call: ToolCall,
    skill_map: dict[str, SkillMetadata],
) -> ToolResult:
    """Handle the ``load_skill`` tool call.

    Reads from ``sandbox_path`` when available (skill was staged into sandbox),
    falling back to the original host ``path``.  The returned body uses the
    sandbox skill directory so the LLM can reference scripts via shell.
    """
    args = parse_tool_args(tool_call.arguments)
    skill_name = args.get("name", "")

    skill = skill_map.get(skill_name)
    if not skill:
        return ToolResult(
            tool_call_id=tool_call.id,
            content=f"Skill '{skill_name}' not found. Available: {list(skill_map.keys())}",
            is_error=True,
        )

    # Prefer sandbox path (staged copy) over host path
    read_path = skill.sandbox_path or skill.path

    try:
        body = read_skill_body(read_path)
    except FileNotFoundError as e:
        return ToolResult(
            tool_call_id=tool_call.id,
            content=f"Skill '{skill_name}' SKILL.md not found: {e}",
            is_error=True,
        )

    return ToolResult(
        tool_call_id=tool_call.id,
        content=f"[Skill loaded: {skill.name} from {read_path}]\n\n{body}",
    )


def inject_skill_catalog(
    messages: list[dict[str, Any]],
    catalog: str,
) -> None:
    """Append skill catalog to the system message."""
    for msg in messages:
        if msg.get("role") == "system":
            msg["content"] = msg.get("content", "") + "\n" + catalog
            return


def evict_skill_bodies(
    messages: list[dict[str, Any]],
    skill_body_indices: list[tuple[int, str]],
) -> None:
    """Replace skill bodies in messages with short eviction notes."""
    for msg_idx, skill_name in skill_body_indices:
        if 0 <= msg_idx < len(messages):
            msg = messages[msg_idx]
            if msg.get("role") == "tool":
                msg["content"] = (
                    f"[Skill {skill_name} was loaded — "
                    f"call load_skill again if needed]"
                )


# ══════════════════════════════════════════════════════════════
# Agent Delegation
# ══════════════════════════════════════════════════════════════


async def delegate_to_agent(
    agent: BaseAgent,
    tool_call: ToolCall,
    context: dict[str, Any],
    logger: Any,
) -> ToolResult:
    """Delegate a tool call to a sub-agent (agent-as-tool)."""
    args = parse_tool_args(tool_call.arguments)
    input_text = args.pop("input", "")

    sub_messages: list[dict[str, Any]] = [
        {"role": "user", "content": input_text}
    ]

    try:
        sub_context = {**context, **args}
        agent_result = await agent.forward(sub_messages, sub_context)
        return ToolResult(
            tool_call_id=tool_call.id,
            content=agent_result.content.output.output_text,
        )
    except Exception as e:
        logger.error(
            "Sub-agent '{name}' execution failed: {err}",
            name=agent.config.name,
            err=e,
        )
        return ToolResult(
            tool_call_id=tool_call.id,
            content=f"Sub-agent '{agent.config.name}' error: {e}",
            is_error=True,
        )


# ══════════════════════════════════════════════════════════════
# Input Filtering
# ══════════════════════════════════════════════════════════════


def apply_input_filter(
    messages: list[dict[str, Any]],
    input_filter: str | None,
) -> list[dict[str, Any]]:
    """Apply an input filter to messages for a handoff."""
    if input_filter is None:
        return messages.copy()

    if input_filter == "last_message_only":
        for msg in reversed(messages):
            if msg.get("role") == "user":
                return [msg]
        return messages[-1:] if messages else []

    if input_filter == "remove_tools":
        return [
            msg
            for msg in messages
            if msg.get("role") not in ("tool",) and "tool_calls" not in msg
        ]

    # Custom filter: "path.to.module::func_name"
    if "::" in input_filter:
        import importlib

        module_path, func_name = input_filter.rsplit("::", 1)
        module = importlib.import_module(module_path)
        filter_fn = getattr(module, func_name)
        return filter_fn(messages)

    return messages.copy()


# ══════════════════════════════════════════════════════════════
# Tool Execution
# ══════════════════════════════════════════════════════════════


async def execute_function_tool(
    tool_call: ToolCall,
    tool_service: ToolService | None,
) -> ToolResult:
    """Execute a function/MCP tool call via ToolService."""
    if tool_service:
        return await tool_service.execute_tool_call(tool_call)

    return ToolResult(
        tool_call_id=tool_call.id,
        content=f"Tool '{tool_call.name}' not available — no ToolService configured",
        is_error=True,
    )


async def execute_tool_calls_batch(
    tool_calls: list[ToolCall],
    execute_fn: Callable[[ToolCall], Awaitable[ToolResult]],
    max_workers: int = 1,
) -> list[ToolResult]:
    """Execute tool calls with concurrency control."""
    if max_workers <= 1 or len(tool_calls) == 1:
        results: list[ToolResult] = []
        for tc in tool_calls:
            results.append(await execute_fn(tc))
        return results

    semaphore = asyncio.Semaphore(max_workers)

    async def _run(tc: ToolCall) -> ToolResult:
        async with semaphore:
            return await execute_fn(tc)

    return list(await asyncio.gather(*[_run(tc) for tc in tool_calls]))


# ══════════════════════════════════════════════════════════════
# Tool Handler Auto-Discovery
# ══════════════════════════════════════════════════════════════


async def auto_build_tool_service(
    tool_yaml_paths: dict[str, Path],
    tool_map: dict[str, dict[str, Any]],
    config: AgentConfig,
    sandbox_session: Any | None,
    sandbox_dir: str | None,
    logger: Any,
) -> tuple[ToolService | None, Any | None]:
    """Auto-discover tool handler ``.py`` files and build a ToolService.

    For each tool YAML that has a sibling ``.py`` file, loads the handler
    and registers it.  Returns ``(None, session)`` if no handlers found.

    Args:
        tool_yaml_paths: Map of tool name → resolved YAML path.
        tool_map: Map of tool name → tool schema dict.
        config: Agent config (used for ``config.sandbox``).
        sandbox_session: Existing SandboxSession (or ``None``).
        sandbox_dir: Sandbox directory path (or ``None``).
        logger: Logger instance.

    Returns:
        Tuple of ``(tool_service_or_none, sandbox_session_or_none)``.
        The sandbox session is returned because it may have been created.
    """
    from fluxibly.tools.executor import ToolService as _ToolService
    from fluxibly.tools.handler_loader import (
        find_handler_path,
        resolve_handler,
    )

    # Discover handler .py files
    handlers_to_load: dict[str, Path] = {}
    for tool_name, yaml_path in tool_yaml_paths.items():
        py_path = find_handler_path(yaml_path)
        if py_path is not None:
            handlers_to_load[tool_name] = py_path

    if not handlers_to_load:
        return None, sandbox_session

    # Build handler context (session, etc.)
    handler_context, sandbox_session = build_handler_context(
        config,
        sandbox_session,
        sandbox_dir,
    )

    # Create ToolService and register handlers
    ts = _ToolService()
    for tool_name, py_path in handlers_to_load.items():
        handler = await resolve_handler(py_path, handler_context)
        schema = tool_map[tool_name]
        ts.register_function_from_schema(schema, handler)

    logger.info(
        "Auto-built ToolService with {n} handler(s): {names}",
        n=len(handlers_to_load),
        names=list(handlers_to_load.keys()),
    )
    return ts, sandbox_session


def build_handler_context(
    config: AgentConfig,
    sandbox_session: Any | None,
    sandbox_dir: str | None,
) -> tuple[dict[str, Any], Any | None]:
    """Build the context dict passed to ``create_handler()`` factories.

    Creates a ``SandboxSession`` if ``config.sandbox`` is set and no
    session exists yet.

    Args:
        config: Agent config.
        sandbox_session: Existing SandboxSession (or ``None``).
        sandbox_dir: Sandbox directory path (or ``None``).

    Returns:
        Tuple of ``(context_dict, sandbox_session_or_none)``.
        The sandbox session is returned because it may have been created.
    """
    context: dict[str, Any] = {}

    if config.sandbox and sandbox_session is None:
        from fluxibly.tools.sandbox import SandboxConfig, SandboxSession

        sandbox_config = SandboxConfig(**config.sandbox)
        resolved_dir = sandbox_dir or f"/tmp/fluxibly/{config.name}"
        sandbox_session = SandboxSession(resolved_dir, sandbox_config)

    if sandbox_session is not None:
        context["session"] = sandbox_session

    return context, sandbox_session


# ══════════════════════════════════════════════════════════════
# Tool List Assembly
# ══════════════════════════════════════════════════════════════


def rebuild_all_tools(
    tool_map: dict[str, dict[str, Any]],
    handoff_map: dict[str, dict[str, Any]],
    agent_tool_map: dict[str, dict[str, Any]],
    skill_map: dict[str, SkillMetadata],
) -> tuple[list[dict[str, Any]], list[str], list[dict[str, Any]]]:
    """Rebuild the flat tool list from all resolution maps.

    Args:
        tool_map: Map of tool name → tool schema dict.
        handoff_map: Map of handoff tool name → handoff entry dict.
        agent_tool_map: Map of agent tool name → agent entry dict.
        skill_map: Map of skill name → SkillMetadata.

    Returns:
        Tuple of ``(all_tools, tool_names, tool_details)`` — the flat
        list of tool dicts for the LLM, the list of tool names, and a
        list of ``{name, description, parameters}`` dicts for monitoring.
    """
    from fluxibly.agent.skills import build_load_skill_tool_schema

    tools: list[dict[str, Any]] = []

    # Tool dicts
    tools.extend(tool_map.values())

    # Handoff tool dicts
    for entry in handoff_map.values():
        tools.append(entry["tool_dict"])

    # Agent-as-tool dicts
    for entry in agent_tool_map.values():
        tools.append(entry["tool_dict"])

    # load_skill tool (if skills are configured)
    if skill_map:
        tools.append(build_load_skill_tool_schema())

    tool_names: list[str] = [
        (t.get("function", {}).get("name") or t.get("name", ""))
        for t in tools
        if (t.get("function", {}).get("name") or t.get("name"))
    ]

    # Build tool details for monitoring (name + description + parameters)
    tool_details: list[dict[str, Any]] = []
    for t in tools:
        fn = t.get("function", {})
        name = fn.get("name") or t.get("name", "")
        if name:
            tool_details.append(
                {
                    "name": name,
                    "description": fn.get("description", ""),
                    "parameters": fn.get("parameters", {}),
                }
            )

    return tools, tool_names, tool_details


# ══════════════════════════════════════════════════════════════
# Logging
# ══════════════════════════════════════════════════════════════


def log_tool_loop_iteration(
    logger: Any,
    iteration: int,
    max_iterations: int,
    tool_calls: list[ToolCall],
) -> None:
    """Log tool loop iteration summary and details."""
    logger.info(
        "Tool loop iteration {i}/{max}: {n} tool call(s) — {names}",
        i=iteration,
        max=max_iterations,
        n=len(tool_calls),
        names=[tc.name for tc in tool_calls],
    )
    logger.debug(
        "Tool call details:\n{details}",
        details=[
            {"name": tc.name, "id": tc.id, "arguments": tc.arguments}
            for tc in tool_calls
        ],
    )
