Copies/copienator/statement_checkpoints.py

"""Extraction checkpoints assume unchanged exam sources until an explicit restart."""
from __future__ import annotations

import json
import re
from datetime import UTC, datetime
from pathlib import Path
from uuid import uuid4

from copienator import CliError, atomic_write_json, atomic_write_text
from copienator.latex_json import repair_latex_json


def archive_response(root: Path, stage: str, raw: str | None) -> Path:
    stamp = datetime.now(UTC).strftime('%Y%m%dT%H%M%S%fZ')
    path = root / 'Cache' / 'statement_responses' / f'{stamp}-{stage}-{uuid4().hex}.json'
    atomic_write_text(path, raw or '')
    return path


def report_error(root: Path, stage: str, error, raw=None, response_path=None) -> None:
    diagnostic = f'{stage}: {error}\n'
    if raw is not None:
        diagnostic += f'Raw response (escaped): {raw!r}\n'
    if response_path:
        diagnostic += f'Raw response saved to: {response_path}\n'
    print(diagnostic, end='', flush=True)
    try:
        path = root / 'Cache' / 'statement_errors.log'
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open('a', encoding='utf-8') as log:
            log.write(f'[{datetime.now(UTC).isoformat()}]\n{diagnostic}\n')
    except OSError as log_error:
        print(f'Could not save extraction error log: {log_error}', flush=True)


def validate_text(text: str) -> None:
    if not text.strip():
        raise ValueError('Empty extracted content')
    bad = [(i, ord(c)) for i, c in enumerate(text)
           if (ord(c) < 32 and c not in '\n\t') or ord(c) == 127]
    # Tabs can be legitimate indentation in code, but in math they often
    # indicate an unknown misescaped TeX command beginning with t.
    for match in re.finditer(r'(\$\$.*?\$\$|\$.*?\$|\\\(.*?\\\)|\\\[.*?\\\])', text, re.DOTALL):
        bad.extend((match.start() + i, ord(c)) for i, c in enumerate(match.group()) if c == '\t')
    if bad:
        positions = ', '.join(f'U+{code:04X} at offset {i}' for i, code in bad)
        raise ValueError(f'Control characters: {positions}; decoded content: {text!r}')


def extract_stage(root: Path, stage: str, item_model, content_field: str, request,
                  *, labels: list[str] | None = None) -> list:
    """Keep valid items; request only pending known labels on the next run.

    Questions are re-requested as a complete ordered inventory because there
    is no independent list of question labels before extraction. Contexts may
    legitimately be empty; an empty questions/solutions response may not.
    """
    checkpoint = root / 'Cache' / f'statement_{stage}.json'
    state = json.loads(checkpoint.read_text()) if checkpoint.exists() else {}
    accepted = state.get(stage, [])
    inventory = state.get('inventory', [item['label'] for item in accepted] if stage == 'questions' else [])
    if state.get('complete'):
        return [item_model.model_validate(item) for item in accepted]
    pending = state.get('pending')
    if stage == 'solutions':
        known = {item['label'] for item in accepted}
        pending = [label for label in labels if label not in known]
    legacy = root / 'Cache' / f'gemini_{"context" if stage == "contexts" else stage}.json'
    raw = None
    response_path = None
    try:
        if not checkpoint.exists() and legacy.exists():
            raw = legacy.read_text(encoding='utf-8')
            print(f'Importing legacy {stage} response from {legacy}')
        else:
            print(f'Extracting {stage}' + (f' for {pending}' if pending else '') + '...')
            raw = request(pending).text
        response_path = archive_response(root, stage, raw)
        payload = json.loads(repair_latex_json(raw))
        if not isinstance(payload, dict) or not isinstance(payload.get(stage), list):
            raise TypeError(f'Response must contain a {stage} list')
    except Exception as error:
        # Mark legacy data as attempted too, so an invalid cache is never
        # replayed indefinitely. Previously accepted items remain untouched.
        atomic_write_json(checkpoint, {stage: accepted, 'complete': False, 'pending': pending,
                                       'inventory': inventory})
        report_error(root, stage, error, raw, response_path)
        raise CliError(f'{stage} extraction failed; saved progress can be resumed') from error

    errors = []
    failed_targets = []
    items = []
    seen = set()
    for index, entry in enumerate(payload[stage]):
        if stage == 'questions' and isinstance(entry, dict):
            label = entry.get('label')
            if isinstance(label, str) and label.strip() and label not in inventory:
                inventory.append(label)
        try:
            item = item_model.model_validate(entry)
            validate_text(getattr(item, content_field))
            if stage == 'contexts':
                first, last = item.target_question_label, item.last_question_label
                if first not in labels or last not in labels or labels.index(first) > labels.index(last):
                    raise ValueError('Invalid context target or range')
                if pending and first not in pending:
                    raise ValueError('Unexpected context target')
                key = (first, last, item.context_content)
            else:
                if not item.label.strip():
                    raise ValueError('Empty question label')
                if labels is not None and item.label not in (pending or labels):
                    raise ValueError('Unexpected solution label')
                key = item.label
            if key in seen:
                raise ValueError('Duplicate extracted item')
            seen.add(key)
            items.append(item.model_dump())
        except (ValueError, TypeError) as error:
            errors.append(f'Item {index}: {error}; item (escaped): {entry!r}')
            if isinstance(entry, dict) and entry.get('target_question_label') in (labels or []):
                failed_targets.append(entry['target_question_label'])

    if stage == 'contexts':
        # Preserve good contexts even when another context for the same target fails.
        combined = accepted + [item for item in items if item not in accepted]
        next_pending = list(dict.fromkeys(failed_targets)) if len(failed_targets) == len(errors) else None
        absent = [label for label in (pending or []) if not any(
            item['target_question_label'] == label for item in items)]
        if absent:
            errors.append(f'Missing context targets: {absent}')
            next_pending = list(dict.fromkeys((next_pending or []) + absent))
    else:
        previous = {item['label']: item for item in accepted}
        previous.update({item['label']: item for item in items if item['label'] not in previous})
        order = [item['label'] for item in items] if stage == 'questions' else labels
        combined = [previous[label] for label in order if label in previous]
        combined += [item for label, item in previous.items() if label not in order]
        missing = [label for label in (labels or []) if label not in previous]
        if stage == 'questions':
            missing = [label for label in inventory if label not in seen]
            if not items:
                errors.append('No valid questions returned')
        if missing:
            errors.append(f'Missing labels: {missing}')
        next_pending = [label for label in (labels or []) if label not in previous] or None
    complete = not errors
    atomic_write_json(checkpoint, {stage: combined, 'complete': complete, 'pending': next_pending,
                                   'inventory': inventory})
    if errors:
        report_error(root, stage, '\n'.join(errors), raw, response_path)
        raise CliError(f'Incomplete {stage}; {len(combined)} valid items saved. Rerun to resume.')
    return [item_model.model_validate(item) for item in combined]