from __future__ import annotations

import importlib.util
import os
from pathlib import Path
from unittest import TestCase, mock


BOT_PATH = Path(__file__).resolve().parents[1] / "atlasbot" / "bot.py"


def load_bot_module():
    env = {
        "BOT_USER": "atlas-smart",
        "BOT_PASS": "smart-pass",
        "BOT_USER_QUICK": "atlas-quick",
        "BOT_PASS_QUICK": "quick-pass",
        "BOT_USER_SMART": "atlas-smart",
        "BOT_PASS_SMART": "smart-pass",
        "BOT_USER_GENIUS": "atlas-genius",
        "BOT_PASS_GENIUS": "genius-pass",
        "OLLAMA_URL": "http://ollama.invalid",
        "OLLAMA_MODEL": "base-model",
        "ATLASBOT_MODEL_FAST": "fast-model",
        "ATLASBOT_MODEL_SMART": "smart-model",
        "ATLASBOT_MODEL_GENIUS": "genius-model",
        "ATLASBOT_QUICK_TIME_BUDGET_SEC": "15",
        "ATLASBOT_SMART_TIME_BUDGET_SEC": "45",
        "ATLASBOT_GENIUS_TIME_BUDGET_SEC": "180",
        "KB_DIR": "",
        "VM_URL": "http://vm.invalid",
        "ARIADNE_STATE_URL": "",
        "ARIADNE_STATE_TOKEN": "",
    }
    with mock.patch.dict(os.environ, env, clear=False):
        spec = importlib.util.spec_from_file_location("atlasbot_bot", BOT_PATH)
        module = importlib.util.module_from_spec(spec)
        assert spec.loader is not None
        spec.loader.exec_module(module)
        return module


class AtlasbotModeTests(TestCase):
    def setUp(self):
        self.bot = load_bot_module()

    def test_bot_accounts_include_genius_mode(self):
        accounts = self.bot._bot_accounts()
        by_user = {account["user"]: account["mode"] for account in accounts}

        self.assertEqual(by_user["atlas-quick"], "fast")
        self.assertEqual(by_user["atlas-smart"], "smart")
        self.assertEqual(by_user["atlas-genius"], "genius")

    def test_objective_cluster_question_uses_fact_pack_without_llm(self):
        fact_lines = [
            "hottest_cpu: longhorn-system (6.69)",
            "hottest_ram: longhorn-system (36.05 GB)",
        ]

        with (
            mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
            mock.patch.object(self.bot, "_ollama_call", side_effect=AssertionError("LLM should not be called")),
        ):
            reply = self.bot.open_ended_answer(
                "what is the hottest cpu node in titan lab currently?",
                inventory=[],
                snapshot=None,
                workloads=[],
                history_lines=[],
                mode="smart",
                allow_tools=True,
            )

        self.assertIn("longhorn-system", reply)
        self.assertIn("Confidence:", reply)

    def test_subjective_genius_answer_uses_genius_model(self):
        fact_lines = [
            "hottest_cpu: longhorn-system (6.69)",
            "worker_nodes: titan-01, titan-02, titan-03",
        ]
        captured: dict[str, object] = {}

        def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None, options=None):
            captured["model"] = model
            captured["timeout"] = timeout
            captured["context"] = context
            return "The worker spread stands out because Titan keeps meaningful capacity on the same cluster. Confidence: high"

        with (
            mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
            mock.patch.object(self.bot, "_ollama_call", side_effect=fake_ollama_call),
        ):
            reply = self.bot.open_ended_answer(
                "what stands out about titan lab?",
                inventory=[],
                snapshot=None,
                workloads=[],
                history_lines=[],
                mode="genius",
                allow_tools=True,
                context='Cluster snapshot (JSON): {"injected":true}',
            )

        self.assertIn("The worker spread stands out", reply)
        self.assertEqual(captured["model"], "genius-model")
        self.assertLessEqual(float(captured["timeout"]), 180.0)
        self.assertIn('Cluster snapshot (JSON): {"injected":true}', str(captured["context"]))

    def test_mode_timeouts_stay_within_budgets(self):
        fact_lines = [
            "hottest_cpu: longhorn-system (6.69)",
            "worker_nodes: titan-01, titan-02, titan-03",
        ]
        seen: list[tuple[str, float]] = []

        def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None, options=None):
            seen.append((str(model), float(timeout or 0)))
            return "Atlas has a clear standout because the worker spread is healthy. Confidence: high"

        with (
            mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
            mock.patch.object(self.bot, "_ollama_call", side_effect=fake_ollama_call),
        ):
            for mode in ("fast", "smart", "genius"):
                reply = self.bot.open_ended_answer(
                    "what stands out about titan lab?",
                    inventory=[],
                    snapshot=None,
                    workloads=[],
                    history_lines=[],
                    mode=mode,
                    allow_tools=True,
                )
                self.assertIn("Confidence:", reply)

        self.assertEqual([model for model, _ in seen], ["fast-model", "smart-model", "genius-model"])
        self.assertLessEqual(seen[0][1], 15.0)
        self.assertLessEqual(seen[1][1], 45.0)
        self.assertLessEqual(seen[2][1], 180.0)

    def test_llm_timeout_still_returns_a_conclusion(self):
        fact_lines = [
            "worker_nodes: titan-01, titan-02, titan-03",
            "hottest_cpu: longhorn-system (6.69)",
        ]

        with (
            mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
            mock.patch.object(self.bot, "_ollama_call", side_effect=TimeoutError("simulated timeout")),
        ):
            reply = self.bot.open_ended_answer(
                "what stands out about the worker nodes?",
                inventory=[],
                snapshot=None,
                workloads=[],
                history_lines=[],
                mode="genius",
                allow_tools=True,
            )

        self.assertIn("worker nodes", reply.lower())
        self.assertIn("Confidence:", reply)

    def test_longhorn_rpi4_subset_beats_generic_rpi4_list(self):
        fact_lines = [
            "rpi4: titan-12, titan-13, titan-14, titan-15, titan-17, titan-18, titan-19",
            "rpi4 armbian longhorn: titan-13/15/17/19",
        ]

        with (
            mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
            mock.patch.object(self.bot, "_ollama_call", side_effect=AssertionError("LLM should not be called")),
        ):
            reply = self.bot.open_ended_answer(
                "which nodes in titan are the rpi4 longhorn nodes?",
                inventory=[],
                snapshot=None,
                workloads=[],
                history_lines=[],
                mode="smart",
                allow_tools=True,
            )

        self.assertIn("titan-13", reply)
        self.assertIn("titan-15", reply)
        self.assertIn("titan-17", reply)
        self.assertIn("titan-19", reply)
        self.assertNotIn("titan-12", reply)
        self.assertNotIn("titan-14", reply)
        self.assertNotIn("titan-18", reply)
        self.assertIn("Confidence:", reply)

    def test_longhorn_query_uses_kb_context_when_snapshot_only_has_generic_rpi4(self):
        fact_lines = [
            "rpi4: titan-12, titan-13, titan-14, titan-15, titan-17, titan-18, titan-19",
        ]
        kb_detail = (
            "Atlas KB (retrieved):\n"
            "- metis (knowledge/metis.md)\n"
            "rpi4 armbian longhorn: titan-13/15/17/19"
        )

        with (
            mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
            mock.patch.object(self.bot, "kb_retrieve", return_value=kb_detail),
            mock.patch.object(self.bot, "_ollama_call", side_effect=AssertionError("LLM should not be called")),
        ):
            reply = self.bot.open_ended_answer(
                "which nodes in titan are the rpi4 longhorn nodes?",
                inventory=[],
                snapshot=None,
                workloads=[],
                history_lines=[],
                mode="smart",
                allow_tools=True,
            )

        self.assertIn("titan-13", reply)
        self.assertIn("titan-15", reply)
        self.assertIn("titan-17", reply)
        self.assertIn("titan-19", reply)
        self.assertNotIn("titan-12", reply)
        self.assertNotIn("titan-14", reply)
        self.assertNotIn("titan-18", reply)
        self.assertIn("Confidence:", reply)
