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):
            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,
            )

        self.assertIn("The worker spread stands out", reply)
        self.assertEqual(captured["model"], "genius-model")
        self.assertLessEqual(float(captured["timeout"]), 180.0)