import json
import tempfile
import unittest
from pathlib import Path
from unittest.mock import Mock, patch

from PIL import Image

from copienator import EvaluationWorkspace, read_json
from copienator.commands import cutleft, gemini_for_labels as labels


class FullPageFallbackTests(unittest.TestCase):
    def setUp(self):
        temporary = tempfile.TemporaryDirectory()
        self.addCleanup(temporary.cleanup)
        self.workspace = EvaluationWorkspace(Path(temporary.name))
        self.workspace.cutleft_dir.mkdir()
        self.workspace.copies_dir.mkdir()
        self.images = [self.workspace.cutleft_dir / f'Copie01_{i:02d}.jpg' for i in (1, 2, 3)]
        for image in self.images:
            image.write_bytes(b'original crop')
        self.schema_path = self.workspace.cutleft_dir / 'Copie01_schema.json'
        self.schema_path.write_text(json.dumps({'columns_per_file': [1, 2, 1],
                                               'full_width_pages': [], 'total_pages': 4}))
        self.pages = [Image.new('RGB', (100, 200), color) for color in ('red', 'green', 'blue', 'white')]
        self.client = Mock()
        self.error = Mock(text='{"error":"insufficient_crop","error_reason":"Exercise heading clipped"}')
        self.success = Mock(text='{"name":"Student","list":[{"box_2d":[1,2,3,4],"label":"Ex 1"}]}')

    def process(self):
        return labels.process_copy_group(
            self.workspace, 'Copie01', self.images, client=self.client,
            labels_text='Ex 1', names_text='Student', valid_labels={'Ex 1'},
            valid_names={'Student'}, overwrite=False, sleep=lambda _: None, target_interval=0,
        )

    def test_regenerates_only_affected_pages_and_updates_metadata(self):
        with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages):
            cutleft.regenerate_full_page_image(self.workspace, self.images[1])
        self.assertEqual(self.images[0].read_bytes(), b'original crop')
        self.assertEqual(self.images[2].read_bytes(), b'original crop')
        with Image.open(self.images[1]) as result:
            self.assertEqual(result.size, (205, 200))
            self.assertGreater(result.getpixel((40, 50))[1], 100)
            self.assertGreater(result.getpixel((150, 50))[2], 200)
        schema = read_json(self.schema_path)
        self.assertEqual(schema['columns_per_file'], [1, 2, 1])
        self.assertEqual(schema['full_width_pages'], [2, 3])
        self.assertEqual([crop['right_px'] for crop in schema['page_crops']], [100, 100])
        with self.assertRaisesRegex(ValueError, 'already contains full pages'):
            cutleft.regenerate_full_page_image(self.workspace, self.images[1])

    def test_error_regenerates_and_resends_with_context(self):
        self.client.models.generate_content.side_effect = [self.success, self.error, self.success, self.success]
        with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages):
            self.assertEqual(self.process(), 3)
        calls = self.client.models.generate_content.call_args_list
        retry_parts = calls[2].kwargs['contents'][0].parts
        self.assertEqual(retry_parts[0].inline_data.data, self.images[1].read_bytes())
        self.assertIn('now contains full pages', retry_parts[1].text)
        self.assertIn('"Ex 1"', retry_parts[1].text)
        result = read_json(self.workspace.copies_dir / 'Copie01_02.json')
        self.assertEqual(set(result), {'name', 'list'})
        self.assertEqual(result['list'][0]['label'], 'Ex 1')

    def test_second_crop_error_stops_and_defers_following_images(self):
        self.client.models.generate_content.side_effect = [self.error, self.error]
        with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages), \
                self.assertRaises(labels.IncompleteCopyError) as caught:
            self.process()
        self.assertEqual(caught.exception.images, self.images)
        self.assertEqual(self.client.models.generate_content.call_count, 2)
        self.assertFalse(list(self.workspace.copies_dir.glob('*.json')))

    def test_bad_mapping_leaves_original_image_and_schema_intact(self):
        original_schema = self.schema_path.read_bytes()
        self.client.models.generate_content.return_value = self.error
        with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages[:1]), \
                self.assertRaises(labels.IncompleteCopyError):
            self.process()
        self.assertEqual(self.client.models.generate_content.call_count, 1)
        self.assertEqual(self.schema_path.read_bytes(), original_schema)
        self.assertEqual(self.images[0].read_bytes(), b'original crop')

    def test_full_page_retry_gets_fresh_bounded_attempt_budget(self):
        self.client.models.generate_content.side_effect = [TimeoutError(), TimeoutError(), self.error,
                                                          TimeoutError(), TimeoutError(), self.success,
                                                          self.success, self.success]
        with patch.object(cutleft, 'get_pdf_pages', return_value=self.pages):
            self.assertEqual(self.process(), 3)
        self.assertEqual(self.client.models.generate_content.call_count, 8)

    def test_resumed_full_pages_use_correct_prompt_without_regenerating(self):
        schema = read_json(self.schema_path)
        schema['full_width_pages'] = [1]
        self.schema_path.write_text(json.dumps(schema))
        self.client.models.generate_content.return_value = self.error
        with patch.object(cutleft, 'get_pdf_pages') as render, self.assertRaises(labels.IncompleteCopyError):
            self.process()
        render.assert_not_called()
        prompt = self.client.models.generate_content.call_args.kwargs['contents'][0].parts[1].text
        self.assertIn('now contains full pages', prompt)

    def test_missing_success_payload_is_not_accepted(self):
        with self.assertRaises(ValueError):
            labels.AnnotationData.model_validate_json('{}')