import unittest
import os
import tempfile
import json
import io
import types
from PIL import Image
from pathlib import Path
from unittest.mock import patch

from qwen21.handler import cached_model_path, handler, sampling


class SamplingTests(unittest.TestCase):
    def test_native_rgba_photo_is_saved_as_rgb_without_recoloring(self):
        source = io.BytesIO()
        Image.new('RGB', (64, 64), 'white').save(source, 'PNG')
        photo = Image.new('RGBA', (64, 64), (123, 145, 167, 254))
        pipe = unittest.mock.MagicMock()
        pipe.return_value.images = [photo]
        pipe.scheduler.timesteps = range(40)
        torch = types.SimpleNamespace(
            cuda=types.SimpleNamespace(synchronize=lambda: None,
                reset_peak_memory_stats=lambda: None, get_device_name=lambda: '48 GB test GPU',
                max_memory_reserved=lambda: 20 * 1024**3),
            Generator=lambda device: types.SimpleNamespace(manual_seed=lambda seed: None))
        with patch.dict('sys.modules', {'torch': torch}), \
                patch('qwen21.handler.pipeline', return_value=pipe), \
                patch('qwen21.handler.requests.get') as get, \
                patch('qwen21.handler.upload', return_value='https://result.png') as save, \
                patch.dict(os.environ, {'QWEN21_MODEL_REVISION': 'pinned'}):
            get.return_value.content = source.getvalue()
            result = handler({'input': {'task': 'edit', 'jobs': [{
                'image_url': 'https://source.png', 'prompt': 'A complete photograph.'}]}})
        self.assertTrue(result['results'][0]['ok'])
        self.assertEqual(result['results'][0]['native_alpha_range'], (254, 254))
        saved = save.call_args.args[0]
        self.assertEqual(saved.mode, 'RGB')
        self.assertEqual(saved.getpixel((0, 0)), (123, 145, 167))

    def test_only_the_pinned_runpod_snapshot_is_used(self):
        with tempfile.TemporaryDirectory() as root, patch.dict(os.environ, {
                'HF_HUB_CACHE': root, 'QWEN21_MODEL_REPO': 'Qwen/example',
                'QWEN21_MODEL_REVISION': 'pinned'}):
            with self.assertRaises(FileNotFoundError):
                cached_model_path()
            wrong = Path(root) / 'models--Qwen--example/snapshots/wrong'
            wrong.mkdir(parents=True)
            (wrong / 'model_index.json').write_text('{}')
            with self.assertRaises(FileNotFoundError):
                cached_model_path()
            expected = wrong.parent / 'pinned'
            expected.mkdir()
            (expected / 'model_index.json').write_text('{}')
            self.assertEqual(cached_model_path(), str(expected))

    def test_official_base_recipe_uses_prefix_cache_without_2511_guidance(self):
        args = sampling({'prompt': 'Preserve the photographed bag.', 'size': '1024*1536'})
        self.assertEqual((args['width'], args['height'], args['num_inference_steps']), (1024, 1536, 40))
        self.assertEqual(args['true_cfg_scale'], 1.0)
        self.assertTrue(args['use_kv_cache'])
        self.assertEqual(args['output_resolution'], 1024)
        self.assertNotIn('guidance_scale', args)
        self.assertNotIn('sigmas', args)

    def test_sampling_matches_the_pinned_official_pipeline_signature(self):
        root = Path(__file__).parents[1]
        contract = json.loads((root / 'tests/pipeline_contract.json').read_text())
        self.assertIn(contract['diffusers_commit'], (root / 'requirements.txt').read_text())
        self.assertLessEqual(set(sampling({'prompt': 'Keep the actual bag.'})),
                             set(contract['call_parameters']))

    def test_invalid_controls_fail_before_loading_a_model(self):
        with patch('qwen21.handler.pipeline') as load:
            result = handler({'input': {'task': 'edit', 'jobs': [
                {'prompt': 'Keep it.', 'image_url': 'https://source.jpg', 'guidance_scale': 4}]}})
        load.assert_not_called()
        self.assertFalse(result['results'][0]['ok'])

    def test_resolution_steps_and_empty_prompt_are_validated(self):
        for values in ({'size': '1025*1536'}, {'num_inference_steps': 200},
                       {'true_cfg_scale': 0}, {'prompt': ''}):
            with self.subTest(values=values), self.assertRaises(ValueError):
                sampling({'prompt': 'Keep it.', **values})


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