"""Durable checkpoint journal for the generic Online Agent worker.

The JSON task state remains a fast cache.  This SQLite journal is the durable,
append-only evidence used for restart/resume/replay decisions.
"""
from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import sqlite3
from typing import Any, Mapping

PHASE_ORDER = {
    "claimed": 10,
    "workspace_prepared": 20,
    "verification_passed": 30,
    "source_capture_done": 40,
    "commit_done": 40,
    "reconciliation_required": 45,
    "result_pending": 50,
    "result_submitted": 60,
}


class CheckpointError(RuntimeError):
    pass


def canonical(value: object) -> bytes:
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str).encode("utf-8")


def sha256_json(value: object) -> str:
    return hashlib.sha256(canonical(value)).hexdigest()


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


@dataclass(frozen=True)
class CheckpointBinding:
    task_id: str
    attempt: int
    execution_digest: str
    job_id: str
    repository: str
    base_oid: str
    expected_branch_oid: str
    target_branch: str
    generation: int

    def material(self) -> dict[str, Any]:
        return {
            "task_id": self.task_id,
            "attempt": self.attempt,
            "execution_digest": self.execution_digest,
            "job_id": self.job_id,
            "repository": self.repository,
            "base_oid": self.base_oid,
            "expected_branch_oid": self.expected_branch_oid,
            "target_branch": self.target_branch,
            "generation": self.generation,
        }

    @property
    def binding_sha256(self) -> str:
        return sha256_json(self.material())


class CheckpointStore:
    def __init__(self, path: Path) -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self._connect() as connection:
            connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS checkpoints(
                  task_id TEXT NOT NULL,
                  attempt INTEGER NOT NULL,
                  sequence INTEGER NOT NULL,
                  phase TEXT NOT NULL,
                  phase_order INTEGER NOT NULL,
                  execution_digest TEXT NOT NULL,
                  job_id TEXT NOT NULL,
                  repository TEXT NOT NULL,
                  base_oid TEXT NOT NULL,
                  expected_branch_oid TEXT NOT NULL,
                  target_branch TEXT NOT NULL,
                  generation INTEGER NOT NULL,
                  binding_sha256 TEXT NOT NULL,
                  payload_json TEXT NOT NULL,
                  payload_sha256 TEXT NOT NULL,
                  created_at TEXT NOT NULL,
                  PRIMARY KEY(task_id, attempt, sequence)
                );
                CREATE INDEX IF NOT EXISTS idx_checkpoints_task_latest
                  ON checkpoints(task_id, attempt DESC, sequence DESC);
                CREATE UNIQUE INDEX IF NOT EXISTS idx_checkpoints_exact_replay
                  ON checkpoints(task_id, attempt, phase, payload_sha256);
                CREATE TABLE IF NOT EXISTS checkpoint_heads(
                  task_id TEXT PRIMARY KEY,
                  attempt INTEGER NOT NULL,
                  sequence INTEGER NOT NULL,
                  phase TEXT NOT NULL,
                  phase_order INTEGER NOT NULL,
                  execution_digest TEXT NOT NULL,
                  binding_sha256 TEXT NOT NULL,
                  payload_sha256 TEXT NOT NULL,
                  updated_at TEXT NOT NULL
                );
                """
            )

    def _connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.path, timeout=30.0)
        connection.row_factory = sqlite3.Row
        connection.execute("PRAGMA foreign_keys=ON")
        connection.execute("PRAGMA journal_mode=WAL")
        connection.execute("PRAGMA synchronous=FULL")
        connection.execute("PRAGMA busy_timeout=30000")
        return connection

    @staticmethod
    def _phase_order(phase: str) -> int:
        try:
            return PHASE_ORDER[phase]
        except KeyError as exc:
            raise CheckpointError("CHECKPOINT_PHASE_INVALID") from exc

    def append(
        self,
        binding: CheckpointBinding,
        phase: str,
        payload: Mapping[str, Any],
    ) -> dict[str, Any]:
        order = self._phase_order(phase)
        payload_value = dict(payload)
        payload_sha = sha256_json(payload_value)
        with self._connect() as connection:
            connection.execute("BEGIN IMMEDIATE")
            replay = connection.execute(
                """
                SELECT * FROM checkpoints
                 WHERE task_id=? AND attempt=? AND phase=? AND payload_sha256=?
                """,
                (binding.task_id, binding.attempt, phase, payload_sha),
            ).fetchone()
            if replay is not None:
                connection.commit()
                return dict(replay)

            head = connection.execute(
                "SELECT * FROM checkpoint_heads WHERE task_id=?",
                (binding.task_id,),
            ).fetchone()
            if head is not None:
                head_attempt = int(head["attempt"])
                head_order = int(head["phase_order"])
                if binding.attempt < head_attempt:
                    raise CheckpointError("CHECKPOINT_ATTEMPT_REGRESSION")
                if binding.attempt == head_attempt:
                    if head["execution_digest"] != binding.execution_digest:
                        raise CheckpointError("CHECKPOINT_EXECUTION_DIGEST_CONFLICT")
                    if head["binding_sha256"] != binding.binding_sha256:
                        raise CheckpointError("CHECKPOINT_SCOPE_MISMATCH_DENY")
                    if order < head_order and phase != "reconciliation_required":
                        raise CheckpointError("CHECKPOINT_PHASE_REGRESSION")
                    sequence = int(head["sequence"]) + 1
                else:
                    sequence = 1
            else:
                sequence = 1

            created_at = now()
            connection.execute(
                """
                INSERT INTO checkpoints(
                  task_id,attempt,sequence,phase,phase_order,execution_digest,job_id,
                  repository,base_oid,expected_branch_oid,target_branch,generation,
                  binding_sha256,payload_json,payload_sha256,created_at
                ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
                """,
                (
                    binding.task_id,
                    binding.attempt,
                    sequence,
                    phase,
                    order,
                    binding.execution_digest,
                    binding.job_id,
                    binding.repository,
                    binding.base_oid,
                    binding.expected_branch_oid,
                    binding.target_branch,
                    binding.generation,
                    binding.binding_sha256,
                    canonical(payload_value).decode("utf-8"),
                    payload_sha,
                    created_at,
                ),
            )
            connection.execute(
                """
                INSERT INTO checkpoint_heads(
                  task_id,attempt,sequence,phase,phase_order,execution_digest,
                  binding_sha256,payload_sha256,updated_at
                ) VALUES (?,?,?,?,?,?,?,?,?)
                ON CONFLICT(task_id) DO UPDATE SET
                  attempt=excluded.attempt,
                  sequence=excluded.sequence,
                  phase=excluded.phase,
                  phase_order=excluded.phase_order,
                  execution_digest=excluded.execution_digest,
                  binding_sha256=excluded.binding_sha256,
                  payload_sha256=excluded.payload_sha256,
                  updated_at=excluded.updated_at
                """,
                (
                    binding.task_id,
                    binding.attempt,
                    sequence,
                    phase,
                    order,
                    binding.execution_digest,
                    binding.binding_sha256,
                    payload_sha,
                    created_at,
                ),
            )
            connection.commit()
            row = connection.execute(
                "SELECT * FROM checkpoints WHERE task_id=? AND attempt=? AND sequence=?",
                (binding.task_id, binding.attempt, sequence),
            ).fetchone()
            return dict(row)

    def latest(self, task_id: str) -> dict[str, Any] | None:
        with self._connect() as connection:
            row = connection.execute(
                """
                SELECT * FROM checkpoints
                 WHERE task_id=?
                 ORDER BY attempt DESC, sequence DESC
                 LIMIT 1
                """,
                (task_id,),
            ).fetchone()
            if row is None:
                return None
            value = dict(row)
            value["payload"] = json.loads(value.pop("payload_json"))
            return value

    def latest_for_attempt(self, task_id: str, attempt: int) -> dict[str, Any] | None:
        with self._connect() as connection:
            row = connection.execute(
                """
                SELECT * FROM checkpoints
                 WHERE task_id=? AND attempt=?
                 ORDER BY sequence DESC
                 LIMIT 1
                """,
                (task_id, attempt),
            ).fetchone()
            if row is None:
                return None
            value = dict(row)
            value["payload"] = json.loads(value.pop("payload_json"))
            return value

    def verify_latest_binding(self, binding: CheckpointBinding) -> dict[str, Any] | None:
        value = self.latest(binding.task_id)
        if value is None:
            return None
        if int(value["attempt"]) > binding.attempt:
            raise CheckpointError("CHECKPOINT_FUTURE_ATTEMPT")
        # A previous attempt may be reused only when the immutable execution/scope binding matches.
        if value["execution_digest"] != binding.execution_digest:
            return None
        expected = CheckpointBinding(
            task_id=value["task_id"],
            attempt=int(value["attempt"]),
            execution_digest=value["execution_digest"],
            job_id=value["job_id"],
            repository=value["repository"],
            base_oid=value["base_oid"],
            expected_branch_oid=value["expected_branch_oid"],
            target_branch=value["target_branch"],
            generation=int(value["generation"]),
        )
        current_scope = binding.material()
        historical_scope = expected.material()
        current_scope.pop("attempt", None)
        historical_scope.pop("attempt", None)
        if sha256_json(current_scope) != sha256_json(historical_scope):
            raise CheckpointError("CHECKPOINT_SCOPE_MISMATCH_DENY")
        return value