from __future__ import annotations

import json
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .alignment import align_page_blocks
from .backends import CodexExecBackend, TranscriptionBackend, TranscriptionRequest
from .builder import build_candidate_page
from .cache import RunCache
from .config import EvaluationProfile
from .context import build_neighbor_context, build_prompt
from .diagnostics import detect_cross_page_duplicates
from .reconciliation import reconcile_adjacent_pages
from .schema_tools import export_model_schema, schema_errors
from .source import SourcePackage
from .util import (
    load_json,
    project_root,
    sha256_file,
    sha256_json,
    utc_now,
    write_json,
    write_text,
)
from .validation import validate_page, validate_reciprocal_pages


@dataclass(frozen=True)
class PageRunResult:
    page: int
    status: str
    candidate_path: Path | None
    errors: list[str]
    cache_hit: bool
    repair_used: bool = False
    attempt_count: int = 0
    elapsed_seconds: float = 0.0
    resumed: bool = False
    infrastructure_retries: int = 0
    infrastructure_events: tuple[dict[str, Any], ...] = ()

    def as_dict(self) -> dict[str, Any]:
        return {
            "page": self.page,
            "status": self.status,
            "candidate_path": str(self.candidate_path) if self.candidate_path else None,
            "errors": self.errors,
            "cache_hit": self.cache_hit,
            "repair_used": self.repair_used,
            "attempt_count": self.attempt_count,
            "elapsed_seconds": round(self.elapsed_seconds, 3),
            "resumed": self.resumed,
            "infrastructure_retries": self.infrastructure_retries,
            "infrastructure_events": list(self.infrastructure_events),
        }


def codex_version() -> str:
    completed = subprocess.run(
        ["codex", "--version"], text=True, capture_output=True, check=False
    )
    return completed.stdout.strip() or completed.stderr.strip() or "unknown"


def _cache_key(
    *,
    source: SourcePackage,
    page: int,
    prompt: str,
    schema_path: Path,
    profile: EvaluationProfile,
    backend_name: str,
) -> str:
    return sha256_json(
        {
            "page": page,
            "image": sha256_file(source.image_path(page)),
            "ocr": sha256_file(source.ocr_text_path(page)),
            "tsv": sha256_file(source.ocr_tsv_path(page)),
            "prompt": prompt,
            "schema": sha256_file(schema_path),
            "backend": backend_name,
            "model": profile.model,
            "reasoning_effort": profile.reasoning_effort,
            "codex_version": codex_version() if backend_name == "codex-exec" else None,
        }
    )


def _parse_response(value: str | None) -> tuple[dict[str, Any] | None, list[str]]:
    if value is None:
        return None, ["Backend produced no final message"]
    try:
        parsed = json.loads(value)
    except json.JSONDecodeError as exc:
        return None, [f"Final message is not JSON: {exc}"]
    if not isinstance(parsed, dict):
        return None, ["Final JSON response must be an object"]
    return parsed, []


def _repair_prompt(base_prompt: str, prior: str | None, errors: list[str]) -> str:
    prior_text = prior or "[no parseable prior response]"
    return f"""{base_prompt.rstrip()}

## Explicit repair attempt

The prior response failed deterministic validation. Return a complete replacement
object, not a patch. Preserve the source reading and fix only the structural or
schema problems listed below.

<validation_errors>
{chr(10).join(f'- {error}' for error in errors)}
</validation_errors>

<prior_response status="invalid">
{prior_text}
</prior_response>
"""


def _infrastructure_category(result: Any) -> str:
    detail = " ".join(
        value
        for value in (result.error, result.stderr, result.stdout)
        if isinstance(value, str)
    ).casefold()
    if "timed out" in detail or "timeout" in detail:
        return "timeout"
    if any(value in detail for value in ("429", "rate limit", "throttl", "too many requests")):
        return "throttling"
    if any(value in detail for value in ("authentication", "unauthorized", "login", "credential")):
        return "authentication"
    if any(value in detail for value in ("server error", "service unavailable", "502", "503", "504")):
        return "service"
    return "process"


def _write_backend_artifacts(
    attempt_root: Path, prompt: str, result: Any, elapsed_seconds: float
) -> None:
    write_text(attempt_root / "prompt.md", prompt)
    write_json(attempt_root / "command.json", {"command": result.command})
    write_text(attempt_root / "stdout.jsonl", result.stdout)
    write_text(attempt_root / "stderr.log", result.stderr)
    write_text(attempt_root / "raw-final-response.txt", result.final_message or "")
    write_json(
        attempt_root / "backend-result.json",
        {
            "success": result.success,
            "exit_code": result.exit_code,
            "metadata": result.metadata,
            "error": result.error,
            "elapsed_seconds": round(elapsed_seconds, 3),
        },
    )


def _invoke_with_infrastructure_retries(
    *,
    backend: TranscriptionBackend,
    request: TranscriptionRequest,
    prompt: str,
    attempt_root: Path,
    profile: EvaluationProfile,
) -> tuple[Any, Path, list[dict[str, Any]]]:
    events: list[dict[str, Any]] = []
    result: Any = None
    used_root = attempt_root
    total_tries = 1 + profile.max_infrastructure_retries
    for infrastructure_try in range(1, total_tries + 1):
        used_root = (
            attempt_root
            if infrastructure_try == 1
            else attempt_root.with_name(
                f"{attempt_root.name}-infra-{infrastructure_try:02d}"
            )
        )
        invocation_started = time.perf_counter()
        result = backend.transcribe(request)
        invocation_elapsed = time.perf_counter() - invocation_started
        _write_backend_artifacts(used_root, prompt, result, invocation_elapsed)
        if result.success:
            break
        category = _infrastructure_category(result)
        retrying = infrastructure_try < total_tries
        delay = (
            profile.infrastructure_backoff_seconds * (2 ** (infrastructure_try - 1))
            if retrying
            else 0.0
        )
        event = {
            "try": infrastructure_try,
            "category": category,
            "error": result.error or "Backend infrastructure failure",
            "exit_code": result.exit_code,
            "retrying": retrying,
            "backoff_seconds": delay,
            "artifact": used_root.name,
        }
        events.append(event)
        write_json(used_root / "infrastructure-failure.json", event)
        if retrying and delay:
            time.sleep(delay)
    return result, used_root, events


def _infrastructure_retry_count(events: list[dict[str, Any]]) -> int:
    return sum(bool(event.get("retrying")) for event in events)


def transcribe_page(
    *,
    source: SourcePackage,
    page: int,
    run_id: str,
    run_root: Path,
    profile: EvaluationProfile,
    backend: TranscriptionBackend | None = None,
    force: bool = False,
) -> PageRunResult:
    started = time.perf_counter()
    source.require_page(page)
    backend = backend or CodexExecBackend()
    run_root.mkdir(parents=True, exist_ok=True)
    schema_path = run_root / "input-schema.json"
    model_schema = load_json(schema_path) if schema_path.is_file() else export_model_schema(schema_path)
    specification_snapshot = run_root / "input-transcription-specification.md"
    if specification_snapshot.is_file():
        specification = specification_snapshot.read_text(encoding="utf-8")
    else:
        specification_path = project_root() / "prompts" / "transcribe-page.md"
        specification = specification_path.read_text(encoding="utf-8")
        write_text(specification_snapshot, specification)
    candidate_search = [
        run_root / "candidates",
        project_root() / "edition" / "owner-approved",
        project_root() / "edition" / "pages",
    ]
    context = build_neighbor_context(
        source=source,
        page=page,
        previous_characters=profile.previous_context_characters,
        next_characters=profile.next_context_characters,
        canonical_search_dirs=candidate_search,
    )
    base_prompt = build_prompt(
        specification=specification,
        page_id=source.page_id(page),
        current_ocr=source.ocr_text(page),
        context=context,
    )
    page_root = run_root / "pages" / source.page_id(page)
    page_root.mkdir(parents=True, exist_ok=True)
    write_text(page_root / "current-ocr.txt", source.ocr_text(page))
    write_json(page_root / "neighbor-context.json", context.as_dict())
    write_json(
        page_root / "input-evidence.json",
        {
            "page": page,
            "source": source.source_evidence(page),
            "context_pages": context.pages(),
        },
    )
    cache = RunCache(project_root() / "artifacts" / "cache")
    key = _cache_key(
        source=source,
        page=page,
        prompt=base_prompt,
        schema_path=schema_path,
        profile=profile,
        backend_name=backend.name,
    )
    cached = None if force else cache.load(key)
    draft: dict[str, Any] | None = None
    used_prompt_path = page_root / "attempt-01" / "prompt.md"
    cache_hit = False
    accumulated_errors: list[str] = []
    last_response: str | None = None
    attempts_used = 0
    total_attempts = 1 + profile.max_repair_attempts
    repaired = False
    infrastructure_events: list[dict[str, Any]] = []

    if cached:
        draft = cached.get("draft")
        cache_hit = draft is not None
        attempts_used = 1 if draft is not None else 0
        last_response = json.dumps(draft, ensure_ascii=False) if draft else None
        attempt_root = page_root / "attempt-01"
        write_text(attempt_root / "prompt.md", base_prompt)
        write_json(attempt_root / "cache.json", {"key": key, "hit": True})
        write_json(attempt_root / "model-draft.json", draft)
    else:
        prompt = base_prompt
        for attempt in range(1, total_attempts + 1):
            attempts_used = attempt
            attempt_root = page_root / f"attempt-{attempt:02d}"
            request = TranscriptionRequest(
                page_id=source.page_id(page),
                image_path=source.image_path(page),
                prompt=prompt,
                output_schema_path=schema_path,
                model=profile.model,
                reasoning_effort=profile.reasoning_effort,
                timeout_seconds=profile.timeout_seconds,
            )
            result, used_attempt_root, events = _invoke_with_infrastructure_retries(
                backend=backend,
                request=request,
                prompt=prompt,
                attempt_root=attempt_root,
                profile=profile,
            )
            infrastructure_events.extend(events)
            used_prompt_path = used_attempt_root / "prompt.md"
            last_response = result.final_message
            if not result.success:
                accumulated_errors.append(
                    f"attempt {attempt}: infrastructure retries exhausted: "
                    f"{result.error or 'Backend failure'}"
                )
                break
            parsed, parse_errors = _parse_response(result.final_message)
            attempt_errors: list[str] = []
            attempt_errors.extend(parse_errors)
            if parsed is not None:
                attempt_errors.extend(schema_errors(parsed, model_schema))
            write_json(
                used_attempt_root / "model-validation.json",
                {"valid": not attempt_errors, "errors": attempt_errors},
            )
            if not attempt_errors and parsed is not None:
                draft = parsed
                write_json(used_attempt_root / "model-draft.json", draft)
                break
            accumulated_errors.extend(
                f"attempt {attempt}: {error}" for error in attempt_errors
            )
            prompt = _repair_prompt(base_prompt, result.final_message, attempt_errors)

    if draft is None:
        elapsed = time.perf_counter() - started
        write_json(
            page_root / "status.json",
            {
                "page": page,
                "status": "failed",
                "errors": accumulated_errors,
                "cache_hit": cache_hit,
                "repair_used": attempts_used > 1,
                "attempt_count": attempts_used,
                "infrastructure_retries": _infrastructure_retry_count(
                    infrastructure_events
                ),
                "infrastructure_events": infrastructure_events,
                "elapsed_seconds": round(elapsed, 3),
            },
        )
        return PageRunResult(
            page,
            "failed",
            None,
            accumulated_errors,
            cache_hit,
            attempts_used > 1,
            attempts_used,
            elapsed,
            False,
            _infrastructure_retry_count(infrastructure_events),
            tuple(infrastructure_events),
        )

    candidate = build_candidate_page(
        draft=draft,
        source=source,
        page=page,
        run_id=run_id,
        backend=backend.name,
        model=profile.model,
        reasoning_effort=profile.reasoning_effort,
        prompt_path=used_prompt_path,
        schema_path=schema_path,
        context_pages=context.pages(),
    )
    alignment_warning: str | None = None
    try:
        align_page_blocks(candidate, source)
    except Exception as exc:  # Alignment is explicitly non-blocking for the PoC.
        alignment_warning = f"TSV alignment skipped: {exc}"
    report = validate_page(candidate, source)
    validation = report.as_dict()
    if alignment_warning:
        validation["warnings"].append(alignment_warning)
    write_json(page_root / "canonical-validation.json", validation)
    if not report.valid and attempts_used < total_attempts:
        # Provider schemas cannot express every cross-field editorial invariant.
        # Spend the configured repair attempt only after preserving the original
        # response and its canonical validation report.
        canonical_errors = [f"canonical: {error}" for error in report.errors]
        repair_attempt = attempts_used + 1
        attempts_used = repair_attempt
        repair_root = page_root / f"attempt-{repair_attempt:02d}"
        repair_prompt = _repair_prompt(base_prompt, last_response, canonical_errors)
        request = TranscriptionRequest(
            page_id=source.page_id(page),
            image_path=source.image_path(page),
            prompt=repair_prompt,
            output_schema_path=schema_path,
            model=profile.model,
            reasoning_effort=profile.reasoning_effort,
            timeout_seconds=profile.timeout_seconds,
        )
        result, used_repair_root, events = _invoke_with_infrastructure_retries(
            backend=backend,
            request=request,
            prompt=repair_prompt,
            attempt_root=repair_root,
            profile=profile,
        )
        infrastructure_events.extend(events)
        used_prompt_path = used_repair_root / "prompt.md"
        parsed, parse_errors = _parse_response(result.final_message) if result.success else (None, [])
        repair_errors = (
            []
            if result.success
            else [
                "infrastructure retries exhausted: "
                f"{result.error or 'Backend failure'}"
            ]
        )
        repair_errors.extend(parse_errors)
        if parsed is not None:
            repair_errors.extend(schema_errors(parsed, model_schema))
        write_json(
            used_repair_root / "model-validation.json",
            {"valid": not repair_errors, "errors": repair_errors},
        )
        if not repair_errors and parsed is not None:
            repaired = True
            draft = parsed
            write_json(used_repair_root / "model-draft.json", draft)
            candidate = build_candidate_page(
                draft=draft,
                source=source,
                page=page,
                run_id=run_id,
                backend=backend.name,
                model=profile.model,
                reasoning_effort=profile.reasoning_effort,
                prompt_path=used_prompt_path,
                schema_path=schema_path,
                context_pages=context.pages(),
            )
            alignment_warning = None
            try:
                align_page_blocks(candidate, source)
            except Exception as exc:
                alignment_warning = f"TSV alignment skipped: {exc}"
            report = validate_page(candidate, source)
            validation = report.as_dict()
            if alignment_warning:
                validation["warnings"].append(alignment_warning)
            write_json(used_repair_root / "canonical-validation.json", validation)
        else:
            accumulated_errors.extend(
                f"attempt {repair_attempt}: {error}" for error in repair_errors
            )

    if not report.valid:
        errors = [f"canonical: {error}" for error in report.errors]
        errors = [*accumulated_errors, *errors]
        elapsed = time.perf_counter() - started
        write_json(
            page_root / "status.json",
            {
                "page": page,
                "status": "failed",
                "errors": errors,
                "cache_hit": cache_hit,
                "repair_used": attempts_used > 1 or repaired,
                "attempt_count": attempts_used,
                "infrastructure_retries": _infrastructure_retry_count(
                    infrastructure_events
                ),
                "infrastructure_events": infrastructure_events,
                "elapsed_seconds": round(elapsed, 3),
            },
        )
        return PageRunResult(
            page,
            "failed",
            None,
            errors,
            cache_hit,
            attempts_used > 1 or repaired,
            attempts_used,
            elapsed,
            False,
            _infrastructure_retry_count(infrastructure_events),
            tuple(infrastructure_events),
        )

    if not cache_hit or repaired:
        cache.store(
            key,
            {
                "key": key,
                "created_at": utc_now(),
                "backend": backend.name,
                "model": profile.model,
                "reasoning_effort": profile.reasoning_effort,
                "draft": draft,
            },
        )

    candidate_path = page_root / "candidate.json"
    shared_candidate_path = run_root / "candidates" / f"page-{page:03d}.json"
    write_json(candidate_path, candidate)
    write_json(shared_candidate_path, candidate)
    elapsed = time.perf_counter() - started
    write_json(
        page_root / "status.json",
        {
            "page": page,
            "status": "valid_candidate",
            "cache_hit": cache_hit,
            "cache_key": key,
            "warnings": validation["warnings"],
            "repair_used": attempts_used > 1 or repaired,
            "attempt_count": attempts_used,
            "infrastructure_retries": _infrastructure_retry_count(
                infrastructure_events
            ),
            "infrastructure_events": infrastructure_events,
            "elapsed_seconds": round(elapsed, 3),
        },
    )
    return PageRunResult(
        page,
        "valid_candidate",
        shared_candidate_path,
        [],
        cache_hit,
        attempts_used > 1 or repaired,
        attempts_used,
        elapsed,
        False,
        _infrastructure_retry_count(infrastructure_events),
        tuple(infrastructure_events),
    )


def _resumed_result(
    *, source: SourcePackage, page: int, run_root: Path
) -> PageRunResult | None:
    candidate_path = run_root / "candidates" / f"page-{page:03d}.json"
    status_path = run_root / "pages" / f"page-{page:03d}" / "status.json"
    if not candidate_path.is_file() or not status_path.is_file():
        return None
    candidate = load_json(candidate_path)
    validation = validate_page(candidate, source)
    if not validation.valid:
        return None
    status = load_json(status_path)
    if status.get("status") != "valid_candidate":
        return None
    return PageRunResult(
        page=page,
        status="valid_candidate",
        candidate_path=candidate_path,
        errors=[],
        cache_hit=bool(status.get("cache_hit", False)),
        repair_used=bool(status.get("repair_used", False)),
        attempt_count=int(status.get("attempt_count", 1)),
        elapsed_seconds=float(status.get("elapsed_seconds", 0.0)),
        resumed=True,
        infrastructure_retries=int(status.get("infrastructure_retries", 0)),
        infrastructure_events=tuple(status.get("infrastructure_events", [])),
    )


def run_batch(
    *,
    source: SourcePackage,
    selection_path: Path,
    profile: EvaluationProfile,
    backend: TranscriptionBackend | None = None,
    force: bool = False,
    allow_large_batch: bool = False,
    run_id: str | None = None,
    run_root: Path | None = None,
    resume: bool = False,
) -> tuple[str, list[PageRunResult]]:
    selection = load_json(selection_path)
    pages = sorted({int(case["page"]) for case in selection["cases"]})
    if len(pages) > profile.max_batch_pages and not allow_large_batch:
        raise ValueError(
            f"Selection contains {len(pages)} pages; the batch limit is "
            f"{profile.max_batch_pages}. Use an explicit large-batch override."
        )
    if run_id is None:
        batch_id = utc_now().replace(":", "").replace("-", "").replace("Z", "Z-eval")
    else:
        batch_id = run_id
    target_root = run_root or project_root() / "artifacts" / "runs" / batch_id
    target_root.mkdir(parents=True, exist_ok=True)
    actual_backend = backend or CodexExecBackend()
    manifest_path = target_root / "run-manifest.json"
    prior_manifest = load_json(manifest_path) if resume and manifest_path.is_file() else None
    current_profile = {
        "backend": actual_backend.name,
        "model": profile.model,
        "reasoning_effort": profile.reasoning_effort,
        "codex_version": codex_version() if actual_backend.name == "codex-exec" else None,
        "started_at": utc_now(),
        "resume": resume,
    }
    profile_history = list((prior_manifest or {}).get("execution_profiles", []))
    if prior_manifest and not profile_history:
        profile_history.append(
            {
                "backend": prior_manifest.get("backend"),
                "model": prior_manifest.get("model"),
                "reasoning_effort": prior_manifest.get("reasoning_effort"),
                "codex_version": prior_manifest.get("codex_version"),
                "started_at": prior_manifest.get("created_at"),
                "resume": False,
            }
        )
    if not profile_history or any(
        profile_history[-1].get(key) != current_profile[key]
        for key in ("backend", "model", "reasoning_effort", "codex_version")
    ):
        profile_history.append(current_profile)
    manifest = {
        "run_id": batch_id,
        "created_at": (prior_manifest or {}).get("created_at", utc_now()),
        "selection": str(selection_path),
        "selection_sha256": sha256_file(selection_path),
        "pages": pages,
        "backend": actual_backend.name,
        "model": profile.model,
        "reasoning_effort": profile.reasoning_effort,
        "codex_version": codex_version() if actual_backend.name == "codex-exec" else None,
        "execution_profiles": profile_history,
        "source_root": str(source.root),
        "source_pdf_sha256": source.lock["source_sha256"],
        "source_lock_sha256": sha256_file(source.lock_path),
        "resume_requested": resume,
        "resume_count": int((prior_manifest or {}).get("resume_count", 0)) + (1 if resume else 0),
        "status": "running",
    }
    write_json(manifest_path, manifest)

    results: list[PageRunResult] = []
    for page in pages:
        result = None if force or not resume else _resumed_result(
            source=source, page=page, run_root=target_root
        )
        if result is None:
            result = transcribe_page(
                source=source,
                page=page,
                run_id=batch_id,
                run_root=target_root,
                profile=profile,
                backend=actual_backend,
                force=force,
            )
        results.append(result)

    valid_pages: dict[int, dict[str, Any]] = {}
    for result in results:
        if result.candidate_path:
            valid_pages[result.page] = load_json(result.candidate_path)
    reconciliation_reports: list[dict[str, Any]] = []
    for left_number in sorted(valid_pages):
        right_number = left_number + 1
        if right_number not in valid_pages:
            continue
        left = valid_pages[left_number]
        right = valid_pages[right_number]
        reconciliation = reconcile_adjacent_pages(left, right)
        reconciliation_reports.append(
            {
                "left_page": left_number,
                "right_page": right_number,
                **reconciliation.as_dict(),
            }
        )
        for number, candidate in ((left_number, left), (right_number, right)):
            write_json(target_root / "candidates" / f"page-{number:03d}.json", candidate)
            write_json(target_root / "pages" / f"page-{number:03d}" / "candidate.json", candidate)

    reciprocal = validate_reciprocal_pages(list(valid_pages.values()))
    write_json(
        target_root / "reconciliation.json",
        {
            "pairs": reconciliation_reports,
            "validation": reciprocal.as_dict(),
        },
    )
    duplicate_reports = detect_cross_page_duplicates(valid_pages)
    write_json(
        target_root / "cross-page-duplicates.json",
        {"threshold_words": 8, "findings": duplicate_reports},
    )

    final_invalid: list[int] = []
    for page, candidate in sorted(valid_pages.items()):
        final_report = validate_page(candidate, source)
        write_json(
            target_root / "pages" / f"page-{page:03d}" / "final-validation.json",
            final_report.as_dict(),
        )
        if not final_report.valid:
            final_invalid.append(page)

    failed = [result.page for result in results if result.status == "failed"]
    manifest = load_json(manifest_path)
    manifest["completed_at"] = utc_now()
    manifest["status"] = (
        "complete"
        if not failed and not final_invalid and reciprocal.valid
        else "partial_failure"
    )
    manifest["failed_pages"] = failed
    manifest["final_validation_failed_pages"] = final_invalid
    manifest["duplicate_cross_page_findings"] = len(duplicate_reports)
    manifest["results"] = [result.as_dict() for result in results]
    schema_path = target_root / "input-schema.json"
    specification_path = target_root / "input-transcription-specification.md"
    manifest["input_schema_sha256"] = sha256_file(schema_path) if schema_path.is_file() else None
    manifest["transcription_specification_sha256"] = (
        sha256_file(specification_path) if specification_path.is_file() else None
    )
    write_json(manifest_path, manifest)
    return batch_id, results
