"""
tools.py — AIPA Tool Registry
Provides a lightweight @tool() decorator, schema export, and tool execution
for agent tool-calling loops.

Agent-management tools allow Evelyn (and any tool-enabled agent) to read and
modify agent definitions without restarting the session.
"""

from __future__ import annotations

import inspect
import json
from pathlib import Path
from typing import Callable, Any

import config as _config

# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------

_REGISTRY: dict[str, dict] = {}   # name → {fn, schema}


def tool(description: str, parameters: dict):
    """
    Decorator that registers a callable as a named tool.

    Parameters
    ----------
    description : str
        Plain-language description shown to the LLM.
    parameters : dict
        JSON Schema object for the tool's parameters (OpenAI function-calling
        format — {"type": "object", "properties": {...}, "required": [...]}).
    """
    def decorator(fn: Callable) -> Callable:
        name = fn.__name__
        _REGISTRY[name] = {
            "fn": fn,
            "schema": {
                "type": "function",
                "function": {
                    "name": name,
                    "description": description,
                    "parameters": parameters,
                },
            },
        }
        return fn
    return decorator


def get_schemas(names: list[str]) -> list[dict]:
    """Return OpenAI-compatible tool schemas for the named tools."""
    result = []
    for name in names:
        entry = _REGISTRY.get(name)
        if entry is None:
            raise ValueError(
                f"Tool {name!r} is not registered. "
                f"Available tools: {list(_REGISTRY)}"
            )
        result.append(entry["schema"])
    return result


def call(name: str, arguments_json: str) -> str:
    """
    Execute a registered tool by name with a JSON-encoded arguments string.
    Returns a JSON-encoded result string (always a string for message history).
    """
    entry = _REGISTRY.get(name)
    if entry is None:
        return json.dumps({"error": f"Unknown tool: {name!r}"})
    try:
        kwargs = json.loads(arguments_json) if arguments_json else {}
        result = entry["fn"](**kwargs)
        return json.dumps(result, ensure_ascii=False)
    except Exception as exc:
        return json.dumps({"error": str(exc)})


# ---------------------------------------------------------------------------
# YAML helpers (ruamel for comment-preserving round-trips)
# ---------------------------------------------------------------------------

def _load_agents_yaml() -> tuple[Any, Any]:
    """
    Load agents.yaml using ruamel.yaml for comment preservation.
    Returns (ruamel_instance, data_dict).
    """
    from ruamel.yaml import YAML
    ryaml = YAML()
    ryaml.preserve_quotes = True
    path = _config._AGENTS_CONFIG_PATH
    with open(path, encoding="utf-8") as fh:
        data = ryaml.load(fh)
    return ryaml, data


def _save_agents_yaml(ryaml: Any, data: Any) -> None:
    """Write agents.yaml back with comment preservation."""
    path = _config._AGENTS_CONFIG_PATH
    with open(path, "w", encoding="utf-8") as fh:
        ryaml.dump(data, fh)


# ---------------------------------------------------------------------------
# Agent-management tools
# ---------------------------------------------------------------------------

@tool(
    description="List all agents currently defined in agents.yaml, including their title, provider, model, temperature, max_tokens, stateful flag, and tools list.",
    parameters={
        "type": "object",
        "properties": {},
        "required": [],
    },
)
def list_agents() -> list[dict]:
    """Return all agent configs as a list of dicts."""
    result = []
    for name, cfg in _config.AGENT_CONFIGS.items():
        result.append({
            "name": name,
            "title": _config.agent_title(name),
            "color": cfg.get("color", ""),
            "prompt_file": cfg.get("prompt_file", ""),
            "provider": _config.agent_provider_name(name),
            "provider_type": _config.agent_provider_type(name),
            "model": _config.agent_model(name),
            "temperature": cfg.get("temperature"),
            "max_tokens": cfg.get("max_tokens"),
            "stateful": cfg.get("stateful"),
            "tools": cfg.get("tools", []),
        })
    return result


@tool(
    description="Read the full configuration for a specific agent by name, including provider type, thinking mode, model, and tools.",
    parameters={
        "type": "object",
        "properties": {
            "agent_name": {
                "type": "string",
                "description": "The agent's short name (e.g. 'evelyn', 'atlas').",
            }
        },
        "required": ["agent_name"],
    },
)
def read_agent_config(agent_name: str) -> dict:
    """Return the full resolved config for one agent."""
    name = agent_name.lower()
    if name not in _config.AGENT_CONFIGS:
        return {"error": f"No agent named {agent_name!r}. Known: {list(_config.AGENT_CONFIGS)}"}
    cfg = _config.AGENT_CONFIGS[name]
    provider_name = _config.agent_provider_name(name)
    provider_def = _config._providers.get(provider_name, {})
    return {
        "name": name,
        "title": _config.agent_title(name),
        "color": cfg.get("color", ""),
        "prompt_file": cfg.get("prompt_file"),
        "provider": provider_name,
        "provider_type": _config.agent_provider_type(name),
        "provider_base_url": provider_def.get("base_url", ""),
        "provider_default_model": provider_def.get("default_model", ""),
        "model": _config.agent_model(name),
        "temperature": cfg.get("temperature"),
        "max_tokens": cfg.get("max_tokens"),
        "stateful": cfg.get("stateful"),
        "tools": cfg.get("tools", []),
    }


@tool(
    description=(
        "Create or update an agent definition in agents.yaml. "
        "Pass only the fields you want to set or change — omitted fields are left unchanged on update, "
        "or get sensible defaults on create. "
        "After calling this tool, the session must be reloaded (/reload) for changes to take effect."
    ),
    parameters={
        "type": "object",
        "properties": {
            "agent_name": {
                "type": "string",
                "description": "Short name for the agent (e.g. 'evelyn'). Used as the YAML key.",
            },
            "title": {
                "type": "string",
                "description": "Human-readable title (e.g. 'Director of Personnel & Systems').",
            },
            "color": {
                "type": "string",
                "description": "Rich terminal color name (e.g. 'white', 'cyan', 'green').",
            },
            "prompt_file": {
                "type": "string",
                "description": "Filename (not path) in agents/prompts/ (e.g. 'evelyn_director_of_personnel.md').",
            },
            "provider": {
                "type": "string",
                "description": "Provider name from agents.yaml providers section (e.g. 'vastblueai_thinking').",
            },
            "model": {
                "type": "string",
                "description": "Model name as the provider expects it. Empty string means use provider default.",
            },
            "temperature": {
                "type": "number",
                "description": "Sampling temperature 0.0–1.0. Use 0.6 for thinking, 0.7 for non-thinking.",
            },
            "max_tokens": {
                "type": "integer",
                "description": "Maximum response tokens.",
            },
            "stateful": {
                "type": "boolean",
                "description": "Whether the agent maintains conversation history across calls.",
            },
            "tools": {
                "type": "array",
                "items": {"type": "string"},
                "description": "List of tool names from the tools registry that this agent can use.",
            },
        },
        "required": ["agent_name"],
    },
)
def upsert_agent_definition(
    agent_name: str,
    title: str | None = None,
    color: str | None = None,
    prompt_file: str | None = None,
    provider: str | None = None,
    model: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
    stateful: bool | None = None,
    tools: list[str] | None = None,
) -> dict:
    """Create or update an agent in agents.yaml."""
    name = agent_name.lower()

    # Validate provider if given
    if provider is not None and provider != "none":
        valid_providers = list(_config._providers.keys())
        if provider not in valid_providers:
            return {
                "error": f"Unknown provider {provider!r}. Valid providers: {valid_providers}"
            }

    # Validate tools if given
    if tools is not None:
        unknown = [t for t in tools if t not in _REGISTRY]
        if unknown:
            return {
                "error": f"Unknown tools: {unknown}. Registered tools: {list(_REGISTRY)}"
            }

    ryaml, data = _load_agents_yaml()
    agents_section = data.get("agents", {})

    creating = name not in agents_section

    if creating:
        # Supply required defaults for new agents
        if not prompt_file:
            return {"error": "prompt_file is required when creating a new agent."}
        agents_section[name] = {
            "title":       title or name.capitalize(),
            "color":       color or "white",
            "prompt_file": prompt_file,
            "provider":    provider or "",
            "model":       model if model is not None else "",
            "temperature": temperature if temperature is not None else 0.7,
            "max_tokens":  max_tokens or 16384,
            "stateful":    stateful if stateful is not None else True,
        }
        if tools:
            agents_section[name]["tools"] = tools
    else:
        entry = agents_section[name]
        if title       is not None: entry["title"]       = title
        if color       is not None: entry["color"]       = color
        if prompt_file is not None: entry["prompt_file"] = prompt_file
        if provider    is not None: entry["provider"]    = provider
        if model       is not None: entry["model"]       = model
        if temperature is not None: entry["temperature"] = temperature
        if max_tokens  is not None: entry["max_tokens"]  = max_tokens
        if stateful    is not None: entry["stateful"]    = stateful
        if tools       is not None: entry["tools"]       = tools

    _save_agents_yaml(ryaml, data)
    action = "created" if creating else "updated"
    return {"ok": True, "action": action, "agent": name, "note": "Run /reload to apply changes."}


@tool(
    description="Read the system prompt file for an agent. Returns the raw Markdown content of the prompt file.",
    parameters={
        "type": "object",
        "properties": {
            "agent_name": {
                "type": "string",
                "description": "The agent's short name (e.g. 'evelyn').",
            }
        },
        "required": ["agent_name"],
    },
)
def read_agent_prompt(agent_name: str) -> dict:
    """Read and return the raw prompt file for an agent."""
    name = agent_name.lower()
    cfg = _config.AGENT_CONFIGS.get(name)
    if cfg is None:
        return {"error": f"No agent named {agent_name!r}."}
    path = _config.PROMPTS_DIR / cfg["prompt_file"]
    if not path.exists():
        return {"error": f"Prompt file not found: {path}"}
    return {"prompt_file": cfg["prompt_file"], "content": path.read_text(encoding="utf-8")}


@tool(
    description=(
        "Write or overwrite a prompt file in the agents/prompts/ directory. "
        "Use the prompt_file name exactly (e.g. 'evelyn_director_of_personnel.md'). "
        "The file will be created if it does not exist."
    ),
    parameters={
        "type": "object",
        "properties": {
            "prompt_file": {
                "type": "string",
                "description": "Filename (not path) in agents/prompts/. Must end with .md.",
            },
            "content": {
                "type": "string",
                "description": "Full Markdown content to write to the file.",
            },
        },
        "required": ["prompt_file", "content"],
    },
)
def write_agent_prompt(prompt_file: str, content: str) -> dict:
    """Write (or overwrite) a prompt file."""
    if not prompt_file.endswith(".md"):
        return {"error": "prompt_file must end with .md"}
    # Disallow path traversal
    if "/" in prompt_file or "\\" in prompt_file:
        return {"error": "prompt_file must be a filename only, not a path."}
    path = _config.PROMPTS_DIR / prompt_file
    path.write_text(content, encoding="utf-8")
    return {"ok": True, "wrote": str(path), "bytes": len(content.encode())}


@tool(
    description="List all provider definitions from agents.yaml, showing type, base_url, default_model, and extra_body settings.",
    parameters={
        "type": "object",
        "properties": {},
        "required": [],
    },
)
def list_providers() -> list[dict]:
    """Return all provider definitions."""
    result = []
    for name, pdef in _config._providers.items():
        entry = {"name": name}
        entry["type"]          = pdef.get("type", "")
        entry["base_url"]      = pdef.get("base_url", "")
        entry["default_model"] = pdef.get("default_model", "")
        entry["extra_body"]    = pdef.get("extra_body", {})
        result.append(entry)
    return result


@tool(
    description="List all tools registered in the tool registry, with their names and descriptions.",
    parameters={
        "type": "object",
        "properties": {},
        "required": [],
    },
)
def list_tools() -> list[dict]:
    """Return all registered tools with name and description."""
    return [
        {
            "name": name,
            "description": entry["schema"]["function"]["description"],
        }
        for name, entry in _REGISTRY.items()
    ]


# ---------------------------------------------------------------------------
# Task dispatch (used by Miranda — orchestrator drains this queue after her turn)
# ---------------------------------------------------------------------------

_pending_tasks: list[dict] = []


def get_pending_tasks() -> list[dict]:
    """
    Drain and return all task specs queued by dispatch_task tool calls.
    Call once per directive, after the orchestrator's tool loop completes.
    """
    tasks = _pending_tasks.copy()
    _pending_tasks.clear()
    return tasks


@tool(
    description=(
        "Dispatch a task to a Lead agent. "
        "Call this tool once per Lead you want to task — multiple calls dispatch multiple tasks in parallel. "
        "Do NOT respond with TASK BRIEF text; use this tool instead. "
        "Results will be returned to you when all Leads complete so you can synthesize."
    ),
    parameters={
        "type": "object",
        "properties": {
            "to": {
                "type": "string",
                "description": "Lead agent name (lowercase): atlas, cole, clio, iris, or evelyn.",
            },
            "directive": {
                "type": "string",
                "description": "What the Principal wants — the full task objective.",
            },
            "scope": {
                "type": "string",
                "description": "Exactly what this Lead is responsible for delivering.",
            },
            "constraints": {
                "type": "string",
                "description": "Deadlines, depth limits, format requirements, or assumptions to hold. Use 'NONE' if none.",
            },
            "dependencies": {
                "type": "string",
                "description": "Other tasks this depends on, or 'NONE'.",
            },
            "return_format": {
                "type": "string",
                "description": "Specific structure or emphasis you need back from this Lead beyond the standard STATUS / SUMMARY / FINDINGS / OPEN ITEMS format.",
            },
        },
        "required": ["to", "directive", "scope"],
    },
)
def dispatch_task(
    to: str,
    directive: str,
    scope: str,
    constraints: str = "NONE",
    dependencies: str = "NONE",
    return_format: str = "",
) -> dict:
    """Queue a task for dispatch to a Lead agent after Miranda's turn completes."""
    lead = to.lower().strip()
    if lead not in _config.LEAD_AGENT_NAMES:
        return {
            "error": f"Unknown lead {to!r}. Valid leads: {sorted(_config.LEAD_AGENT_NAMES)}"
        }
    _pending_tasks.append({
        "to":            lead,
        "directive":     directive,
        "scope":         scope,
        "constraints":   constraints,
        "dependencies":  dependencies,
        "return_format": return_format,
    })
    return {
        "queued": True,
        "to":     lead,
        "note":   "Task queued. You will receive lead outputs to synthesize when all tasks complete.",
    }


# ---------------------------------------------------------------------------
# Knowledge base tools
# ---------------------------------------------------------------------------

@tool(
    description=(
        "Write or overwrite a Markdown document in the knowledge base. "
        "The document is stored as a .md file on the filesystem. "
        "Existing changelogs are preserved and a new entry is appended automatically. "
        "Use Markdown with a top-level '# Title' heading as the first line."
    ),
    parameters={
        "type": "object",
        "properties": {
            "filename": {
                "type": "string",
                "description": "Filename for the document (e.g. 'project_goals.md'). Must end with .md.",
            },
            "content": {
                "type": "string",
                "description": "Full Markdown content to write (excluding changelog — that is managed automatically).",
            },
            "collection": {
                "type": "string",
                "description": "Collection (subdirectory) to store the document in. Defaults to the configured default collection.",
            },
            "agent": {
                "type": "string",
                "description": "Name of the agent writing this document — recorded in the changelog.",
            },
            "changelog_entry": {
                "type": "string",
                "description": "Short description of what changed, for the changelog (e.g. 'Added Q2 goals'). Auto-generated if omitted.",
            },
        },
        "required": ["filename", "content"],
    },
)
def kb_write(
    filename: str,
    content: str,
    collection: str = "",
    agent: str = "",
    changelog_entry: str = "",
) -> dict:
    """Write or overwrite a KB Markdown document."""
    import knowledge_base as _kb
    return _kb.kb_write(
        filename=filename,
        content=content,
        collection=collection,
        agent=agent,
        changelog_entry=changelog_entry,
    )


@tool(
    description=(
        "Read a Markdown document from the knowledge base. "
        "Returns the document body and its changelog history separately."
    ),
    parameters={
        "type": "object",
        "properties": {
            "filename": {
                "type": "string",
                "description": "Filename of the document to read (e.g. 'project_goals.md').",
            },
            "collection": {
                "type": "string",
                "description": "Collection containing the document. Defaults to the configured default collection.",
            },
        },
        "required": ["filename"],
    },
)
def kb_read(filename: str, collection: str = "") -> dict:
    """Read a KB document, returning body and changelog separately."""
    import knowledge_base as _kb
    return _kb.kb_read(filename=filename, collection=collection)


@tool(
    description=(
        "Search the knowledge base by semantic similarity. "
        "Syncs the embedding index from the filesystem before querying. "
        "Returns matching documents ordered by relevance (lower distance = more similar)."
    ),
    parameters={
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "Natural-language search query.",
            },
            "n_results": {
                "type": "integer",
                "description": "Maximum number of results to return. Defaults to 5.",
            },
            "collection": {
                "type": "string",
                "description": "Collection to search. Defaults to the configured default collection.",
            },
        },
        "required": ["query"],
    },
)
def kb_search(query: str, n_results: int = 5, collection: str = "") -> list[dict]:
    """Semantic search in the knowledge base."""
    import knowledge_base as _kb
    return _kb.kb_search(query=query, n_results=n_results, collection=collection)


@tool(
    description="List all Markdown documents in a knowledge base collection with their titles and modification times.",
    parameters={
        "type": "object",
        "properties": {
            "collection": {
                "type": "string",
                "description": "Collection to list. Defaults to the configured default collection.",
            },
        },
        "required": [],
    },
)
def kb_list(collection: str = "") -> list[dict]:
    """List all documents in a KB collection."""
    import knowledge_base as _kb
    return _kb.kb_list(collection=collection)


@tool(
    description=(
        "Delete a Markdown document from the knowledge base. "
        "The file is removed from disk; the embedding index is cleaned up on next search."
    ),
    parameters={
        "type": "object",
        "properties": {
            "filename": {
                "type": "string",
                "description": "Filename of the document to delete (e.g. 'old_notes.md').",
            },
            "collection": {
                "type": "string",
                "description": "Collection containing the document. Defaults to the configured default collection.",
            },
        },
        "required": ["filename"],
    },
)
def kb_delete(filename: str, collection: str = "") -> dict:
    """Delete a KB document from the filesystem."""
    import knowledge_base as _kb
    return _kb.kb_delete(filename=filename, collection=collection)


@tool(
    description="List all collections in the knowledge base with their document counts.",
    parameters={
        "type": "object",
        "properties": {},
        "required": [],
    },
)
def kb_list_collections() -> list[dict]:
    """List all knowledge base collections."""
    import knowledge_base as _kb
    return _kb.kb_list_collections()