"""Check prediction identity, duplicate detection, and safe evaluation resume."""
import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

import torch

from eval.m4seg import score_predictions
from eval.tasks import parse_args
from tools.predictions import (checkpoint_digest, generation_identity,
                               index_prediction_files, prepare_prediction_directory,
                               sample_identifier)


class PredictionTests(unittest.TestCase):
    def test_conflicting_shard_outputs_are_rejected(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for shard, value in [('first', 1), ('second', 2)]:
                (root / shard).mkdir()
                torch.save({'value': value}, root / shard / 'example.pth')
            with self.assertRaisesRegex(ValueError, 'Duplicate predictions for example'):
                index_prediction_files(root)

    def test_disjoint_shard_outputs_are_indexed_together(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for sample in ['first', 'second']:
                (root / sample).mkdir()
                (root / sample / f'{sample}.pth').write_bytes(b'prediction')
            self.assertEqual(set(index_prediction_files(root)), {'first', 'second'})

    def test_scorer_uses_the_same_unique_identifier_as_generation(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            record = {'pair_id': 'source-pair', 'example_id': 'test-example',
                      'img_ids': [1], 'dataset_name': 'source',
                      'groundings': {'img_1': [{'name': 'chair',
                                               'segmentation': [[0, 0, 4, 0, 4, 4, 0, 4]]}]}}
            torch.save({'output': ['<p> chair </p> [SEG] (IMAGE1)'],
                        'pred_masks': [None, torch.ones(1, 4, 4)],
                        'image_sizes': [[4, 4]]}, root / 'test-example.pth')
            result = score_predictions([root], [record])['results'][0]
            self.assertEqual(result['samples'], 1)
            self.assertEqual(result['miou'], 1.0)

    def test_duplicate_annotation_ids_are_rejected(self):
        with self.assertRaisesRegex(ValueError, 'unique example identifiers'):
            score_predictions([], [{'pair_id': 'same'}, {'pair_id': 'same'}])

    def test_identifiers_reject_unsafe_or_missing_names(self):
        self.assertEqual(sample_identifier({'pair_id': 17}), '17')
        for value in ('', '.', '..', '../sample', '/sample', 'a\\b'):
            with self.subTest(value=value), self.assertRaises(ValueError):
                sample_identifier({'pair_id': value})
        with self.assertRaises(ValueError):
            sample_identifier({})

    def test_resume_requires_provenance_and_preserves_existing_outputs(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            existing = root / 'sample.pth'
            existing.write_bytes(b'previous output')
            with self.assertRaisesRegex(ValueError, 'no generation provenance'):
                prepare_prediction_directory(root, {'model': 'new'}, resume=True)
            self.assertEqual(existing.read_bytes(), b'previous output')
            self.assertFalse((root / 'generation.json').exists())

    def test_resume_requires_an_exact_identity_without_overwriting_it(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            identity = {'model': 'original', 'settings': {'max_length': 3400}}
            prepare_prediction_directory(root, identity, resume=False)
            provenance = root / 'generation.json'
            before = provenance.read_bytes()
            prepare_prediction_directory(root, identity, resume=True)
            self.assertEqual(provenance.read_bytes(), before)
            with self.assertRaisesRegex(ValueError, 'does not match'):
                prepare_prediction_directory(root, {'model': 'other'}, resume=True)
            self.assertEqual(provenance.read_bytes(), before)
            with self.assertRaisesRegex(ValueError, 'already has generation provenance'):
                prepare_prediction_directory(root, identity, resume=False)

    def test_malformed_provenance_and_parent_shard_directories_are_rejected(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / 'generation.json').write_text('{broken')
            with self.assertRaisesRegex(ValueError, 'Cannot verify'):
                prepare_prediction_directory(root, {}, resume=True)
            (root / 'shard').mkdir()
            (root / 'shard/sample.pth').write_bytes(b'prediction')
            with self.assertRaisesRegex(ValueError, 'separate output directory'):
                prepare_prediction_directory(root, {}, resume=True)

    def test_checkpoint_identity_includes_nested_tokenizer_assets(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / 'model.safetensors').write_bytes(b'weights')
            (root / 'qformer_tokenizer').mkdir()
            vocabulary = root / 'qformer_tokenizer/vocab.txt'
            vocabulary.write_text('original')
            before = checkpoint_digest(root)
            vocabulary.write_text('changed')
            self.assertNotEqual(checkpoint_digest(root), before)

    def test_generation_identity_tracks_settings_environment_and_assets(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / 'model').mkdir()
            config = root / 'model/config.json'
            config.write_text('{}')
            annotation = root / 'test.json'
            annotation.write_text('[]')
            args = parse_args(['--merged_ckpt_path', str(root / 'model'),
                               '--m4seg_annpath', str(annotation), '--dataset_dir', str(root)])
            with patch.dict('os.environ', {'PRIMA_TRUNCATE_FROM_LEFT': '0'}):
                before = generation_identity(args, root / 'model')
                for name, value in [('conv_type', 'other'), ('model_max_length', 3000)]:
                    previous = getattr(args, name)
                    setattr(args, name, value)
                    self.assertNotEqual(generation_identity(args, root / 'model'), before)
                    setattr(args, name, previous)
                with patch.dict('os.environ', {'PRIMA_TRUNCATE_FROM_LEFT': '1'}):
                    self.assertNotEqual(generation_identity(args, root / 'model'), before)
                config.write_text('{"changed": true}')
                self.assertNotEqual(generation_identity(args, root / 'model'), before)