"""Model-aware agent tools — query models, materials, and facility packs.

These tools bridge the NeutronOS Model Corral into the chat agent,
letting the LLM answer questions about registered physics models,
verified nuclear materials, and installed facility packs.

Gracefully degrades when NeutronOS is not installed (Axiom-only mode).
"""

from __future__ import annotations

from axiom.infra.orchestrator.actions import ActionCategory

from ..tools import ToolDef

# ---------------------------------------------------------------------------
# NeutronOS availability check
# ---------------------------------------------------------------------------

_NEUTRON_OS_AVAILABLE = False
try:
    import neutron_os.extensions.builtins.model_corral  # noqa: F401

    _NEUTRON_OS_AVAILABLE = True
except ImportError:
    pass


def _no_neutron_os() -> dict:
    return {
        "error": (
            "NeutronOS model_corral extension is not installed. "
            "Install NeutronOS to use model and material tools."
        )
    }


# ---------------------------------------------------------------------------
# Tool definitions
# ---------------------------------------------------------------------------

TOOLS = [
    ToolDef(
        name="model_list",
        description=(
            "List all physics models in the local registry. Use when the user "
            "asks about their models, what's registered, or wants to see their "
            "model collection."
        ),
        category=ActionCategory.READ,
        parameters={
            "type": "object",
            "properties": {
                "reactor_type": {
                    "type": "string",
                    "description": "Filter by reactor type (e.g., TRIGA, PWR, MSR)",
                },
                "physics_code": {
                    "type": "string",
                    "description": "Filter by physics code (e.g., MCNP, MPACT, Serpent)",
                },
                "status": {
                    "type": "string",
                    "description": "Filter by status (e.g., draft, validated, published)",
                },
            },
        },
    ),
    ToolDef(
        name="model_show",
        description=(
            "Get full details for a specific physics model including versions, "
            "validation status, and metadata. Use when the user asks about a "
            "specific model by name or ID."
        ),
        category=ActionCategory.READ,
        parameters={
            "type": "object",
            "properties": {
                "model_id": {
                    "type": "string",
                    "description": "Model identifier (e.g., 'triga-netl-mcnp-v3')",
                },
            },
            "required": ["model_id"],
        },
    ),
    ToolDef(
        name="model_search",
        description=(
            "Search the physics model registry by keyword. Searches model names, "
            "descriptions, tags, reactor types, and physics codes."
        ),
        category=ActionCategory.READ,
        parameters={
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query (e.g., 'TRIGA transient MCNP')",
                },
            },
            "required": ["query"],
        },
    ),
    ToolDef(
        name="model_materials",
        description=(
            "Search or list verified nuclear material compositions from the "
            "materials database. Use when the user asks about available materials, "
            "material properties, or wants to find a specific composition. "
            "The database contains verified compositions from authoritative "
            "nuclear data references."
        ),
        category=ActionCategory.READ,
        parameters={
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search term (e.g., 'uranium', 'steel', 'boron')",
                },
                "category": {
                    "type": "string",
                    "enum": ["fuel", "moderator", "coolant", "structural", "absorber", "other"],
                    "description": "Filter by material category",
                },
            },
        },
    ),
    ToolDef(
        name="material_card",
        description=(
            "Generate an MCNP or MPACT material card for a verified nuclear "
            "material. Use when the user asks for material cards, isotope "
            "compositions, or MCNP/MPACT input for a specific material like "
            "UZrH-20, H2O, B4C, SS304, etc."
        ),
        category=ActionCategory.READ,
        parameters={
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "Material name (e.g., UZrH-20, H2O, SS304, B4C)",
                },
                "format": {
                    "type": "string",
                    "enum": ["mcnp", "mpact"],
                    "description": "Output format (default: mcnp)",
                },
                "mat_number": {
                    "type": "integer",
                    "description": "MCNP material number (default: 1)",
                },
            },
            "required": ["name"],
        },
    ),
    ToolDef(
        name="model_generate",
        description=(
            "Generate all material cards for a model from its model.yaml. "
            "Use when the user wants to generate MCNP or MPACT material "
            "definitions for an entire model directory."
        ),
        category=ActionCategory.READ,
        parameters={
            "type": "object",
            "properties": {
                "model_path": {
                    "type": "string",
                    "description": "Path to the model directory containing model.yaml",
                },
                "format": {
                    "type": "string",
                    "enum": ["mcnp", "mpact"],
                    "description": "Output format (default: mcnp)",
                },
            },
            "required": ["model_path"],
        },
    ),
    ToolDef(
        name="facility_list",
        description=(
            "List installed facility packs. Facility packs contain materials, "
            "templates, and parameters for specific reactor facilities. Use when "
            "the user asks what facilities or facility packs are available."
        ),
        category=ActionCategory.READ,
        parameters={"type": "object", "properties": {}},
    ),
    ToolDef(
        name="facility_materials",
        description=(
            "List materials from a specific facility pack. Use when the user "
            "asks what materials a particular facility (e.g., NETL-TRIGA, MSRE) "
            "provides."
        ),
        category=ActionCategory.READ,
        parameters={
            "type": "object",
            "properties": {
                "pack_name": {
                    "type": "string",
                    "description": "Facility pack name (e.g., NETL-TRIGA, MSRE, PWR-generic)",
                },
            },
            "required": ["pack_name"],
        },
    ),
]


# ---------------------------------------------------------------------------
# Execution
# ---------------------------------------------------------------------------


def execute(name: str, params: dict) -> dict:
    """Dispatch a model/material/facility tool call."""
    if name == "model_list":
        return _tool_model_list(params)
    if name == "model_show":
        return _tool_model_show(params)
    if name == "model_search":
        return _tool_model_search(params)
    if name == "model_materials":
        return _tool_model_materials(params)
    if name == "material_card":
        return _tool_material_card(params)
    if name == "model_generate":
        return _tool_model_generate(params)
    if name == "facility_list":
        return _tool_facility_list(params)
    if name == "facility_materials":
        return _tool_facility_materials(params)
    return {"error": f"Unknown tool: {name}"}


# ---------------------------------------------------------------------------
# Tool implementations
# ---------------------------------------------------------------------------


def _tool_model_list(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from neutron_os.extensions.builtins.model_corral.cli import _get_service

    svc = _get_service()
    models = svc.list_models(
        reactor_type=params.get("reactor_type"),
        physics_code=params.get("physics_code"),
        status=params.get("status"),
    )
    return {"models": models, "count": len(models)}


def _tool_model_show(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from neutron_os.extensions.builtins.model_corral.cli import _get_service

    svc = _get_service()
    info = svc.show(params["model_id"])
    if info is None:
        return {"error": f"Model not found: {params['model_id']}"}
    return info


def _tool_model_search(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from neutron_os.extensions.builtins.model_corral.cli import _get_service

    svc = _get_service()
    results = svc.search(params.get("query", ""))
    return {"count": len(results), "models": results}


def _tool_model_materials(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from neutron_os.extensions.builtins.model_corral.materials_db import (
        list_materials,
        search_materials,
    )

    query = params.get("query", "")
    category = params.get("category", "")

    if query:
        mats = search_materials(query)
    elif category:
        mats = list_materials(category=category)
    else:
        mats = list_materials()

    return {"materials": [m.to_dict() for m in mats], "count": len(mats)}


def _tool_material_card(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from neutron_os.extensions.builtins.model_corral.materials_db import get_material

    mat = get_material(params["name"])
    if mat is None:
        return {"error": f"Material not found: {params['name']}"}

    fmt = params.get("format", "mcnp")
    mat_number = params.get("mat_number", 1)

    if fmt == "mpact":
        return {"card": mat.mpact_card(), "material": mat.to_dict()}
    return {"card": mat.mcnp_cards(mat_number=mat_number), "material": mat.to_dict()}


def _tool_model_generate(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from pathlib import Path

    from neutron_os.extensions.builtins.model_corral.commands.generate import (
        generate_materials,
    )

    try:
        output = generate_materials(
            Path(params["model_path"]),
            output_format=params.get("format", "mcnp"),
        )
        return {"output": output, "format": params.get("format", "mcnp")}
    except Exception as e:
        return {"error": str(e)}


def _tool_facility_list(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from neutron_os.extensions.builtins.model_corral.facilities.registry import (
        discover_packs,
    )

    packs = discover_packs()
    return {"packs": [p.to_dict() for p in packs], "count": len(packs)}


def _tool_facility_materials(params: dict) -> dict:
    if not _NEUTRON_OS_AVAILABLE:
        return _no_neutron_os()
    from neutron_os.extensions.builtins.model_corral.facilities.registry import (
        get_pack,
    )
    from neutron_os.extensions.builtins.model_corral.materials_db import (
        YamlMaterialSource,
    )

    pack = get_pack(params["pack_name"])
    if pack is None:
        return {"error": f"Pack not found: {params['pack_name']}"}

    source = YamlMaterialSource(
        pack.materials_path, source_name=f"facility:{params['pack_name']}"
    )
    mats = source.load()
    return {
        "materials": [m.to_dict() for m in mats],
        "count": len(mats),
        "pack": pack.manifest.to_dict(),
    }
