import json
import tempfile
import unittest
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np

from eval_segmentqa_vllm import (
    VideoCache,
    load_fixed64_rows,
    load_fixed80_absmeta_requests,
    pending_rows,
    resolve_checkpoint,
    row_id,
    sample_frames,
)


class EvalSegmentQATest(unittest.TestCase):
    def test_single_frame_is_duplicated(self):
        frame = np.ones((2, 2, 3), dtype=np.uint8)
        sampled = sample_frames([frame], 64)
        self.assertEqual(len(sampled), 2)
        self.assertTrue(np.array_equal(sampled[0], sampled[1]))

    def test_cache_concurrent_copy(self):
        with tempfile.TemporaryDirectory() as root:
            src = Path(root) / 'dataset' / 'videos' / 'v.bin'
            src.parent.mkdir(parents=True); src.write_bytes(b'video')
            cache = VideoCache(Path(root) / 'cache', 1)
            with ThreadPoolExecutor(max_workers=4) as pool:
                paths = list(pool.map(cache.get, [src] * 4))
            self.assertEqual(len(set(paths)), 1)
            self.assertEqual(paths[0].read_bytes(), b'video')
            self.assertEqual(cache.bytes, 5)

    def test_resume_skips_ids_and_only_changes_execution_order(self):
        rows = [
            {'id': 1, '_dataset': 'd', 'video': 'b', 'timestamp_start': '00:00:02', '_input_index': 0},
            {'id': 2, '_dataset': 'd', 'video': 'a', 'timestamp_start': '00:00:03', '_input_index': 1},
            {'id': 3, '_dataset': 'd', 'video': 'a', 'timestamp_start': '00:00:01', '_input_index': 2},
        ]
        by_id, done, pending = pending_rows(rows, [{'id': 1}], 'video_time')
        self.assertEqual(done, {'1'})
        self.assertEqual([r['id'] for r in pending], [3, 2])
        self.assertEqual([by_id[str(i)]['_input_index'] for i in (1, 2, 3)], [0, 1, 2])

    def test_load_fixed64_rows_validates_counts_and_root(self):
        with tempfile.TemporaryDirectory() as root:
            dataset = Path(root) / 'test.jsonl'
            dataset.write_text(json.dumps({
                'id': 'q1', 'dataset': 'heico-focus-vqa', 'derived_video': 'videos/a.mp4',
                'derived_frame_count': 64, 'verified_frame_count': 64, 'output_fps': 2.0,
                'question': 'q', 'answer': 'a', 'answer_format': 'free-form'
            }) + '\n')
            rows = load_fixed64_rows(dataset)
            self.assertEqual(rows[0]['_dataset'], 'heico')
            self.assertEqual(rows[0]['_fixed64_root'], root)

            dataset.write_text(json.dumps({
                'id': 'bad', 'dataset': 'heico-focus-vqa', 'derived_frame_count': 63,
                'verified_frame_count': 64
            }) + '\n')
            with self.assertRaises(ValueError):
                load_fixed64_rows(dataset)

    def test_load_fixed80_requests_uses_segment_metadata_id_and_rejects_answers(self):
        with tempfile.TemporaryDirectory() as root:
            video = Path(root) / 'clip.mp4'
            video.write_bytes(b'placeholder')
            dataset = Path(root) / 'requests.jsonl'
            row = {
                'messages': [{'role': 'system', 'content': 'sys'}, {'role': 'user', 'content': '<video>\nquestion'}],
                'videos': [str(video)],
                'chat_template_kwargs': {'video_metadata': {
                    'fps': 1000, 'frames_indices': list(range(80)), 'total_num_frames': 80,
                    'width': 420, 'height': 224, 'duration': 0.08,
                    'video_backend': 'precomputed_absolute_pts'}},
                'segment_metadata': {'id': 'q1', 'dataset': 'heico-focus-vqa',
                                     'video': 'v.mp4', 'timestamp_start': '00:00:01',
                                     'answer_format': 'binary'},
            }
            dataset.write_text(json.dumps(row) + '\n')
            rows = load_fixed80_absmeta_requests(dataset, 80)
            self.assertEqual(row_id(rows[0]), 'q1')
            self.assertTrue(rows[0]['_fixed80_absmeta'])

            row['answer'] = 'yes'
            dataset.write_text(json.dumps(row) + '\n')
            with self.assertRaises(ValueError):
                load_fixed80_absmeta_requests(dataset, 80)

    def test_resolve_checkpoint_rejects_non_final_complete_gate(self):
        with tempfile.TemporaryDirectory() as root:
            ckpt = Path(root) / 'checkpoint-1'
            ckpt.mkdir()
            for name in ('adapter_config.json', 'adapter_model.safetensors'):
                (ckpt / name).write_text('{}')
            (ckpt / 'trainer_state.json').write_text(json.dumps({'global_step': 1, 'max_steps': 2}))
            with self.assertRaises(ValueError):
                resolve_checkpoint(str(ckpt), require_complete=True)
            resolved, info = resolve_checkpoint(str(ckpt), require_complete=False)
            self.assertEqual(resolved, ckpt)
            self.assertEqual(info['global_step'], 1)


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