from pathlib import Path
import sys
import unittest
from unittest import mock


ROOT = Path(__file__).resolve().parents[1]
APP_DIR = ROOT / "app"
sys.path.insert(0, str(APP_DIR))

import keycheck_runner
from keycheckers.gemini import geminiKeycheck as gemini
from keycheckers.openai import Keycheck as openai
from keycheckers.xai import xaiKeycheck as xai


class ModelProbePreferenceTests(unittest.TestCase):
    def test_gemini_prefers_current_pro_then_current_flash(self):
        models = ["gemini-2.5-flash", "gemini-3.7-flash", "gemini-3.1-pro-preview"]
        self.assertEqual(gemini.choose_probe_model(models), "gemini-3.1-pro-preview")
        self.assertEqual(
            gemini.choose_probe_model(["gemini-2.5-flash", "gemini-3.7-flash"]),
            "gemini-3.7-flash",
        )
        self.assertEqual(gemini.choose_probe_model(["gemini-4-preview"]), "gemini-4-preview")

    def test_xai_prefers_grok_4_6_and_falls_back_to_advertised_model(self):
        self.assertEqual(xai.choose_chat_model(["grok-4.5", "grok-4.6"]), "grok-4.6")
        self.assertEqual(
            xai.choose_chat_model(["grok-4.20-0309-non-reasoning", "grok-4.5"]),
            "grok-4.5",
        )

    def test_openai_prefers_advertised_gpt_5_6_sol(self):
        models = ["gpt-5", "gpt-5.6-terra", "gpt-5.6-sol", "o3-pro"]
        self.assertEqual(openai.choose_probe_model(models), "gpt-5.6-sol")
        self.assertEqual(openai.choose_probe_model(["gpt-5", "o3-pro"]), "gpt-5")
        self.assertEqual(openai.choose_probe_model(["gpt-4o", "text-embedding-3-large"]), "gpt-4o")
        self.assertEqual(
            openai.reportable_target_models(models),
            {"gpt-5", "gpt-5.6-terra", "gpt-5.6-sol", "o3-pro"},
        )

    def test_openai_probe_uses_supported_completion_budget(self):
        response = mock.Mock(status_code=200, text='{"choices":[{}]}')
        response.json.return_value = {"choices": [{}], "service_tier": "default"}
        with mock.patch.object(openai.requests, "post", return_value=response) as request:
            status, tier = openai.check_balance_and_tier("fixture-key", "gpt-5.6-sol", None)

        self.assertEqual((status, tier), ("ok", "default"))
        payload = request.call_args.kwargs["json"]
        self.assertEqual(payload["model"], "gpt-5.6-sol")
        self.assertEqual(payload["max_completion_tokens"], 16)
        self.assertNotIn("max_tokens", payload)

    def test_postgres_projection_detail_includes_probe_and_inventory(self):
        detail = keycheck_runner.postgres_status_projection_detail({
            "message": "accepted",
            "llm_probe_model": "gpt-5.6-sol",
            "llm_probe_status": "GENERATION_OK",
            "model_count": 3,
            "model_inventory": ["gpt-5.6-sol", "gpt-5", "o3-pro"],
        })
        self.assertEqual(
            detail,
            "accepted; probe_model=gpt-5.6-sol; probe_status=GENERATION_OK; "
            "model_count=3; models=gpt-5.6-sol,gpt-5,o3-pro",
        )


if __name__ == "__main__":
    unittest.main()