"""PostgreSQL processing state for restart-safe Agent dispatch and callbacks."""

from __future__ import annotations

import hashlib
import os
from dataclasses import dataclass
from datetime import date
from typing import Any, Callable, Optional, Tuple

from arr_ingestion.contracts import ArtifactRef
from arr_ingestion.repository import IngestionOutcome
from arr_processing.contracts import ProcessingRequest, ProcessingResult
from arr_processing.errors import ProcessingError
from arr_processing.ledger import ProcessingRecord


TARGET_DATABASE = "booking_test"
DATABASE_ENV = "ARR_DATABASE_URL"
_RUN_TERMINAL = frozenset({"accepted", "rejected", "failed", "cancelled"})
_SAFE_START_FAILURE_CODES = frozenset(
    {
        "PROCESSING_REMOTE_PROTOCOL_INVALID",
        "PROCESSING_REMOTE_RATE_LIMITED",
        "PROCESSING_REMOTE_REJECTED",
        "PROCESSING_REMOTE_SUBMISSION_AMBIGUOUS",
        "PROCESSING_REMOTE_UNAVAILABLE",
    }
)


def _default_connect(dsn: str) -> Any:
    try:
        import psycopg  # type: ignore[import-not-found]
    except ImportError:
        raise ProcessingError(
            "DATABASE_DRIVER_UNAVAILABLE",
            "PostgreSQL driver is unavailable",
            retryable=True,
        ) from None
    try:
        return psycopg.connect(dsn, autocommit=False)
    except Exception:
        raise ProcessingError(
            "DATABASE_UNAVAILABLE",
            "processing database is unavailable",
            retryable=True,
        ) from None


@dataclass(frozen=True)
class ProcessingDatabaseConfig:
    dsn: str

    @classmethod
    def from_environment(cls) -> "ProcessingDatabaseConfig":
        dsn = os.environ.get(DATABASE_ENV, "").strip()
        if not dsn:
            raise ProcessingError(
                "DATABASE_CONFIG_MISSING",
                f"{DATABASE_ENV} is required",
            )
        return cls(dsn)


class PostgresProcessingState:
    """Implements ProcessingLedger plus source and committed-outcome lookup."""

    def __init__(
        self,
        config: ProcessingDatabaseConfig,
        connect: Optional[Callable[[str], Any]] = None,
    ) -> None:
        self._config = config
        self._connect = connect or _default_connect

    def _run(self, operation: Callable[[Any], Any], *, read_only: bool = False) -> Any:
        try:
            connection = self._connect(self._config.dsn)
        except ProcessingError:
            raise
        except Exception:
            raise ProcessingError(
                "DATABASE_UNAVAILABLE",
                "processing database is unavailable",
                retryable=True,
            ) from None
        try:
            with connection.transaction():
                with connection.cursor() as cursor:
                    cursor.execute(
                        "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"
                        + (" READ ONLY" if read_only else "")
                    )
                    cursor.execute("SET LOCAL statement_timeout = '15s'")
                    cursor.execute("SET LOCAL lock_timeout = '5s'")
                    cursor.execute("SELECT current_database()")
                    row = cursor.fetchone()
                    if not row or row[0] != TARGET_DATABASE:
                        raise ProcessingError(
                            "DATABASE_TARGET_INVALID",
                            "processing database target is invalid",
                        )
                    return operation(cursor)
        except ProcessingError:
            raise
        except Exception:
            raise ProcessingError(
                "DATABASE_WRITE_FAILED" if not read_only else "DATABASE_QUERY_FAILED",
                "processing state could not be persisted"
                if not read_only
                else "processing state could not be read",
                retryable=True,
            ) from None
        finally:
            connection.close()

    def reserve(
        self,
        request: ProcessingRequest,
        request_sha256: str,
        idempotency_key: str,
    ) -> ProcessingRecord:
        def operation(cursor: Any) -> ProcessingRecord:
            row = self._attempt_row(
                cursor,
                request.job_id,
                request.attempt_no,
                for_update=True,
            )
            if (
                row[1] != "opera_daily"
                or row[2] != request.processor_version
                or str(row[3]) != request.rule_set_sha256
                or str(row[5]) != idempotency_key
            ):
                raise ProcessingError(
                    "PROCESSING_JOB_CONFLICT",
                    "processing attempt identity conflicts",
                )
            return ProcessingRecord(
                request=request,
                request_sha256=request_sha256,
                idempotency_key=idempotency_key,
                remote_run_id=str(row[6]) if row[6] is not None else None,
                remote_status=self._remote_status(str(row[7])),
                failure_code=str(row[9]) if row[9] is not None else None,
            )

        return self._run(operation)

    def get(self, job_id: str, attempt_no: int) -> Optional[ProcessingRecord]:
        def operation(cursor: Any) -> Optional[ProcessingRecord]:
            try:
                row = self._attempt_row(cursor, job_id, attempt_no)
            except ProcessingError as error:
                if error.code == "PROCESSING_JOB_NOT_FOUND":
                    return None
                raise
            placeholder = "persisted_" + hashlib.sha256(
                f"{job_id}\x1f{attempt_no}".encode("utf-8")
            ).hexdigest()[:48]
            request = ProcessingRequest(
                job_id=job_id,
                attempt_no=attempt_no,
                source_file_id=placeholder,
                processor_version=str(row[2]),
                rule_set_sha256=str(row[3]),
            )
            return ProcessingRecord(
                request=request,
                request_sha256=request.sha256(),
                idempotency_key=str(row[5]),
                remote_run_id=str(row[6]) if row[6] is not None else None,
                remote_status=self._remote_status(str(row[7])),
                failure_code=str(row[9]) if row[9] is not None else None,
            )

        return self._run(operation, read_only=True)

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

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

    def fail_start(
        self,
        job_id: str,
        attempt_no: int,
        failure_code: str,
    ) -> ProcessingRecord:
        if failure_code not in _SAFE_START_FAILURE_CODES:
            failure_code = "PROCESSING_REMOTE_PROTOCOL_INVALID"

        def operation(cursor: Any) -> ProcessingRecord:
            row = self._attempt_row(cursor, job_id, attempt_no, for_update=True)
            if row[6] is not None:
                raise ProcessingError(
                    "PROCESSING_RUN_CONFLICT",
                    "processing attempt already has a remote run",
                )
            if str(row[4]) not in _RUN_TERMINAL:
                cursor.execute(
                    """
                    UPDATE ingestion.processing_attempts
                    SET attempt_status = 'failed',
                        failure_code = %s,
                        failure_message = 'remote processing submission did not complete',
                        finished_at = now()
                    WHERE id = %s
                    """,
                    (failure_code, row[8]),
                )
                cursor.execute(
                    """
                    UPDATE ingestion.processing_runs
                    SET run_status = 'failed',
                        failure_code = %s,
                        failure_message = 'remote processing submission did not complete',
                        updated_at = now(),
                        finished_at = now()
                    WHERE id = %s
                    """,
                    (failure_code, row[0]),
                )
                persisted_failure_code = failure_code
            else:
                persisted_failure_code = (
                    str(row[9]) if row[9] is not None else failure_code
                )
            return self._record_from_row(
                row,
                job_id=job_id,
                attempt_no=attempt_no,
                remote_status="failed",
                failure_code=persisted_failure_code,
            )

        return self._run(operation)

    def attach_result(self, result: ProcessingResult) -> Tuple[ProcessingRecord, bool]:
        def operation(cursor: Any) -> Tuple[ProcessingRecord, bool]:
            row = self._attempt_row(
                cursor,
                result.job_id,
                result.attempt_no,
                for_update=True,
            )
            if (
                row[6] != result.remote_run_id
                or row[2] != result.processor_version
                or str(row[3]) != result.rule_set_sha256
            ):
                raise ProcessingError(
                    "PROCESSING_RESULT_MISMATCH",
                    "processing result does not match its request",
                )
            replayed = self._delivery_matches(cursor, result)
            if str(row[4]) not in _RUN_TERMINAL:
                cursor.execute(
                    """
                    UPDATE ingestion.processing_attempts
                    SET attempt_status = 'delivered',
                        failure_code = NULL,
                        failure_message = NULL
                    WHERE id = %s
                    """,
                    (row[8],),
                )
                cursor.execute(
                    """
                    UPDATE ingestion.processing_runs
                    SET run_status = 'validating',
                        updated_at = now()
                    WHERE id = %s
                    """,
                    (row[0],),
                )
            record = self._record_from_row(row, remote_status="delivered", result=result)
            return record, replayed

        return self._run(operation)

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

        def operation(cursor: Any) -> bool:
            cursor.execute(
                """
                SELECT 1
                FROM ingestion.processing_deliveries
                WHERE delivery_key = %s
                """,
                (delivery_id,),
            )
            return cursor.fetchone() is not None

        return self._run(operation, read_only=True)

    def source_for_attempt(self, job_id: str, attempt_no: int) -> ArtifactRef:
        def operation(cursor: Any) -> ArtifactRef:
            cursor.execute(
                """
                SELECT
                    artifact.object_key,
                    artifact.original_filename,
                    artifact.sha256,
                    artifact.byte_size,
                    artifact.mime_type
                FROM ingestion.processing_runs AS run
                JOIN ingestion.processing_attempts AS attempt
                  ON attempt.processing_run_id = run.id
                 AND attempt.attempt_no = %s
                JOIN ingestion.artifacts AS artifact
                  ON artifact.id = run.source_artifact_id
                WHERE run.run_key = %s
                  AND run.pipeline_type = 'opera_daily'
                  AND artifact.artifact_kind = 'opera_xml'
                """,
                (attempt_no, job_id),
            )
            row = cursor.fetchone()
            if row is None:
                raise ProcessingError(
                    "PROCESSING_JOB_NOT_FOUND",
                    "processing attempt is not registered",
                )
            return ArtifactRef.from_dict(
                "source_xml",
                {
                    "object_key": str(row[0]),
                    "original_filename": str(row[1]),
                    "sha256": str(row[2]),
                    "byte_size": int(row[3]),
                    "mime_type": str(row[4]),
                },
            )

        return self._run(operation, read_only=True)

    def outcome_for_delivery(self, delivery_id: str) -> Optional[IngestionOutcome]:
        def operation(cursor: Any) -> Optional[IngestionOutcome]:
            cursor.execute(
                """
                SELECT
                    delivery.delivery_status,
                    run.run_key,
                    delivery.business_date,
                    delivery.daily_version_id,
                    version.version_no
                FROM ingestion.processing_deliveries AS delivery
                JOIN ingestion.processing_runs AS run
                  ON run.id = delivery.processing_run_id
                LEFT JOIN finance.daily_versions AS version
                  ON version.id = delivery.daily_version_id
                WHERE delivery.delivery_key = %s
                """,
                (delivery_id,),
            )
            row = cursor.fetchone()
            if row is None or row[0] not in {"committed", "recorded_failure"}:
                return None
            return IngestionOutcome(
                status=(
                    "already_committed"
                    if row[0] == "committed"
                    else "recorded_failure"
                ),
                job_id=str(row[1]),
                business_date=row[2] if isinstance(row[2], date) else None,
                daily_version_id=int(row[3]) if row[3] is not None else None,
                version_no=int(row[4]) if row[4] is not None else None,
            )

        return self._run(operation, read_only=True)

    def _set_remote_status(
        self,
        job_id: str,
        attempt_no: int,
        remote_status: str,
        *,
        remote_run_id: Optional[str] = None,
    ) -> ProcessingRecord:
        attempt_status, run_status, failure_code, terminal = self._database_status(
            remote_status
        )

        def operation(cursor: Any) -> ProcessingRecord:
            row = self._attempt_row(cursor, job_id, attempt_no, for_update=True)
            persisted_run_id = str(row[6]) if row[6] is not None else None
            if remote_run_id is not None:
                cursor.execute(
                    """
                    SELECT processing_run_id
                    FROM ingestion.processing_attempts
                    WHERE remote_run_id = %s
                    """,
                    (remote_run_id,),
                )
                owner = cursor.fetchone()
                if owner is not None and int(owner[0]) != int(row[0]):
                    raise ProcessingError(
                        "PROCESSING_RUN_CONFLICT",
                        "remote run is already correlated",
                    )
                if persisted_run_id is not None and persisted_run_id != remote_run_id:
                    raise ProcessingError(
                        "PROCESSING_RUN_CONFLICT",
                        "processing attempt has a different remote run",
                    )
                persisted_run_id = remote_run_id
            if persisted_run_id is None:
                raise ProcessingError(
                    "PROCESSING_RUN_NOT_STARTED",
                    "processing attempt has no remote run",
                )
            if str(row[4]) not in _RUN_TERMINAL:
                cursor.execute(
                    """
                    UPDATE ingestion.processing_attempts
                    SET remote_run_id = %s,
                        attempt_status = %s,
                        failure_code = %s,
                        failure_message = CASE
                            WHEN %s::text IS NULL THEN NULL
                            ELSE 'remote processing did not complete'
                        END,
                        started_at = CASE
                            WHEN %s IN ('dispatched', 'running')
                                THEN COALESCE(started_at, now())
                            ELSE started_at
                        END,
                        finished_at = CASE WHEN %s THEN now() ELSE NULL END
                    WHERE id = %s
                    """,
                    (
                        persisted_run_id,
                        attempt_status,
                        failure_code,
                        failure_code,
                        attempt_status,
                        terminal,
                        row[8],
                    ),
                )
                cursor.execute(
                    """
                    UPDATE ingestion.processing_runs
                    SET run_status = %s,
                        failure_code = %s,
                        failure_message = CASE
                            WHEN %s::text IS NULL THEN NULL
                            ELSE 'remote processing did not complete'
                        END,
                        updated_at = now(),
                        finished_at = CASE WHEN %s THEN now() ELSE NULL END
                    WHERE id = %s
                    """,
                    (run_status, failure_code, failure_code, terminal, row[0]),
                )
            return self._record_from_row(
                row,
                job_id=job_id,
                attempt_no=attempt_no,
                remote_run_id=persisted_run_id,
                remote_status=remote_status,
            )

        return self._run(operation)

    @staticmethod
    def _database_status(remote_status: str) -> Tuple[str, str, Optional[str], bool]:
        if remote_status in {"pending", "queued"}:
            return "dispatched", "queued", None, False
        if remote_status in {"running", "cancelling"}:
            return "running", "running", None, False
        if remote_status in {"success", "completed", "succeeded"}:
            return "delivered", "validating", None, False
        if remote_status in {"failed", "error"}:
            return "failed", "failed", "PROCESSING_REMOTE_FAILED", True
        if remote_status in {"cancelled", "canceled", "interrupted"}:
            return "cancelled", "cancelled", None, True
        raise ProcessingError(
            "PROCESSING_REMOTE_PROTOCOL_INVALID",
            "remote processing status is invalid",
        )

    @staticmethod
    def _remote_status(attempt_status: str) -> str:
        return {
            "queued": "queued",
            "dispatched": "pending",
            "running": "running",
            "delivered": "success",
            "succeeded": "success",
            "failed": "failed",
            "cancelled": "cancelled",
        }.get(attempt_status, "failed")

    @staticmethod
    def _attempt_row(
        cursor: Any,
        job_id: str,
        attempt_no: int,
        *,
        for_update: bool = False,
    ) -> Tuple[Any, ...]:
        cursor.execute(
            """
            SELECT
                run.id,
                run.pipeline_type,
                run.requested_processor_version,
                run.requested_rule_set_sha256,
                run.run_status,
                attempt.idempotency_key,
                attempt.remote_run_id,
                attempt.attempt_status,
                attempt.id,
                attempt.failure_code
            FROM ingestion.processing_runs AS run
            JOIN ingestion.processing_attempts AS attempt
              ON attempt.processing_run_id = run.id
             AND attempt.attempt_no = %s
            WHERE run.run_key = %s
            """
            + (" FOR UPDATE OF run, attempt" if for_update else ""),
            (attempt_no, job_id),
        )
        row = cursor.fetchone()
        if row is None:
            raise ProcessingError(
                "PROCESSING_JOB_NOT_FOUND",
                "processing attempt is not registered",
            )
        return tuple(row)

    @staticmethod
    def _record_from_row(
        row: Tuple[Any, ...],
        *,
        remote_status: str,
        job_id: Optional[str] = None,
        attempt_no: Optional[int] = None,
        remote_run_id: Optional[str] = None,
        result: Optional[ProcessingResult] = None,
        failure_code: Optional[str] = None,
    ) -> ProcessingRecord:
        resolved_job_id = result.job_id if result is not None else job_id
        resolved_attempt_no = result.attempt_no if result is not None else attempt_no
        if resolved_job_id is None or resolved_attempt_no is None:
            raise ProcessingError(
                "PROCESSING_JOB_NOT_FOUND",
                "processing attempt identity is unavailable",
            )
        request = ProcessingRequest(
            job_id=resolved_job_id,
            attempt_no=resolved_attempt_no,
            source_file_id="persisted-source",
            processor_version=str(row[2]),
            rule_set_sha256=str(row[3]),
        )
        return ProcessingRecord(
            request=request,
            request_sha256=request.sha256(),
            idempotency_key=str(row[5]),
            remote_run_id=remote_run_id
            if remote_run_id is not None
            else (str(row[6]) if row[6] is not None else None),
            remote_status=remote_status,
            failure_code=(
                failure_code
                if failure_code is not None
                else (str(row[9]) if row[9] is not None else None)
            ),
            result=result,
        )

    @staticmethod
    def _delivery_matches(cursor: Any, result: ProcessingResult) -> bool:
        cursor.execute(
            """
            SELECT
                run.run_key,
                attempt.attempt_no,
                attempt.remote_run_id,
                delivery.result_status,
                delivery.business_date,
                delivery.processor_version,
                delivery.rule_set_sha256,
                delivery.result_schema_version
            FROM ingestion.processing_deliveries AS delivery
            JOIN ingestion.processing_runs AS run
              ON run.id = delivery.processing_run_id
            LEFT JOIN ingestion.processing_attempts AS attempt
              ON attempt.id = delivery.attempt_id
            WHERE delivery.delivery_key = %s
            """,
            (result.delivery_id,),
        )
        row = cursor.fetchone()
        if row is None:
            return False
        expected_date = result.business_date
        if (
            row[0] != result.job_id
            or int(row[1]) != result.attempt_no
            or row[2] != result.remote_run_id
            or row[3] != result.status
            or row[4] != expected_date
            or row[5] != result.processor_version
            or str(row[6]) != result.rule_set_sha256
            or row[7] != result.result_schema_version
        ):
            raise ProcessingError(
                "PROCESSING_CALLBACK_CONFLICT",
                "delivery identifier conflicts",
            )
        return True