"""Core ground-truth export logic.

The service reads from ``jobs`` and ``job_corrections`` and produces
``GroundTruthSample`` rows. The HTTP layer and the CLI both go through
this module, so the export logic stays in exactly one place.

Design notes
------------
* **Reverse-replay for ``initial_result``.** Audit rows store both
  ``old_value`` and ``new_value``. We start from the current
  ``final_result`` and undo each correction in reverse chronological
  order to recover the pipeline's original pre-HITL output. This is
  what fine-tuning needs: the raw mistake next to the reviewer fix.
* **No deep copying in the caller.** The service always returns freshly
  computed dicts so mutating one sample can't leak into the next.
* **Generator by design.** Large exports shouldn't load the whole table
  into memory; the HTTP route streams JSONL line-by-line.
"""

from __future__ import annotations

import copy
import re
from collections.abc import Iterator
from dataclasses import dataclass
from datetime import datetime
from typing import Any

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from ocr_sprint.db.models import JobCorrectionRow, JobRow
from ocr_sprint.schemas.document import DocumentStatus
from ocr_sprint.schemas.ground_truth import (
    FieldCorrectionCount,
    GroundTruthCorrection,
    GroundTruthSample,
    GroundTruthStats,
)


@dataclass(frozen=True)
class GroundTruthFilters:
    """HTTP / CLI filters for ``iter_ground_truth_samples``.

    All fields are optional; ``approved_only=True`` is the default because
    only approved samples are actually safe to train on.
    """

    since: datetime | None = None
    until: datetime | None = None
    approved_only: bool = True
    has_corrections: bool | None = None
    limit: int | None = None


def build_initial_result(
    *,
    final_result: dict[str, Any] | None,
    corrections: list[JobCorrectionRow],
) -> dict[str, Any] | None:
    """Replay ``corrections`` backwards over ``final_result`` to recover
    the pre-HITL version. Returns ``None`` if there's nothing to
    reconstruct.
    """
    if final_result is None:
        return None
    restored = copy.deepcopy(final_result)
    for event in reversed(corrections):
        _assign_path(restored, event.field_path, event.old_value)
    return restored


def _assign_path(data: dict[str, Any], path: str, value: Any) -> None:
    """Set ``data[path] = value`` using the same dotted syntax the HITL
    repository understands. Silently ignores paths that no longer resolve
    (the dict shape may have drifted since the event was recorded; in
    that case we just skip the replay step rather than raising).
    """
    segments = _parse(path)
    if not segments:
        return

    cursor: Any = data
    for name, idx in segments[:-1]:
        if not isinstance(cursor, dict) or name not in cursor:
            return
        cursor = cursor[name]
        if idx is not None:
            if not isinstance(cursor, list) or idx >= len(cursor):
                return
            cursor = cursor[idx]

    name, idx = segments[-1]
    if idx is not None:
        if not isinstance(cursor, dict) or name not in cursor:
            return
        container = cursor[name]
        if not isinstance(container, list) or idx >= len(container):
            return
        container[idx] = value
        return

    if not isinstance(cursor, dict):
        return
    cursor[name] = value


_SEGMENT_RE = re.compile(r"^([a-zA-Z_][a-zA-Z0-9_]*)(?:\[(\d+)\])?$")


def _parse(path: str) -> list[tuple[str, int | None]]:
    """Forgiving dotted-path parser — returns ``[]`` on malformed input
    instead of raising, since audit-log replay must not blow up on a
    stale path. The strict variant used by the HITL PATCH route lives
    in ``ocr_sprint.db.repositories``.
    """
    if not path or path.startswith(".") or path.endswith("."):
        return []
    out: list[tuple[str, int | None]] = []
    for part in path.split("."):
        match: re.Match[str] | None = _SEGMENT_RE.match(part)
        if match is None:
            return []
        out.append((match.group(1), int(match.group(2)) if match.group(2) else None))
    return out


def iter_ground_truth_samples(
    session: Session,
    filters: GroundTruthFilters,
) -> Iterator[GroundTruthSample]:
    """Yield ``GroundTruthSample`` rows matching ``filters``.

    Rows are ordered by ``created_at ASC`` so repeated exports are
    deterministic and downstream diffing works.
    """
    stmt = select(JobRow).order_by(JobRow.created_at, JobRow.job_id)
    if filters.approved_only:
        stmt = stmt.where(JobRow.approved.is_(True))
    if filters.since is not None:
        stmt = stmt.where(JobRow.created_at >= filters.since)
    if filters.until is not None:
        stmt = stmt.where(JobRow.created_at <= filters.until)

    remaining = filters.limit
    for job_row in session.scalars(stmt):
        if remaining is not None and remaining <= 0:
            return

        correction_rows = list(
            session.scalars(
                select(JobCorrectionRow)
                .where(JobCorrectionRow.job_id == job_row.job_id)
                .order_by(JobCorrectionRow.corrected_at, JobCorrectionRow.id)
            )
        )

        if filters.has_corrections is True and not correction_rows:
            continue
        if filters.has_corrections is False and correction_rows:
            continue

        initial = build_initial_result(final_result=job_row.result, corrections=correction_rows)
        sample = GroundTruthSample(
            job_id=job_row.job_id,
            filename=job_row.filename,
            source_kind=job_row.source_kind,
            approved=bool(job_row.approved),
            reviewed_by=job_row.reviewed_by,
            reviewed_at=job_row.reviewed_at,
            created_at=job_row.created_at,
            initial_result=initial,
            # Use an ``is None`` check to stay consistent with
            # ``build_initial_result``; otherwise an empty-dict result
            # would produce ``initial_result={}`` but ``final_result=None``.
            final_result=(copy.deepcopy(job_row.result) if job_row.result is not None else None),
            corrections=[
                GroundTruthCorrection(
                    field_path=c.field_path,
                    old_value=c.old_value,
                    new_value=c.new_value,
                    corrected_by=c.corrected_by,
                    reason=c.reason,
                    corrected_at=c.corrected_at,
                )
                for c in correction_rows
            ],
            review_flags=list(job_row.review_flags or []),
            confidence=job_row.confidence,
        )
        if remaining is not None:
            remaining -= 1
        yield sample


def ground_truth_stats(session: Session, *, top_n: int = 10) -> GroundTruthStats:
    """Compute aggregate statistics over the jobs + corrections tables."""
    by_status: dict[str, int] = {
        row[0]: int(row[1])
        for row in session.execute(
            select(JobRow.status, func.count(JobRow.job_id)).group_by(JobRow.status)
        ).all()
    }
    total_jobs = int(sum(by_status.values()))
    approved_jobs = int(
        session.execute(
            select(func.count(JobRow.job_id)).where(JobRow.approved.is_(True))
        ).scalar_one()
    )

    total_corrections = int(session.execute(select(func.count(JobCorrectionRow.id))).scalar_one())
    jobs_with_corrections = int(
        session.execute(select(func.count(func.distinct(JobCorrectionRow.job_id)))).scalar_one()
    )

    # Field-path histogram. A single SQL ``GROUP BY`` is cheaper than
    # pulling every row over the wire, but we still cap at ``top_n`` so
    # a pathological dataset can't drown the response.
    top_rows = session.execute(
        select(
            JobCorrectionRow.field_path,
            func.count(JobCorrectionRow.id).label("n"),
        )
        .group_by(JobCorrectionRow.field_path)
        .order_by(func.count(JobCorrectionRow.id).desc(), JobCorrectionRow.field_path)
        .limit(top_n)
    ).all()

    return GroundTruthStats(
        total_jobs=total_jobs,
        completed_jobs=by_status.get(DocumentStatus.COMPLETED.value, 0),
        needs_review_jobs=by_status.get(DocumentStatus.NEEDS_REVIEW.value, 0),
        failed_jobs=by_status.get(DocumentStatus.FAILED.value, 0),
        approved_jobs=approved_jobs,
        total_corrections=total_corrections,
        jobs_with_corrections=jobs_with_corrections,
        top_corrected_fields=[
            FieldCorrectionCount(field_path=row[0], count=int(row[1])) for row in top_rows
        ],
    )


def serialize_sample_to_jsonl(sample: GroundTruthSample) -> str:
    """Serialize one sample as a newline-terminated JSON string.

    Kept standalone so the HTTP and CLI layers share the exact same
    representation — diffs between a curl download and a CLI dump stay
    byte-identical.
    """
    return sample.model_dump_json() + "\n"
