"""Thread-safe reference ledger for job/run and callback correlation."""

from __future__ import annotations

import copy
import threading
from dataclasses import dataclass
from typing import Dict, Optional, Protocol, Tuple

from arr_processing.contracts import ProcessingRequest, ProcessingResult
from arr_processing.errors import ProcessingError


@dataclass(frozen=True)
class ProcessingRecord:
    request: ProcessingRequest
    request_sha256: str
    idempotency_key: str
    remote_run_id: Optional[str]
    remote_status: str
    result: Optional[ProcessingResult] = None
    failure_code: Optional[str] = None


class ProcessingLedger(Protocol):
    def reserve(
        self,
        request: ProcessingRequest,
        request_sha256: str,
        idempotency_key: str,
    ) -> ProcessingRecord:
        ...

    def get(self, job_id: str, attempt_no: int) -> Optional[ProcessingRecord]:
        ...

    def bind_run(
        self,
        job_id: str,
        attempt_no: int,
        remote_run_id: str,
        remote_status: str,
    ) -> ProcessingRecord:
        ...

    def update_status(
        self,
        job_id: str,
        attempt_no: int,
        remote_status: str,
    ) -> ProcessingRecord:
        ...

    def fail_start(
        self,
        job_id: str,
        attempt_no: int,
        failure_code: str,
    ) -> ProcessingRecord:
        ...

    def attach_result(self, result: ProcessingResult) -> Tuple[ProcessingRecord, bool]:
        ...

    def accept_callback(self, delivery_id: str, signed_sha256: str) -> bool:
        ...


class InMemoryProcessingLedger:
    """Reference implementation; production should map this protocol to PostgreSQL."""

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._records: Dict[Tuple[str, int], ProcessingRecord] = {}
        self._runs: Dict[str, Tuple[str, int]] = {}
        self._callbacks: Dict[str, str] = {}

    def reserve(
        self,
        request: ProcessingRequest,
        request_sha256: str,
        idempotency_key: str,
    ) -> ProcessingRecord:
        key = (request.job_id, request.attempt_no)
        with self._lock:
            existing = self._records.get(key)
            if existing is not None:
                if (
                    existing.request != request
                    or existing.request_sha256 != request_sha256
                    or existing.idempotency_key != idempotency_key
                ):
                    raise ProcessingError(
                        "PROCESSING_JOB_CONFLICT", "processing attempt identity conflicts"
                    )
                return copy.deepcopy(existing)
            record = ProcessingRecord(
                request=request,
                request_sha256=request_sha256,
                idempotency_key=idempotency_key,
                remote_run_id=None,
                remote_status="reserved",
            )
            self._records[key] = record
            return copy.deepcopy(record)

    def get(self, job_id: str, attempt_no: int) -> Optional[ProcessingRecord]:
        with self._lock:
            record = self._records.get((job_id, attempt_no))
            return copy.deepcopy(record) if record is not None else None

    def bind_run(
        self,
        job_id: str,
        attempt_no: int,
        remote_run_id: str,
        remote_status: str,
    ) -> ProcessingRecord:
        key = (job_id, attempt_no)
        with self._lock:
            record = self._required(key)
            run_owner = self._runs.get(remote_run_id)
            if run_owner is not None and run_owner != key:
                raise ProcessingError(
                    "PROCESSING_RUN_CONFLICT", "remote run is already correlated"
                )
            if record.remote_run_id is not None and record.remote_run_id != remote_run_id:
                raise ProcessingError(
                    "PROCESSING_RUN_CONFLICT", "processing attempt has a different remote run"
                )
            updated = ProcessingRecord(
                request=record.request,
                request_sha256=record.request_sha256,
                idempotency_key=record.idempotency_key,
                remote_run_id=remote_run_id,
                remote_status=remote_status,
                result=record.result,
                failure_code=record.failure_code,
            )
            self._records[key] = updated
            self._runs[remote_run_id] = key
            return copy.deepcopy(updated)

    def update_status(
        self,
        job_id: str,
        attempt_no: int,
        remote_status: str,
    ) -> ProcessingRecord:
        key = (job_id, attempt_no)
        with self._lock:
            record = self._required(key)
            updated = ProcessingRecord(
                request=record.request,
                request_sha256=record.request_sha256,
                idempotency_key=record.idempotency_key,
                remote_run_id=record.remote_run_id,
                remote_status=remote_status,
                result=record.result,
                failure_code=record.failure_code,
            )
            self._records[key] = updated
            return copy.deepcopy(updated)

    def fail_start(
        self,
        job_id: str,
        attempt_no: int,
        failure_code: str,
    ) -> ProcessingRecord:
        key = (job_id, attempt_no)
        with self._lock:
            record = self._required(key)
            if record.remote_run_id is not None:
                raise ProcessingError(
                    "PROCESSING_RUN_CONFLICT",
                    "processing attempt already has a remote run",
                )
            if record.remote_status == "failed":
                return copy.deepcopy(record)
            updated = ProcessingRecord(
                request=record.request,
                request_sha256=record.request_sha256,
                idempotency_key=record.idempotency_key,
                remote_run_id=None,
                remote_status="failed",
                result=record.result,
                failure_code=failure_code,
            )
            self._records[key] = updated
            return copy.deepcopy(updated)

    def attach_result(self, result: ProcessingResult) -> Tuple[ProcessingRecord, bool]:
        key = (result.job_id, result.attempt_no)
        with self._lock:
            record = self._required(key)
            if record.remote_run_id != result.remote_run_id:
                raise ProcessingError(
                    "PROCESSING_RESULT_MISMATCH", "result remote run does not match its attempt"
                )
            if record.result is not None:
                if record.result != result:
                    raise ProcessingError(
                        "PROCESSING_RESULT_CONFLICT", "processing attempt has a different result"
                    )
                return copy.deepcopy(record), True
            updated = ProcessingRecord(
                request=record.request,
                request_sha256=record.request_sha256,
                idempotency_key=record.idempotency_key,
                remote_run_id=record.remote_run_id,
                remote_status="delivered",
                result=result,
                failure_code=None,
            )
            self._records[key] = updated
            return copy.deepcopy(updated), False

    def accept_callback(self, delivery_id: str, signed_sha256: str) -> bool:
        with self._lock:
            existing = self._callbacks.get(delivery_id)
            if existing is None:
                self._callbacks[delivery_id] = signed_sha256
                return False
            if existing != signed_sha256:
                raise ProcessingError(
                    "PROCESSING_CALLBACK_CONFLICT", "delivery identifier conflicts"
                )
            return True

    def _required(self, key: Tuple[str, int]) -> ProcessingRecord:
        try:
            return self._records[key]
        except KeyError:
            raise ProcessingError(
                "PROCESSING_JOB_NOT_FOUND", "processing attempt is not registered"
            ) from None