"""Thin data-access layer over the ORM.

Repositories encapsulate the SQL so the API + Celery task code never has to
know about sessions, transactions, or the row → schema mapping.
"""

from __future__ import annotations

import copy
import re
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
from uuid import UUID

from sqlalchemy import select
from sqlalchemy.orm import Session

from ocr_sprint.db.models import JobCorrectionRow, JobRow
from ocr_sprint.schemas.document import DocumentStatus, SourceKind


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


class JobNotFoundError(LookupError):
    """Raised by API code when GET /documents/{id} hits a missing row."""


class InvalidFieldPathError(ValueError):
    """Raised when a PATCH request references an unsupported field path."""


class JobAlreadyApprovedError(RuntimeError):
    """Raised when a PATCH is attempted against an already-approved job."""


class JobNotCompletedError(RuntimeError):
    """Raised when a PATCH/approve is attempted against a job that hasn't
    produced a ``result`` payload yet (e.g. still pending or failed).
    """


@dataclass(frozen=True)
class AppliedCorrection:
    """Internal record of a correction that successfully applied to the
    in-memory ``result`` dict. The repository turns this into a persisted
    ``JobCorrectionRow`` after the whole batch is validated.
    """

    field_path: str
    old_value: Any
    new_value: Any
    reason: str | None


# Allow-list of top-level keys we let reviewers edit. Keeps the attack
# surface small: they can't inject arbitrary fields into the JSON blob.
_ALLOWED_ROOTS: frozenset[str] = frozenset({"header", "ttd", "personel", "untuk"})

# Matches a single path segment like ``personel[3]`` — supports at most one
# index per segment, enough for the list fields we care about.
_SEGMENT_RE = re.compile(r"^([a-zA-Z_][a-zA-Z0-9_]*)(?:\[(\d+)\])?$")


def _split_path(path: str) -> list[tuple[str, int | None]]:
    """Parse ``header.nomor_sprint`` or ``personel[2].nrp`` into segments.

    Returns list of ``(name, index_or_none)`` tuples. Raises
    ``InvalidFieldPathError`` on malformed input so the caller can surface
    a 400 to the client.
    """
    if not path or path.startswith(".") or path.endswith("."):
        raise InvalidFieldPathError(f"Invalid field path: {path!r}")

    parts = path.split(".")
    out: list[tuple[str, int | None]] = []
    for part in parts:
        match = _SEGMENT_RE.match(part)
        if match is None:
            raise InvalidFieldPathError(f"Invalid segment in path: {part!r}")
        name = match.group(1)
        idx_raw = match.group(2)
        idx = int(idx_raw) if idx_raw is not None else None
        out.append((name, idx))

    if out[0][0] not in _ALLOWED_ROOTS:
        raise InvalidFieldPathError(
            f"Field path root {out[0][0]!r} not in allowed roots {sorted(_ALLOWED_ROOTS)!r}"
        )
    return out


def _apply_path(data: dict[str, Any], path: str, new_value: Any) -> Any:
    """Apply a single correction to ``data`` in place. Returns the old
    value so the caller can record it in the audit row.

    Does NOT validate that the new value matches the field's expected
    type — that's the reviewer's responsibility; the whole point of HITL
    is to let humans override the model's typing.
    """
    segments = _split_path(path)
    cursor: Any = data
    for name, idx in segments[:-1]:
        if not isinstance(cursor, dict) or name not in cursor:
            raise InvalidFieldPathError(f"Cannot traverse to {path!r}: missing {name!r}")
        cursor = cursor[name]
        if idx is not None:
            if not isinstance(cursor, list) or idx >= len(cursor):
                raise InvalidFieldPathError(
                    f"Cannot traverse to {path!r}: index [{idx}] out of range"
                )
            cursor = cursor[idx]

    name, idx = segments[-1]
    if idx is not None:
        # Terminal segment is a list-element, e.g. ``untuk[2]``.
        if not isinstance(cursor, dict) or name not in cursor:
            raise InvalidFieldPathError(f"Cannot apply to {path!r}: missing container {name!r}")
        container = cursor[name]
        if not isinstance(container, list) or idx >= len(container):
            raise InvalidFieldPathError(f"Cannot apply to {path!r}: index [{idx}] out of range")
        old = container[idx]
        container[idx] = new_value
        return old

    if not isinstance(cursor, dict):
        raise InvalidFieldPathError(f"Cannot apply to {path!r}: parent is not an object")
    old = cursor.get(name)
    cursor[name] = new_value
    return old


class JobRepository:
    """SQL-backed repository for `jobs` rows."""

    def __init__(self, session: Session) -> None:
        self.session = session

    # ---------- writes ----------

    def create(
        self,
        *,
        job_id: UUID,
        filename: str,
        source_kind: SourceKind,
        blob_key: str,
    ) -> JobRow:
        row = JobRow(
            job_id=job_id,
            status=DocumentStatus.PENDING.value,
            source_kind=source_kind.value,
            filename=filename,
            blob_key=blob_key,
        )
        self.session.add(row)
        self.session.flush()
        return row

    def mark_processing(self, job_id: UUID) -> None:
        row = self._get_or_raise(job_id)
        row.status = DocumentStatus.PROCESSING.value
        row.updated_at = _utcnow()

    def mark_completed(
        self,
        job_id: UUID,
        *,
        status: DocumentStatus,
        confidence: float,
        result: dict[str, Any],
        review_flags: list[str],
    ) -> None:
        row = self._get_or_raise(job_id)
        row.status = status.value
        row.confidence = confidence
        row.result = result
        row.review_flags = review_flags
        row.error = None
        row.updated_at = _utcnow()

    def mark_failed(self, job_id: UUID, *, error: str) -> None:
        row = self._get_or_raise(job_id)
        row.status = DocumentStatus.FAILED.value
        row.error = error[:2048]
        row.updated_at = _utcnow()

    # ---------- reads ----------

    def get(self, job_id: UUID) -> JobRow | None:
        stmt = select(JobRow).where(JobRow.job_id == job_id)
        return self.session.scalar(stmt)

    def get_or_raise(self, job_id: UUID) -> JobRow:
        return self._get_or_raise(job_id)

    def _get_or_raise(self, job_id: UUID) -> JobRow:
        row = self.get(job_id)
        if row is None:
            raise JobNotFoundError(f"Job not found: {job_id}")
        return row

    # ---------- Phase 6 — HITL ----------

    def apply_corrections(
        self,
        job_id: UUID,
        *,
        corrections: list[tuple[str, Any, str | None]],
        corrected_by: str | None,
    ) -> list[JobCorrectionRow]:
        """Apply a batch of field corrections atomically.

        ``corrections`` is a list of ``(path, new_value, reason)`` tuples.
        Returns the persisted audit rows so the caller can surface them in
        the response.

        Raises
        ------
        JobNotFoundError
            If the row doesn't exist.
        JobNotCompletedError
            If the job hasn't produced a result yet (status pending /
            processing / failed).
        JobAlreadyApprovedError
            If the job has been approved — edits are locked.
        InvalidFieldPathError
            If any path is malformed or references a disallowed root.
        """
        row = self._get_or_raise(job_id)
        if row.result is None:
            raise JobNotCompletedError(
                f"Job {job_id} has no result to correct (status={row.status})"
            )
        if row.approved:
            raise JobAlreadyApprovedError(f"Job {job_id} is already approved; edits are locked")

        # Deep-copy so we can roll back in memory if any correction fails.
        # The underlying JSON column will only be re-assigned once every
        # path applied cleanly.
        working = copy.deepcopy(row.result)
        applied: list[AppliedCorrection] = []
        for path, new_value, reason in corrections:
            old_value = _apply_path(working, path, new_value)
            applied.append(
                AppliedCorrection(
                    field_path=path, old_value=old_value, new_value=new_value, reason=reason
                )
            )

        # Persist audit rows first; if they fail the session rollback also
        # undoes the result-column update we're about to do.
        persisted: list[JobCorrectionRow] = []
        for event in applied:
            row_event = JobCorrectionRow(
                job_id=job_id,
                field_path=event.field_path,
                old_value=event.old_value,
                new_value=event.new_value,
                corrected_by=corrected_by,
                reason=event.reason,
            )
            self.session.add(row_event)
            persisted.append(row_event)

        # Clear review flags that the correction has resolved. Right now we
        # only auto-clear MISSING_FIELD when any corrected field previously
        # held a null/empty value — the reviewer explicitly filled a gap.
        row.result = working
        row.review_flags = _recompute_flags(
            original_flags=list(row.review_flags or []),
            applied=applied,
            working_result=working,
        )
        row.updated_at = _utcnow()

        self.session.flush()
        return persisted

    def list_corrections(self, job_id: UUID) -> list[JobCorrectionRow]:
        """Return the full audit trail for ``job_id`` in chronological order."""
        # ``get_or_raise`` so callers get a 404 instead of an empty list
        # when the job itself doesn't exist.
        self._get_or_raise(job_id)
        stmt = (
            select(JobCorrectionRow)
            .where(JobCorrectionRow.job_id == job_id)
            .order_by(JobCorrectionRow.corrected_at, JobCorrectionRow.id)
        )
        return list(self.session.scalars(stmt))

    def approve(self, job_id: UUID, *, reviewed_by: str | None) -> JobRow:
        """Mark a job as approved. Idempotent — re-approving is a no-op
        that keeps the original ``reviewed_at`` (so the audit trail stays
        intact).
        """
        row = self._get_or_raise(job_id)
        if row.result is None:
            raise JobNotCompletedError(
                f"Job {job_id} has no result to approve (status={row.status})"
            )
        if row.approved:
            return row
        row.approved = True
        row.reviewed_by = reviewed_by
        row.reviewed_at = _utcnow()
        row.updated_at = row.reviewed_at
        return row


def _recompute_flags(
    *,
    original_flags: list[str],
    applied: list[AppliedCorrection],
    working_result: dict[str, Any],
) -> list[str]:
    """Update review flags in light of the corrections just applied.

    Keeps the policy simple on purpose:
    * ``missing_field`` is removed if after the edit every required
      header field is non-empty.
    * Other flags stay untouched — the reviewer should either correct the
      underlying issue (which this helper can detect) or explicitly
      approve the result as-is (which bypasses the flag list).
    """
    flags = list(original_flags)
    if "missing_field" in flags:
        header = working_result.get("header") or {}
        filled = all(bool(header.get(key)) for key in ("nomor_sprint", "satuan_penerbit"))
        if filled:
            flags = [f for f in flags if f != "missing_field"]

    # ``applied`` isn't used directly in this MVP rule, but we keep the
    # parameter so future policies can inspect exactly what changed
    # without re-diffing the blob.
    _ = applied
    return flags
