import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch

from google.genai import types

from copienator import configuration, prompting
from copienator.commands import correction, gemini_for_enonce, gemini_for_labels
from copienator.gemini_config import thinking_config


class GeminiConfigTests(unittest.TestCase):
    def assert_request_config(self, config, level):
        data = config.model_dump(mode="json", by_alias=True, exclude_none=True)
        for key in ("temperature", "topP", "topK", "thinkingBudget"):
            self.assertNotIn(key, data)
            self.assertNotIn(key, data.get("thinkingConfig", {}))
        if level is None:
            self.assertNotIn("thinkingConfig", data)
        else:
            self.assertEqual(data["thinkingConfig"], {"thinkingLevel": level.upper()})

    def test_grading_live_and_batch_settings(self):
        with tempfile.TemporaryDirectory() as directory:
            image = Path(directory) / "answer.jpg"
            image.write_bytes(b"image")
            for level in ("minimal", "low", "medium", "high", None):
                with self.subTest(level=level), patch.object(prompting, "make_prompt", return_value="Grade"):
                    _, config = prompting.generate_request(directory, image, "Ex 1", thinking_level=level)
                    self.assert_request_config(config, level)
                    batch = prompting.batch_generation_config(level)
                    self.assert_request_config(types.GenerateContentConfig.model_validate(batch), level)
                    self.assertEqual(batch["responseSchema"], prompting.UNROLLED_SCHEMA)

    def test_labels_and_statement_use_their_own_settings(self):
        with tempfile.TemporaryDirectory() as directory:
            image = Path(directory) / "answer.jpg"
            image.write_bytes(b"image")
            with patch.object(configuration, "THINKING_LEVEL_FOR_LABEL", "low"):
                _, config = gemini_for_labels.generate_request(image, "Ex 1", "Student", [], [])
                self.assert_request_config(config, "low")
        client = Mock()
        with patch.object(configuration, "THINKING_LEVEL_LITE", "high"):
            gemini_for_enonce._request_rubrics(client, "Question")
        self.assert_request_config(client.models.generate_content.call_args.kwargs["config"], "high")

    def test_correction_helpers_use_flash_level(self):
        with patch.object(configuration, "THINKING_LEVEL_FLASH", "low"), patch.object(
            prompting, "get_single_image_bytes", return_value=b"image"
        ):
            requests = [
                prompting.request_for_box_correction("answer.pdf", []),
                prompting.request_for_wrong_label("answer.pdf", "Ex 1", "Question", "Ex 1"),
                prompting.request_for_additional_answer("answer.pdf", "Ex 1", "Question", "Ex 1"),
            ]
        for _, config in requests:
            self.assert_request_config(config, "low")

    def test_quota_fallback_switches_thinking_level_without_mutating_request(self):
        client = Mock()
        client.models.generate_content_stream.side_effect = [
            RuntimeError("429 daily quota exhausted"), [SimpleNamespace(text="result")],
        ]
        config = types.GenerateContentConfig(thinking_config=thinking_config("high"))
        with patch.object(correction, "client", client), patch.object(
            correction, "MODEL_ID_pro", "pro-model"
        ), patch.object(correction, "pro_quota_exhausted", False), patch.object(
            configuration, "THINKING_LEVEL_FLASH", "low"
        ), patch.object(correction, "tprint"), patch.object(correction, "stop_requested") as stop:
            stop.is_set.return_value = False
            result = correction.call_gemini_with_retries("pro-model", [], config, "flash-model")
        self.assertEqual(result, "result")
        calls = client.models.generate_content_stream.call_args_list
        self.assert_request_config(calls[0].kwargs["config"], "high")
        self.assert_request_config(calls[1].kwargs["config"], "low")
        self.assertEqual(calls[1].kwargs["model"], "flash-model")
        self.assert_request_config(config, "high")

    def test_invalid_level_fails_before_submission(self):
        with self.assertRaisesRegex(ValueError, "Invalid Gemini thinking level"):
            thinking_config("ultra")