import json
import tempfile
import unittest
from pathlib import Path

from PIL import Image

from scripts.editing_protocol import DEFAULT_SEEDS, image_diagnostics
from scripts.editing_report import validate_suite
from scripts.integrity import sha256


class EditingReleaseGateTests(unittest.TestCase):
    def fixture(self, directory, case_count=5, seeds=DEFAULT_SEEDS):
        root = Path(directory)
        suite = root/'evaluation'/'editing-v2'
        inputs = suite/'inputs'
        inputs.mkdir(parents=True)
        source_root = root/'scripts'
        source_root.mkdir()
        hashes = {}
        for name in ('benchmark_edits_v2.py', 'editing_protocol.py', 'benchmark.py', 'runtime.py'):
            (source_root/name).write_text('# CPU release validation fixture\n')
            hashes[name] = sha256(source_root/name)

        definitions = Path(__file__).resolve().parents[1]/'benchmarks'/'editing-v2.json'
        cases = json.loads(definitions.read_text(encoding='utf-8'))
        for case in cases:
            Image.new('RGBA', (32, 32), 'white').save(inputs/case['input_file'])
            case['input_sha256'] = sha256(inputs/case['input_file'])
        (root/'benchmarks').mkdir()
        (root/'benchmarks'/'editing-v2.json').write_text(json.dumps(cases), encoding='utf-8')
        selected = cases[:case_count]
        image = Image.new('RGBA', (1024, 1024), 'white')
        diagnostics = image_diagnostics(image)
        isolation = []
        for index, precision in enumerate(('bf16', 'int8')):
            run = suite/precision
            run.mkdir()
            image.save(run/'result.png')
            (run/'cases.json').write_text(json.dumps(selected), encoding='utf-8')
            env = dict(
                protocol='editing-v2', seeds=list(seeds), steps=40, resolution=1024,
                kv_cache=True, vae_tiling=False, offload='model', offload_aux_fix=True,
                warmup=True, generator_device='cpu', pid=index+100,
                before_load_memory={'allocated_bytes': 0},
                cases_sha256=sha256(run/'cases.json'),
                model_identity={'kind': precision, 'base_revision': 'fixture',
                                'fingerprint': precision},
                gpu='fixture', cuda='fixture', packages={}, source_sha256=hashes,
                upstream_source_sha256={}, scheduler_config={}, vae_dtype='torch.bfloat16',
                vae_tiles={})
            (run/'environment.json').write_text(json.dumps(env), encoding='utf-8')
            rows = []
            for case in selected:
                for seed in seeds:
                    rows.append(dict(
                        protocol='editing-v2', case_id=case['id'], prompt=case['prompt'],
                        source_seed=case['source_seed'], seed=seed,
                        input_file=case['input_file'], input_sha256=case['input_sha256'],
                        width=1024, height=1024, output_resolution=1024, steps=40, cfg=1.,
                        kv_cache=True, vae_tiling=False, offload='model',
                        image='result.png', image_sha256=sha256(run/'result.png'),
                        sigmas=[1., 0.], **diagnostics))
            (run/'records.jsonl').write_text(
                ''.join(json.dumps(row)+'\n' for row in rows), encoding='utf-8')
            (run/'COMPLETE.json').write_text(json.dumps(dict(
                records=len(rows), records_sha256=sha256(run/'records.jsonl'))))
            isolation.append(dict(precision=precision, pid=env['pid'], exit_code=0,
                                  started=index*2, ended=index*2+1))
        (suite/'process-isolation.jsonl').write_text(
            ''.join(json.dumps(row)+'\n' for row in isolation), encoding='utf-8')
        return suite, source_root

    def test_accepts_complete_release_suite(self):
        with tempfile.TemporaryDirectory() as directory:
            suite, source_root = self.fixture(directory)
            runs = validate_suite(suite, source_root=source_root)
            self.assertEqual(len(runs['bf16'][1]), 10)
            self.assertEqual(len(runs['int8'][1]), 10)

    def test_windows_launcher_pid_can_differ_from_recorded_inference_pid(self):
        with tempfile.TemporaryDirectory() as directory:
            suite, source_root = self.fixture(directory)
            path = suite/'process-isolation.jsonl'
            rows = [json.loads(line) for line in path.read_text().splitlines()]
            for row in rows:
                row['inference_pid'] = row['pid']
                row['pid'] += 1000
            path.write_text(''.join(json.dumps(row)+'\n' for row in rows))
            self.assertEqual(len(validate_suite(suite, source_root)['int8'][1]), 10)
            rows[0]['inference_pid'] = 99999
            path.write_text(''.join(json.dumps(row)+'\n' for row in rows))
            with self.assertRaises(ValueError): validate_suite(suite, source_root)

    def test_rejects_partial_case_matrix_for_release(self):
        with tempfile.TemporaryDirectory() as directory:
            suite, source_root = self.fixture(directory, case_count=1)
            with self.assertRaises(ValueError):
                validate_suite(suite, source_root=source_root)

    def test_rejects_single_seed_matrix_for_release(self):
        with tempfile.TemporaryDirectory() as directory:
            suite, source_root = self.fixture(directory, seeds=DEFAULT_SEEDS[:1])
            with self.assertRaises(ValueError):
                validate_suite(suite, source_root=source_root)

    def test_partial_suite_remains_usable_for_diagnostics(self):
        with tempfile.TemporaryDirectory() as directory:
            suite, _ = self.fixture(directory, case_count=1, seeds=DEFAULT_SEEDS[:1])
            runs = validate_suite(suite)
            self.assertEqual(len(runs['bf16'][1]), 1)

    def test_rejects_empty_measured_source_inventory_for_release(self):
        with tempfile.TemporaryDirectory() as directory:
            suite, source_root = self.fixture(directory)
            for precision in ('bf16', 'int8'):
                path = suite/precision/'environment.json'
                env = json.loads(path.read_text(encoding='utf-8'))
                env['source_sha256'] = {}
                path.write_text(json.dumps(env), encoding='utf-8')
            with self.assertRaises(ValueError):
                validate_suite(suite, source_root=source_root)


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