import asyncio
import csv
import json
from pathlib import Path
import tempfile
import unittest

from assembly_dataset.config import Config
from assembly_dataset.pipeline import generate, export, validate


class CsvExportTests(unittest.TestCase):
    def test_each_shard_has_readable_csv_with_same_records(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=25, shard_size=2)
            config.prompts.teacher_system = 'Правила'
            config.prompts.training_system['full'] = 'Текст; с "кавычками"\nи новой строкой'
            asyncio.run(generate(config))
            export(config)
            for path in Path(tmp).glob('sft/*/*/*.jsonl'):
                counterpart = path.with_suffix('.csv')
                self.assertTrue(counterpart.exists())
                self.assertTrue(counterpart.read_bytes().startswith(b'\xef\xbb\xbf'))
                originals = [json.loads(line) for line in path.read_text().splitlines()]
                with counterpart.open(encoding='utf-8-sig', newline='') as file:
                    rows = list(csv.DictReader(file, delimiter=';'))
                self.assertEqual(len(rows), len(originals))
                for original, row in zip(originals, rows):
                    self.assertEqual(row['ID записи'], original['record_id'])
                    self.assertEqual(row['Системная инструкция'], original['messages'][0]['content'])
                    self.assertEqual(row['План'], original['messages'][2]['content'])
                    self.assertNotRegex(row['План'], r'part_\d|slot_\d')
                    if original['output_format'] == 'text':
                        self.assertNotIn('"skill"', row['План'])
                    self.assertTrue(row['План'])
            self.assertEqual(validate(config)['examples_checked'], 25)

    def test_corrupt_csv_is_detected(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=1)
            asyncio.run(generate(config))
            export(config)
            files = list(Path(tmp).glob('sft/*/*/*.csv'))
            self.assertEqual(len(files), 2)
            files[0].write_text('broken')
            with self.assertRaisesRegex(ValueError, 'Повреждён'):
                validate(config)