"""Tests for model-aware agent tools (tools_ext/model_tools.py).

These tests exercise the tools directly without an LLM — verifying they
return structured data the agent can interpret.
"""

from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path
from unittest.mock import patch

import pytest
import yaml


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _skip_if_no_neutron_os():
    """Skip test if NeutronOS model_corral is not importable."""
    try:
        import neutron_os.extensions.builtins.model_corral  # noqa: F401
    except ImportError:
        pytest.skip("NeutronOS model_corral not installed")


def _execute(name: str, params: dict) -> dict:
    from axiom.extensions.builtins.neut_agent.tools_ext.model_tools import execute

    return execute(name, params)


# ---------------------------------------------------------------------------
# Tool schema validation
# ---------------------------------------------------------------------------


class TestToolSchemas:
    """Verify tool definitions are well-formed."""

    def test_tools_list_is_nonempty(self):
        from axiom.extensions.builtins.neut_agent.tools_ext.model_tools import TOOLS

        assert len(TOOLS) >= 7

    def test_each_tool_has_required_fields(self):
        from axiom.extensions.builtins.neut_agent.tools_ext.model_tools import TOOLS

        for tool in TOOLS:
            assert tool.name, "Tool must have a name"
            assert tool.description, "Tool must have a description"
            assert tool.category is not None, "Tool must have a category"
            assert "type" in tool.parameters
            assert tool.parameters["type"] == "object"

    def test_tool_names_are_unique(self):
        from axiom.extensions.builtins.neut_agent.tools_ext.model_tools import TOOLS

        names = [t.name for t in TOOLS]
        assert len(names) == len(set(names)), f"Duplicate tool names: {names}"

    def test_parameters_are_valid_json_schema(self):
        """Each tool's parameters dict should be serializable as JSON Schema."""
        from axiom.extensions.builtins.neut_agent.tools_ext.model_tools import TOOLS

        for tool in TOOLS:
            # Must be JSON-serializable
            serialized = json.dumps(tool.parameters)
            parsed = json.loads(serialized)
            assert parsed["type"] == "object"

    def test_tools_discovered_by_scanner(self):
        """The tools_ext scanner should discover model_tools."""
        from axiom.extensions.builtins.neut_agent.tools import _scan_extensions

        ext_tools = _scan_extensions()
        assert "model_list" in ext_tools or "model_materials" in ext_tools


# ---------------------------------------------------------------------------
# Material tools (no DB needed — pure in-memory)
# ---------------------------------------------------------------------------


class TestMaterialTools:
    """Test material-related tools that use the in-memory materials_db."""

    def test_model_materials_returns_all(self):
        _skip_if_no_neutron_os()
        result = _execute("model_materials", {})
        assert "materials" in result
        assert "count" in result
        assert result["count"] >= 11  # 11 builtin materials

    def test_model_materials_by_category_fuel(self):
        _skip_if_no_neutron_os()
        result = _execute("model_materials", {"category": "fuel"})
        assert result["count"] >= 3  # UZrH-20, UO2-3.1, UO2-4.95, MSRE-salt
        for mat in result["materials"]:
            assert mat["category"] == "fuel"

    def test_model_materials_search(self):
        _skip_if_no_neutron_os()
        result = _execute("model_materials", {"query": "uranium"})
        assert result["count"] >= 1

    def test_material_card_uzrh20_mcnp(self):
        _skip_if_no_neutron_os()
        result = _execute("material_card", {"name": "UZrH-20"})
        assert "card" in result
        assert "material" in result
        card = result["card"]
        assert "m1" in card  # default mat_number=1
        assert "92235" in card  # U-235 ZAID
        assert "UZrH-20" in card

    def test_material_card_uzrh20_mpact(self):
        _skip_if_no_neutron_os()
        result = _execute("material_card", {"name": "UZrH-20", "format": "mpact"})
        assert "card" in result
        assert "mat UZrH-20" in result["card"]

    def test_material_card_custom_mat_number(self):
        _skip_if_no_neutron_os()
        result = _execute("material_card", {"name": "H2O", "mat_number": 5})
        assert "m5" in result["card"]

    def test_material_card_not_found(self):
        _skip_if_no_neutron_os()
        result = _execute("material_card", {"name": "nonexistent"})
        assert "error" in result
        assert "not found" in result["error"].lower()


# ---------------------------------------------------------------------------
# Model registry tools (need SQLite service)
# ---------------------------------------------------------------------------


class TestModelRegistryTools:
    """Test model registry tools using the SQLite fallback."""

    def test_model_list_empty(self):
        _skip_if_no_neutron_os()
        # Force a fresh SQLite DB so the registry is empty
        with tempfile.TemporaryDirectory() as tmp:
            db_path = Path(tmp) / "test.db"
            db_url = f"sqlite:///{db_path}"
            # Reset the cached service
            import neutron_os.extensions.builtins.model_corral.cli as cli_mod

            old_service = cli_mod._SERVICE
            cli_mod._SERVICE = None
            try:
                with patch.dict(os.environ, {"AXIOM_DB_URL": db_url}):
                    result = _execute("model_list", {})
                assert "models" in result
                assert result["count"] == 0
                assert result["models"] == []
            finally:
                cli_mod._SERVICE = old_service

    def test_model_show_not_found(self):
        _skip_if_no_neutron_os()
        with tempfile.TemporaryDirectory() as tmp:
            db_path = Path(tmp) / "test.db"
            db_url = f"sqlite:///{db_path}"
            import neutron_os.extensions.builtins.model_corral.cli as cli_mod

            old_service = cli_mod._SERVICE
            cli_mod._SERVICE = None
            try:
                with patch.dict(os.environ, {"AXIOM_DB_URL": db_url}):
                    result = _execute("model_show", {"model_id": "nonexistent"})
                assert "error" in result
            finally:
                cli_mod._SERVICE = old_service

    def test_model_search_empty(self):
        _skip_if_no_neutron_os()
        with tempfile.TemporaryDirectory() as tmp:
            db_path = Path(tmp) / "test.db"
            db_url = f"sqlite:///{db_path}"
            import neutron_os.extensions.builtins.model_corral.cli as cli_mod

            old_service = cli_mod._SERVICE
            cli_mod._SERVICE = None
            try:
                with patch.dict(os.environ, {"AXIOM_DB_URL": db_url}):
                    result = _execute("model_search", {"query": "triga"})
                assert result["count"] == 0
            finally:
                cli_mod._SERVICE = old_service


# ---------------------------------------------------------------------------
# Facility tools
# ---------------------------------------------------------------------------


class TestFacilityTools:
    """Test facility pack tools."""

    def test_facility_list(self):
        _skip_if_no_neutron_os()
        result = _execute("facility_list", {})
        assert "packs" in result
        assert "count" in result
        # Builtin packs should include NETL-TRIGA, MSRE, PWR-generic
        assert result["count"] >= 3
        names = [p["name"] for p in result["packs"]]
        assert "NETL-TRIGA" in names

    def test_facility_materials_netl_triga(self):
        _skip_if_no_neutron_os()
        result = _execute("facility_materials", {"pack_name": "NETL-TRIGA"})
        if "error" in result:
            pytest.skip(f"NETL-TRIGA pack issue: {result['error']}")
        assert "materials" in result
        assert "pack" in result
        assert result["pack"]["name"] == "NETL-TRIGA"

    def test_facility_materials_not_found(self):
        _skip_if_no_neutron_os()
        result = _execute("facility_materials", {"pack_name": "nonexistent-pack"})
        assert "error" in result
        assert "not found" in result["error"].lower()


# ---------------------------------------------------------------------------
# Model generate tool
# ---------------------------------------------------------------------------


class TestModelGenerateTool:
    """Test model_generate tool with a temp model directory."""

    def test_generate_with_model_yaml(self):
        _skip_if_no_neutron_os()
        with tempfile.TemporaryDirectory() as tmp:
            model_dir = Path(tmp)
            model_yaml = {
                "model_id": "test-model",
                "version": "1.0.0",
                "physics_code": "mcnp",
                "materials": [
                    {"name": "H2O", "mat_number": 1},
                    {"name": "UZrH-20", "mat_number": 2},
                ],
            }
            (model_dir / "model.yaml").write_text(
                yaml.dump(model_yaml), encoding="utf-8"
            )
            result = _execute("model_generate", {"model_path": str(model_dir)})
            assert "error" not in result or result.get("output")
            if "output" in result:
                assert "m1" in result["output"] or "m2" in result["output"]

    def test_generate_missing_model_yaml(self):
        _skip_if_no_neutron_os()
        with tempfile.TemporaryDirectory() as tmp:
            result = _execute("model_generate", {"model_path": tmp})
            assert "error" in result


# ---------------------------------------------------------------------------
# Graceful degradation without NeutronOS
# ---------------------------------------------------------------------------


class TestWithoutNeutronOS:
    """Verify tools return helpful errors when NeutronOS is not available."""

    def test_graceful_error(self):
        import axiom.extensions.builtins.neut_agent.tools_ext.model_tools as mod

        original = mod._NEUTRON_OS_AVAILABLE
        try:
            mod._NEUTRON_OS_AVAILABLE = False
            result = mod.execute("model_list", {})
            assert "error" in result
            assert "not installed" in result["error"].lower()

            result = mod.execute("material_card", {"name": "H2O"})
            assert "error" in result
        finally:
            mod._NEUTRON_OS_AVAILABLE = original


# ---------------------------------------------------------------------------
# Unknown tool name
# ---------------------------------------------------------------------------


class TestUnknownTool:
    def test_unknown_tool(self):
        result = _execute("nonexistent_tool", {})
        assert "error" in result
        assert "Unknown tool" in result["error"]
