import ast
import hashlib
from pathlib import Path
import sys
import unittest

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model_samples import AESTHETIC, NOTES, conditions, image_path, records, sample_data


class TestModelSamples(unittest.TestCase):
    def test_all_twelve_assets_have_verified_hashes(self):
        from PIL import Image
        self.assertEqual(len(records()), 10)
        self.assertEqual(len(records(True)), 2)
        for recommended in (False, True):
            for record in records(recommended):
                path = Path(image_path(record, recommended))
                self.assertEqual(hashlib.sha256(path.read_bytes()).hexdigest(), record['sha256'])
                with Image.open(path) as image:
                    self.assertEqual(image.size, (1024, 1024))
                    image.verify()
                self.assertFalse(record['loras'])
                self.assertFalse(record['embeddings'])

    def test_recommended_changes_only_positive_quality_prefix(self):
        common = [r for r in records() if r['model'] == AESTHETIC]
        for old, new in zip(common, records(True)):
            self.assertEqual(new['prompt'], old['prompt'].removeprefix('masterpiece, best quality, '))
            for key in ('kind', 'seed', 'width', 'height', 'steps', 'cfg', 'sampler',
                        'scheduler', 'negative_prompt', 'loras', 'embeddings'):
                self.assertEqual(old[key], new[key])
            self.assertNotIn('score_', new['prompt'] + new['negative_prompt'])

    def test_selection_and_legacy_alias(self):
        for model in NOTES:
            note, common, detail, extra, rec_detail, visible = sample_data(model)
            self.assertTrue(all(common))
            self.assertIn(model, detail)
            self.assertEqual(visible, model == AESTHETIC)
            self.assertEqual(all(extra), visible)
        self.assertEqual(sample_data('anima_aestheticV11_INT8.safetensors'), sample_data(AESTHETIC))
        self.assertEqual(sample_data('unknown')[1], [None, None])
        self.assertFalse(sample_data('unknown')[-1])

    def test_no_path_escape(self):
        with self.assertRaises(ValueError):
            image_path({'file': '../../../../../app.py'})

    def test_builder_read_only_and_at_correct_position(self):
        tree = ast.parse((ROOT / 'app.py').read_text(encoding='utf-8'))
        builder = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == 'build_generation_page')
        calls = [n for n in ast.walk(builder) if isinstance(n, ast.Call)]
        sample = next(n for n in calls if isinstance(n.func, ast.Name) and n.func.id == 'build_model_samples')
        prompt = next(n for n in calls if isinstance(n.func, ast.Name) and n.func.id == 'prompt_with_completion')
        self.assertLess(sample.lineno, prompt.lineno)
        source = (ROOT / 'model_samples.py').read_text(encoding='utf-8')
        self.assertIn('open=False', source)
        self.assertIn('api_visibility="private"', source)
        self.assertNotIn('generate_image(', source)


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