import json
import tempfile
import unittest
from pathlib import Path
from PIL import Image
from scripts.integrity import sha256
from scripts.editing_protocol import image_diagnostics
from scripts.editing_report import validate_run


class EditingReportTests(unittest.TestCase):
    def fixture(self, root):
        image = Image.new('RGBA', (32, 32), 'white')
        image.save(root/'input.png')
        image.save(root/'result.png')
        case = dict(id='test', prompt='Keep unchanged', source_seed=42,
                    input_file='input.png', input_sha256=sha256(root/'input.png'))
        (root/'cases.json').write_text(json.dumps([case]))
        env = dict(protocol='editing-v2', seeds=[1000042], steps=40, resolution=32,
                   kv_cache=True, vae_tiling=False, offload='model',
                   before_load_memory={'allocated_bytes': 0},
                   cases_sha256=sha256(root/'cases.json'))
        (root/'environment.json').write_text(json.dumps(env))
        row = dict(protocol='editing-v2', case_id='test', prompt=case['prompt'],
                   seed=1000042, source_seed=42, steps=40, output_resolution=32,
                   width=32, height=32, cfg=1., kv_cache=True, vae_tiling=False,
                   offload='model', input_file='input.png', input_sha256=case['input_sha256'],
                   image='result.png', image_sha256=sha256(root/'result.png'),
                   **image_diagnostics(image))
        self.write_rows(root, [row])
        return row

    def write_rows(self, root, rows):
        (root/'records.jsonl').write_text(''.join(json.dumps(r)+'\n' for r in rows))
        (root/'COMPLETE.json').write_text(json.dumps(dict(records=len(rows),
                                 records_sha256=sha256(root/'records.jsonl'))))

    def test_valid_run_and_rejects_duplicate_missing_or_wrong_seed(self):
        with tempfile.TemporaryDirectory() as d:
            root = Path(d); row = self.fixture(root)
            self.assertEqual(len(validate_run(root, root)[1]), 1)
            for rows in ([row, row], [], [dict(row, seed=42)]):
                self.write_rows(root, rows)
                with self.assertRaises(ValueError): validate_run(root, root)

    def test_rejects_diagnostics_that_claim_white_is_transparent(self):
        with tempfile.TemporaryDirectory() as d:
            root = Path(d); row = self.fixture(root)
            self.write_rows(root, [dict(row, alpha_transparent_fraction=1.)])
            with self.assertRaisesRegex(ValueError, 'diagnostic'): validate_run(root, root)

    def test_rejects_incomplete_run_and_mutated_input(self):
        with tempfile.TemporaryDirectory() as d:
            root = Path(d); self.fixture(root)
            (root/'COMPLETE.json').unlink()
            with self.assertRaises((ValueError, FileNotFoundError)): validate_run(root, root)
            self.fixture(root)
            Image.new('RGB', (32, 32), 'red').save(root/'input.png')
            with self.assertRaisesRegex(ValueError, 'Input'): validate_run(root, root)

    def test_rejects_png_dimensions_even_if_diagnostics_agree(self):
        with tempfile.TemporaryDirectory() as d:
            root = Path(d); row = self.fixture(root)
            image = Image.new('RGBA', (64, 32), 'white')
            image.save(root/'result.png')
            row.update(image_sha256=sha256(root/'result.png'), **image_diagnostics(image))
            self.write_rows(root, [row])
            with self.assertRaisesRegex(ValueError, 'dimensions'): validate_run(root, root)