"""Bounded generation, transactional checkpoints, streaming sharded export."""
import asyncio
import csv
from collections import Counter
from contextlib import contextmanager
import fcntl
import hashlib
import json
from itertools import zip_longest
from pathlib import Path
import shutil
import sqlite3
import tempfile

from .config import Config
from .language import validate_texts
from .csv_export import COLUMNS, csv_record
from .progress import ProgressDisplay
from .world import canonical, digest, make_episode, training_rows, validate_plan, validate_branch

VERSION = 3
STORAGE_VERSION = 2


@contextmanager
def storage(config):
    root = config.output_dir
    root.mkdir(parents=True, exist_ok=True)
    with (root / '.lock').open('a') as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise ValueError('Другой процесс уже работает с output_dir') from None
        db = sqlite3.connect(root / 'episodes.sqlite3')
        try:
            db.execute('PRAGMA journal_mode=WAL')
            db.execute('CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)')
            db.execute('CREATE TABLE IF NOT EXISTS episodes (seq INTEGER PRIMARY KEY, signature TEXT UNIQUE NOT NULL, payload TEXT NOT NULL, rows INTEGER NOT NULL)')
            settings = config.model_dump(mode='json', exclude={'target_examples', 'output_dir', 'shard_size', 'prompts_file'})
            # Structured prompts only affect export, so existing v2 episodes remain reusable.
            settings['prompts'].pop('structured_system', None)
            fingerprint = digest({'version': STORAGE_VERSION, 'settings': settings})
            saved = db.execute("SELECT value FROM meta WHERE key='fingerprint'").fetchone()
            if saved and saved[0] != fingerprint:
                raise ValueError('Изменён config: для других настроек используйте новый output_dir')
            if not saved:
                with db:
                    db.execute('INSERT INTO meta VALUES (?, ?)', ('fingerprint', fingerprint))
                    db.execute('INSERT INTO meta VALUES (?, ?)', ('settings', canonical(settings)))
                    db.execute('INSERT INTO meta VALUES (?, ?)', ('next_index', '0'))
            yield db
        finally:
            db.close()


async def generate(config: Config, agent=None, *, show_progress=False, teacher_only=False):
    with storage(config) as db:
        count = db.execute('SELECT COALESCE(SUM(rows), 0) FROM episodes').fetchone()[0]
        index = int(db.execute("SELECT value FROM meta WHERE key='next_index'").fetchone()[0])
        resume_until = db.execute('SELECT COALESCE(MAX(seq), 0) FROM episodes').fetchone()[0]
        with ProgressDisplay('Генерация', config.target_examples, initial=count, enabled=show_progress) as progress:
            if count >= config.target_examples and index >= resume_until:
                return count
            if config.teacher.enabled and config.teacher.fraction > 0 and agent is None:
                from .teacher import create_agent
                progress.update(status='подключение LLM')
                agent = await create_agent(config.teacher, config.prompts)
            batch_size = config.teacher.concurrency if agent is not None else 32
            duplicates = 0
            while count < config.target_examples or index < resume_until:
                # Never construct 500k coroutines or keep the entire dataset in RAM.
                size = min(batch_size, max(1, config.target_examples - count))
                episodes = [make_episode(i, config) for i in range(index, index + size)]
                async def save_episode(episode):
                    nonlocal count, duplicates
                    # A restarted batch can contain responses already committed before cancellation.
                    if db.execute('SELECT 1 FROM episodes WHERE signature=?',
                                  (episode['signature'],)).fetchone():
                        duplicates += 1
                        return
                    if agent is not None:
                        from .teacher import enrich
                        episode, _ = await enrich(episode, config.teacher, agent, config.prompts)
                    if teacher_only and (config.teacher.llm is None or
                                         episode['source']['provider'] != config.teacher.llm.provider or
                                         'teacher_error' in episode):
                        raise ValueError('Режим --teacher-only: ответ модели не прошёл проверку; символическая замена не сохранена')
                    rows = sum(1 for _ in training_rows(episode, config))
                    # No await between receiving a validated response and committing it.
                    with db:
                        cursor = db.execute('INSERT OR IGNORE INTO episodes(seq,signature,payload,rows) VALUES (?,?,?,?)',
                                            (episode['candidate_index'] + 1, episode['signature'], canonical(episode), rows))
                    if cursor.rowcount:
                        count += rows
                        duplicates = 0
                    else:
                        duplicates += 1
                    progress.update(count, status='ожидание LLM' if agent is not None else '')

                if agent is not None:
                    progress.update(status='ожидание LLM')
                tasks = [asyncio.create_task(save_episode(episode)) for episode in episodes]
                try:
                    await asyncio.gather(*tasks)
                finally:
                    for task in tasks:
                        if not task.done():
                            task.cancel()
                    await asyncio.gather(*tasks, return_exceptions=True)
                # Advance only after the whole batch succeeds; holes are retried on resume.
                index += size
                with db:
                    db.execute("UPDATE meta SET value=? WHERE key='next_index'", (str(index),))
                if duplicates > 10000:
                    raise ValueError('Пространство сценариев исчерпано: расширьте диапазон компонентов или domains')
                progress.update(count, status='')
            return count


class Shards:
    def __init__(self, root, size):
        self.root, self.size = root, size
        self.counts = Counter()
        self.handles = {}
        self.csv_handles = {}
        self.csv_writers = {}
        self.files = []

    def write(self, row, structured_row):
        key = row['split'], row['mode']
        count = self.counts[key]
        for variant, record in (('text', row), ('structured', structured_row)):
            handle_key = (*key, variant)
            if count % self.size == 0:
                if handle_key in self.handles:
                    self.handles[handle_key].close()
                    self.csv_handles[handle_key].close()
                suffix = '.jsonl' if variant == 'text' else '.structured.jsonl'
                path = self.root / key[0] / key[1] / f'{count // self.size:06d}{suffix}'
                path.parent.mkdir(parents=True, exist_ok=True)
                self.handles[handle_key] = path.open('w', encoding='utf-8')
                self.files.append(path)
                csv_path = path.with_suffix('.csv')
                self.csv_handles[handle_key] = csv_path.open('w', encoding='utf-8-sig', newline='')
                self.csv_writers[handle_key] = csv.DictWriter(self.csv_handles[handle_key], fieldnames=COLUMNS, delimiter=';')
                self.csv_writers[handle_key].writeheader()
                self.files.append(csv_path)
            self.handles[handle_key].write(canonical(record) + '\n')
            self.csv_writers[handle_key].writerow(csv_record(record))
        self.counts[key] += 1

    def close(self):
        for file in (*self.handles.values(), *self.csv_handles.values()):
            file.close()


def file_digest(path):
    hasher = hashlib.sha256()
    with path.open('rb') as file:
        for block in iter(lambda: file.read(1024 * 1024), b''):
            hasher.update(block)
    return hasher.hexdigest()


def export(config, *, show_progress=False, allow_partial=False):
    with storage(config) as db, ProgressDisplay('Экспорт', config.target_examples, enabled=show_progress) as progress:
        capacity = db.execute('SELECT COALESCE(SUM(rows), 0) FROM episodes').fetchone()[0]
        if capacity < config.target_examples and not allow_partial:
            raise ValueError('Недостаточно сценариев: сначала выполните generate')
        destination = config.output_dir / 'sft'
        backup = config.output_dir / '.sft-previous'
        # Recover an interrupted directory swap before making another export.
        if backup.exists() and not destination.exists():
            backup.rename(destination)
        if destination.exists() and not (destination / 'manifest.json').exists():
            raise ValueError('sft не является управляемым экспортом; выберите другой output_dir')
        stage = Path(tempfile.mkdtemp(prefix='.sft-', dir=config.output_dir))
        writer = Shards(stage, config.shard_size)
        total, scenarios, failures = 0, 0, 0
        domains, sources = Counter(), Counter()
        try:
            for (payload,) in db.execute('SELECT payload FROM episodes ORDER BY seq'):
                if total >= config.target_examples:
                    break
                episode = json.loads(payload)
                scenarios += 1
                domains[episode['scene']['domain']] += 1
                sources[episode['source']['provider']] += 1
                failures += int('teacher_error' in episode)
                for record, structured_record in zip(training_rows(episode, config),
                                                     training_rows(episode, config, output_format='structured'), strict=True):
                    if total >= config.target_examples:
                        break
                    writer.write(record, structured_record)
                    total += 1
                progress.update(total)
            progress.update(status="контрольные суммы")
            writer.close()
            manifest = {'schema_version': VERSION, 'examples': total, 'scenarios': scenarios,
                        'partial': total < config.target_examples,
                        'examples_per_variant': {'text': total, 'structured': total},
                        'counts': {f'{split}/{mode}': n for (split, mode), n in sorted(writer.counts.items())},
                        'domains': dict(domains), 'sources': dict(sources), 'teacher_failures': failures,
                        'validation_level': 'logical_only',
                        'config': config.model_dump(mode='json'),
                        'files': {str(p.relative_to(stage)): file_digest(p) for p in writer.files}}
            (stage / 'manifest.json').write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')
            if backup.exists():
                shutil.rmtree(backup)
            if destination.exists():
                destination.rename(backup)
            stage.rename(destination)
            if backup.exists():
                shutil.rmtree(backup)
            progress.update(total, status='')
            return manifest
        finally:
            writer.close()
            if stage.exists():
                shutil.rmtree(stage)


def validate(config):
    with storage(config) as db:
        episodes, paths = 0, 0
        for signature, payload, stored_rows in db.execute('SELECT signature,payload,rows FROM episodes'):
            episode = json.loads(payload)
            if signature != digest(episode['scene']) or episode['signature'] != signature or episode['scenario_id'] != signature:
                raise ValueError('Неверная сигнатура сценария')
            if 'step_texts' in episode:
                validate_texts(episode, episode['step_texts'], config.prompts)
            trace = validate_plan(episode, episode['plan'])
            if trace != episode['trace']:
                raise ValueError('Trace не соответствует плану')
            paths += validate_branch(episode, episode['branch'])
            if sum(1 for _ in training_rows(episode, config)) != stored_rows:
                raise ValueError('Неверное число записей эпизода')
            episodes += 1
        root = config.output_dir / 'sft'
        manifest_path = root / 'manifest.json'
        if not manifest_path.exists():
            raise ValueError('Нет экспорта SFT: выполните generate или export')
        manifest = json.loads(manifest_path.read_text())
        total = 0
        expected_files = set(manifest['files'])
        actual_files = {str(p.relative_to(root)) for pattern in ('*/*/*.jsonl', '*/*/*.csv') for p in root.glob(pattern)}
        if expected_files != actual_files:
            raise ValueError('Набор шардов не соответствует manifest')
        for relative, checksum in manifest['files'].items():
            path = root / relative
            if path.resolve().is_relative_to(root.resolve()) is False or file_digest(path) != checksum:
                raise ValueError(f'Повреждён шард {relative}')
        for relative in manifest['files']:
            path = root / relative
            if path.suffix == '.jsonl' and '.structured.' not in path.name:
                counterpart = path.with_suffix('.structured.jsonl')
                if str(counterpart.relative_to(root)) not in expected_files:
                    raise ValueError('Отсутствует структурированная пара JSONL')
                with path.open(encoding='utf-8') as file, counterpart.open(encoding='utf-8') as other:
                    for text_line, structured_line in zip_longest(file, other):
                        if text_line is None or structured_line is None:
                            raise ValueError('Разное число записей в паре JSONL')
                        text_row, structured_row = json.loads(text_line), json.loads(structured_line)
                        if (text_row['record_id'] != structured_row['record_id'] or
                                text_row['messages'][1] != structured_row['messages'][1]):
                            raise ValueError('Записи в паре JSONL не совпадают')
                        total += 1
        if total != manifest['examples']:
            raise ValueError('Неверное число экспортированных записей')
        return {'episodes_checked': episodes, 'branch_paths_checked': paths, 'examples_checked': total}