"""Short-transaction PostgreSQL adapter for lease-fenced document indexing."""

from __future__ import annotations

import math
import re
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from typing import Any, cast

import psycopg
from pgvector.psycopg import register_vector
from pgvector.vector import Vector
from psycopg import Connection
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from app.core.config import Settings
from app.core.secrets import SecretFileError
from app.persistence.job_queue import JobLease, LeaseLostError
from app.persistence.retrieval import ActiveEmbeddingProfile
from app.ports.model_providers import ProviderUsage
from app.services.indexing import (
    ApprovedIndexingPlan,
    AssignmentProgress,
    AssignmentStatus,
    CachedEmbedding,
    EmbeddingCacheLookup,
    EmbeddingWrite,
    IndexingItem,
    InvocationStatus,
    embedding_cache_key,
)

type IndexingRow = dict[str, Any]
type ConnectionFactory = Callable[[str, int], Connection[IndexingRow]]
type VectorRegistrar = Callable[[Connection[Any]], None]

_HASH = re.compile(r"^[0-9a-f]{64}$")
_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/-]{0,511}$")
_ERROR_CODE = re.compile(r"^[A-Z][A-Z0-9_]{0,127}$")


class IndexingPersistenceError(RuntimeError):
    """Sanitized infrastructure failure with no SQL, DSN, text, or credential data."""

    def __init__(self) -> None:
        super().__init__("indexing persistence unavailable")


class IndexingPersistenceConflict(RuntimeError):
    """A governed profile, manifest, invocation, or count contract did not match."""

    def __init__(self) -> None:
        super().__init__("indexing state does not satisfy the governed contract")


LEASE_FENCE_SQL = """
SELECT job.resource_id
FROM rag.background_jobs AS job
WHERE job.id = %s
  AND job.status = 'RUNNING'
  AND job.job_type = 'EMBED_DOCUMENT'
  AND job.required_capability = 'embedding'
  AND job.resource_type = 'document_version'
  AND job.lease_owner = %s
  AND job.lease_token = %s
  AND job.lease_until >= now()
FOR UPDATE
"""

LOAD_PLAN_SQL = """
SELECT
  document.knowledge_base_id,
  version.id AS document_version_id,
  version.review_state,
  version.outbound_manifest_sha256,
  version.expected_chunk_count,
  profile.profile_hash,
  profile.model,
  profile.dimension,
  profile.synthetic,
  (
    SELECT count(*)::integer
    FROM rag.outbound_manifest_items AS item
    WHERE item.document_version_id = version.id
      AND item.outbound_manifest_sha256 = version.outbound_manifest_sha256
  ) AS manifest_count,
  (
    SELECT count(*)::integer
    FROM rag.chunks AS chunk
    JOIN rag.outbound_manifest_items AS item
      ON item.document_version_id = chunk.document_version_id
     AND item.ordinal = chunk.ordinal
     AND item.outbound_manifest_sha256 = chunk.outbound_manifest_sha256
     AND item.cloud_text_sha256 = chunk.cloud_text_sha256
     AND item.embedding_text_sha256 = chunk.embedding_text_sha256
    WHERE chunk.document_version_id = version.id
      AND chunk.approval_status = 'CLOUD_APPROVED'
      AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
      AND chunk.embedding_profile_hash = profile.profile_hash
      AND chunk.embedding_model = profile.model
      AND chunk.embedding_dimension = 1024
      AND chunk.deleted_at IS NULL
  ) AS eligible_chunk_count
FROM rag.document_versions AS version
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
WHERE version.id = %s
  AND version.review_state = 'CLOUD_APPROVED'
  AND version.status IN ('PENDING', 'PROCESSING', 'READY')
  AND version.outbound_manifest_sha256 IS NOT NULL
  AND version.expected_chunk_count IS NOT NULL
  AND knowledge_base.active_embedding_profile_kind = 'embedding'
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
  AND document.status IN (
    'CLOUD_APPROVED', 'EMBEDDING', 'INDEXING', 'VALIDATING', 'READY'
  )
FOR SHARE OF version, document, knowledge_base, profile
"""

LOAD_PLAN_ITEMS_SQL = """
SELECT
  chunk.id AS chunk_id,
  chunk.ordinal,
  chunk.embedding_text,
  chunk.embedding_text_sha256,
  CASE
    WHEN assignment.status = 'READY'
     AND assignment.embedding_text_sha256 = chunk.embedding_text_sha256
     AND assignment.cache_profile_hash = profile.profile_hash
     AND assignment.cache_embedding_text_sha256 = chunk.embedding_text_sha256
     AND cache.resolved_model = profile.model
     AND vector_dims(cache.embedding) = 1024
     AND vector_norm(cache.embedding) > 0
     AND chunk.index_status = 'READY'
     AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
     AND chunk.embedding = cache.embedding
      THEN 'READY'
    WHEN assignment.status = 'READY' THEN 'STALE'
    ELSE COALESCE(assignment.status, 'PENDING')
  END AS assignment_status
FROM rag.chunks AS chunk
JOIN rag.document_versions AS version
  ON version.id = chunk.document_version_id
JOIN rag.documents AS document
  ON document.id = chunk.document_id
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = chunk.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
JOIN rag.outbound_manifest_items AS item
  ON item.document_version_id = chunk.document_version_id
 AND item.ordinal = chunk.ordinal
 AND item.outbound_manifest_sha256 = chunk.outbound_manifest_sha256
 AND item.cloud_text_sha256 = chunk.cloud_text_sha256
 AND item.embedding_text_sha256 = chunk.embedding_text_sha256
LEFT JOIN rag.chunk_embedding_assignments AS assignment
 ON assignment.chunk_id = chunk.id
 AND assignment.profile_hash = profile.profile_hash
LEFT JOIN rag.embedding_cache AS cache
  ON cache.profile_hash = assignment.cache_profile_hash
 AND cache.embedding_text_sha256 = assignment.cache_embedding_text_sha256
WHERE version.id = %s
  AND version.review_state = 'CLOUD_APPROVED'
  AND chunk.approval_status = 'CLOUD_APPROVED'
  AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
  AND chunk.embedding_profile_hash = profile.profile_hash
  AND chunk.embedding_model = profile.model
  AND chunk.embedding_dimension = 1024
  AND chunk.deleted_at IS NULL
  AND document.deleted_at IS NULL
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
  AND knowledge_base.active_embedding_profile_kind = 'embedding'
ORDER BY chunk.ordinal, chunk.id
FOR SHARE OF chunk
"""

CACHE_LOOKUP_SQL = """
SELECT
  cache.profile_hash,
  cache.embedding_text_sha256,
  cache.resolved_model,
  vector_dims(cache.embedding)::integer AS dimension
FROM rag.embedding_cache AS cache
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = cache.profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
JOIN rag.document_versions AS version
  ON version.id = %s
 AND version.review_state = 'CLOUD_APPROVED'
 AND version.embedding_profile_hash = profile.profile_hash
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
 AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
 AND knowledge_base.active_embedding_profile_kind = 'embedding'
WHERE cache.profile_hash = %s
  AND cache.embedding_text_sha256 = %s
  AND cache.resolved_model = profile.model
  AND vector_dims(cache.embedding) = 1024
  AND vector_norm(cache.embedding) > 0
FOR SHARE OF cache
"""

BEGIN_INVOCATION_SQL = """
INSERT INTO rag.model_invocations (
  trace_id,
  caller,
  operation,
  profile_hash,
  model,
  status,
  item_count
)
SELECT
  %s,
  'worker',
  'embedding',
  profile.profile_hash,
  profile.model,
  'STARTED',
  %s
FROM rag.document_versions AS version
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
WHERE version.id = %s
  AND version.review_state = 'CLOUD_APPROVED'
  AND profile.profile_hash = %s
  AND profile.model = %s
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
  AND knowledge_base.active_embedding_profile_kind = 'embedding'
RETURNING id
"""

FINISH_INVOCATION_SQL = """
UPDATE rag.model_invocations AS invocation
SET status = %s,
    provider_request_id = %s,
    prompt_tokens = %s,
    completion_tokens = %s,
    total_tokens = %s,
    elapsed_ms = %s,
    error_code = %s,
    finished_at = now()
FROM rag.document_versions AS version
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
WHERE invocation.id = %s
  AND invocation.trace_id = %s
  AND invocation.operation = 'embedding'
  AND invocation.status = 'STARTED'
  AND invocation.profile_hash = profile.profile_hash
  AND invocation.model = profile.model
  AND version.id = %s
  AND version.review_state = 'CLOUD_APPROVED'
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
RETURNING invocation.id
"""

INSERT_CACHE_SQL = """
INSERT INTO rag.embedding_cache (
  profile_hash,
  embedding_text_sha256,
  embedding,
  resolved_model,
  provider_request_id,
  usage,
  elapsed_ms
)
SELECT %s, %s, %s, %s, %s, %s, %s
FROM rag.document_versions AS version
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
WHERE version.id = %s
  AND version.review_state = 'CLOUD_APPROVED'
  AND version.embedding_profile_hash = %s
  AND profile.model = %s
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
  AND vector_dims(%s::vector) = 1024
  AND vector_norm(%s::vector) > 0
ON CONFLICT (profile_hash, embedding_text_sha256) DO NOTHING
RETURNING profile_hash
"""

VERIFY_CACHE_FOR_WRITE_SQL = """
SELECT 1 AS available
FROM rag.embedding_cache AS cache
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = cache.profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
JOIN rag.document_versions AS version
  ON version.id = %s
 AND version.review_state = 'CLOUD_APPROVED'
 AND version.embedding_profile_hash = profile.profile_hash
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
 AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
WHERE cache.profile_hash = %s
  AND cache.embedding_text_sha256 = %s
  AND cache.resolved_model = %s
  AND vector_dims(cache.embedding) = 1024
  AND vector_norm(cache.embedding) > 0
"""

UPSERT_ASSIGNMENT_SQL = """
INSERT INTO rag.chunk_embedding_assignments (
  chunk_id,
  profile_hash,
  embedding_text_sha256,
  cache_profile_hash,
  cache_embedding_text_sha256,
  status,
  error_code,
  completed_at
)
SELECT
  chunk.id,
  profile.profile_hash,
  chunk.embedding_text_sha256,
  cache.profile_hash,
  cache.embedding_text_sha256,
  'READY',
  NULL,
  now()
FROM rag.chunks AS chunk
JOIN rag.document_versions AS version
  ON version.id = chunk.document_version_id
JOIN rag.documents AS document
  ON document.id = chunk.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = chunk.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
JOIN rag.embedding_cache AS cache
  ON cache.profile_hash = profile.profile_hash
 AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
 AND cache.resolved_model = profile.model
WHERE chunk.id = %s
  AND chunk.document_version_id = %s
  AND chunk.approval_status = 'CLOUD_APPROVED'
  AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
  AND chunk.embedding_profile_hash = profile.profile_hash
  AND chunk.embedding_model = profile.model
  AND chunk.embedding_dimension = 1024
  AND chunk.embedding_text_sha256 = %s
  AND chunk.deleted_at IS NULL
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
  AND profile.profile_hash = %s
  AND vector_dims(cache.embedding) = 1024
  AND vector_norm(cache.embedding) > 0
ON CONFLICT (chunk_id, profile_hash) DO UPDATE
SET cache_profile_hash = EXCLUDED.cache_profile_hash,
    cache_embedding_text_sha256 = EXCLUDED.cache_embedding_text_sha256,
    status = 'READY',
    error_code = NULL,
    completed_at = COALESCE(rag.chunk_embedding_assignments.completed_at, now()),
    updated_at = now()
WHERE rag.chunk_embedding_assignments.embedding_text_sha256 = EXCLUDED.embedding_text_sha256
RETURNING chunk_id
"""

UPDATE_CHUNK_FROM_CACHE_SQL = """
UPDATE rag.chunks AS chunk
SET embedding = cache.embedding,
    embedded_text_sha256 = chunk.embedding_text_sha256,
    embedding_profile_hash = profile.profile_hash,
    embedding_model = profile.model,
    embedding_dimension = 1024,
    index_status = 'READY',
    searchable = false,
    updated_at = now()
FROM rag.document_versions AS version
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
JOIN rag.embedding_cache AS cache
  ON cache.profile_hash = profile.profile_hash
WHERE chunk.id = %s
  AND chunk.document_version_id = version.id
  AND version.id = %s
  AND version.review_state = 'CLOUD_APPROVED'
  AND chunk.approval_status = 'CLOUD_APPROVED'
  AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
  AND chunk.embedding_profile_hash = profile.profile_hash
  AND chunk.embedding_text_sha256 = %s
  AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
  AND cache.resolved_model = profile.model
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
  AND chunk.deleted_at IS NULL
  AND chunk.index_status <> 'READY'
  AND vector_dims(cache.embedding) = 1024
  AND vector_norm(cache.embedding) > 0
RETURNING chunk.id
"""

PREPARE_STALE_CHUNK_SQL = """
UPDATE rag.chunks AS chunk
SET searchable = false,
    index_status = 'EMBEDDING',
    updated_at = now()
FROM rag.embedding_cache AS cache
WHERE chunk.id = %s
  AND chunk.document_version_id = %s
  AND chunk.embedding_profile_hash = %s
  AND chunk.embedding_text_sha256 = %s
  AND cache.profile_hash = chunk.embedding_profile_hash
  AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
  AND cache.resolved_model = chunk.embedding_model
  AND chunk.index_status = 'READY'
  AND (
    chunk.embedded_text_sha256 IS DISTINCT FROM chunk.embedding_text_sha256
    OR chunk.embedding IS DISTINCT FROM cache.embedding
    OR vector_dims(chunk.embedding) IS DISTINCT FROM 1024
    OR vector_norm(chunk.embedding) <= 0
  )
RETURNING chunk.id
"""

VERIFY_READY_CHUNK_SQL = """
SELECT chunk.id
FROM rag.chunks AS chunk
JOIN rag.embedding_cache AS cache
  ON cache.profile_hash = chunk.embedding_profile_hash
 AND cache.embedding_text_sha256 = chunk.embedding_text_sha256
WHERE chunk.id = %s
  AND chunk.document_version_id = %s
  AND chunk.embedding_profile_hash = %s
  AND chunk.embedding_text_sha256 = %s
  AND chunk.embedding_model = cache.resolved_model
  AND chunk.embedding_dimension = 1024
  AND chunk.index_status = 'READY'
  AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
  AND chunk.embedding = cache.embedding
  AND chunk.approval_status = 'CLOUD_APPROVED'
  AND chunk.deleted_at IS NULL
  AND vector_dims(chunk.embedding) = 1024
  AND vector_norm(chunk.embedding) > 0
"""

PROGRESS_SQL = """
SELECT
  version.expected_chunk_count AS expected_count,
  (
    SELECT count(*)::integer
    FROM rag.chunks AS chunk
    WHERE chunk.document_version_id = version.id
      AND chunk.approval_status = 'CLOUD_APPROVED'
      AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
      AND chunk.embedding_profile_hash = profile.profile_hash
      AND chunk.embedding_model = profile.model
      AND chunk.embedding_dimension = 1024
      AND chunk.deleted_at IS NULL
  ) AS chunk_count,
  (
    SELECT count(*)::integer
    FROM rag.chunk_embedding_assignments AS assignment
    JOIN rag.chunks AS chunk ON chunk.id = assignment.chunk_id
    WHERE chunk.document_version_id = version.id
      AND assignment.profile_hash = profile.profile_hash
  ) AS assignment_count,
  (
    SELECT count(*)::integer
    FROM rag.chunks AS chunk
    JOIN rag.chunk_embedding_assignments AS assignment
      ON assignment.chunk_id = chunk.id
     AND assignment.profile_hash = profile.profile_hash
     AND assignment.status = 'READY'
     AND assignment.embedding_text_sha256 = chunk.embedding_text_sha256
     AND assignment.cache_profile_hash = profile.profile_hash
     AND assignment.cache_embedding_text_sha256 = chunk.embedding_text_sha256
    JOIN rag.embedding_cache AS cache
      ON cache.profile_hash = assignment.cache_profile_hash
     AND cache.embedding_text_sha256 = assignment.cache_embedding_text_sha256
     AND cache.resolved_model = profile.model
    WHERE chunk.document_version_id = version.id
      AND chunk.approval_status = 'CLOUD_APPROVED'
      AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
      AND chunk.embedding_profile_hash = profile.profile_hash
      AND chunk.embedding_model = profile.model
      AND chunk.embedding_dimension = 1024
      AND chunk.index_status = 'READY'
      AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
      AND chunk.embedding = cache.embedding
      AND chunk.deleted_at IS NULL
      AND vector_dims(cache.embedding) = 1024
      AND vector_norm(cache.embedding) > 0
  ) AS ready_count
FROM rag.document_versions AS version
JOIN rag.documents AS document
  ON document.id = version.document_id
 AND document.deleted_at IS NULL
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = version.embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
 AND profile.dimension = 1024
WHERE version.id = %s
  AND version.review_state = 'CLOUD_APPROVED'
  AND version.expected_chunk_count IS NOT NULL
  AND profile.profile_hash = %s
  AND knowledge_base.active_embedding_profile_hash = profile.profile_hash
  AND knowledge_base.active_embedding_profile_kind = 'embedding'
FOR SHARE OF version, document, knowledge_base, profile
"""

MARK_VERSION_READY_SQL = """
UPDATE rag.document_versions AS version
SET status = 'READY',
    error_code = NULL,
    completed_at = COALESCE(version.completed_at, now())
FROM rag.documents AS document
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = knowledge_base.active_embedding_profile_hash
 AND profile.kind = 'embedding'
 AND profile.enabled IS TRUE
WHERE version.id = %s
  AND version.document_id = document.id
  AND version.review_state = 'CLOUD_APPROVED'
  AND version.embedding_profile_hash = %s
  AND profile.profile_hash = version.embedding_profile_hash
  AND document.deleted_at IS NULL
RETURNING version.document_id
"""

MARK_DOCUMENT_ACTIVE_SQL = """
UPDATE rag.documents AS document
SET active_version_id = %s,
    status = 'READY',
    updated_at = now()
FROM rag.document_versions AS version
WHERE document.id = version.document_id
  AND version.id = %s
  AND version.status = 'READY'
  AND version.review_state = 'CLOUD_APPROVED'
  AND document.deleted_at IS NULL
RETURNING document.id
"""

DEACTIVATE_OLD_CHUNKS_SQL = """
UPDATE rag.chunks
SET searchable = false,
    updated_at = now()
WHERE document_id = %s
  AND document_version_id <> %s
  AND searchable IS TRUE
"""

ACTIVATE_CURRENT_CHUNKS_SQL = """
UPDATE rag.chunks AS chunk
SET searchable = true,
    updated_at = now()
FROM rag.document_versions AS version
JOIN rag.documents AS document
  ON document.id = version.document_id
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = document.knowledge_base_id
WHERE chunk.document_version_id = version.id
  AND version.id = %s
  AND version.status = 'READY'
  AND version.review_state = 'CLOUD_APPROVED'
  AND version.embedding_profile_hash = %s
  AND document.active_version_id = version.id
  AND document.status = 'READY'
  AND document.deleted_at IS NULL
  AND knowledge_base.active_embedding_profile_hash = version.embedding_profile_hash
  AND chunk.approval_status = 'CLOUD_APPROVED'
  AND chunk.outbound_manifest_sha256 = version.outbound_manifest_sha256
  AND chunk.embedding_profile_hash = version.embedding_profile_hash
  AND chunk.index_status = 'READY'
  AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
  AND chunk.embedding IS NOT NULL
  AND chunk.deleted_at IS NULL
RETURNING chunk.id
"""


def _default_connection_factory(dsn: str, timeout: int) -> Connection[IndexingRow]:
    return psycopg.connect(
        dsn,
        connect_timeout=timeout,
        row_factory=dict_row,
        application_name="geological-rag-indexing-worker",
    )


class PostgresIndexingRepository:
    """Implement every protocol method as one independently committed transaction."""

    def __init__(
        self,
        settings: Settings,
        *,
        connect_timeout: int = 5,
        connection_factory: ConnectionFactory = _default_connection_factory,
        vector_registrar: VectorRegistrar = register_vector,
    ) -> None:
        if isinstance(connect_timeout, bool) or not 1 <= connect_timeout <= 60:
            raise ValueError("connect_timeout must be between 1 and 60")
        self._settings = settings
        self._connect_timeout = connect_timeout
        self._connection_factory = connection_factory
        self._vector_registrar = vector_registrar

    def _dsn(self) -> str:
        return (
            self._settings.database_url()
            .set(drivername="postgresql")
            .render_as_string(hide_password=False)
        )

    @contextmanager
    def _transaction(self) -> Iterator[Connection[IndexingRow]]:
        try:
            with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
                with connection.transaction():
                    self._vector_registrar(cast(Connection[Any], connection))
                    yield connection
        except (LeaseLostError, IndexingPersistenceConflict):
            raise
        except (OSError, SecretFileError, psycopg.Error, KeyError, TypeError, ValueError):
            raise IndexingPersistenceError from None

    @staticmethod
    def _fence(connection: Connection[IndexingRow], lease: JobLease) -> uuid.UUID:
        row = connection.execute(
            LEASE_FENCE_SQL,
            (lease.job_id, lease.worker_id, lease.lease_token),
        ).fetchone()
        if row is None:
            raise LeaseLostError("indexing job lease is no longer owned")
        resource_id = row.get("resource_id")
        if not isinstance(resource_id, uuid.UUID):
            raise IndexingPersistenceConflict
        return resource_id

    def load_approved_plan(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
    ) -> ApprovedIndexingPlan:
        with self._transaction() as connection:
            resource_id = self._fence(connection, lease)
            if resource_id != document_version_id:
                raise IndexingPersistenceConflict
            row = connection.execute(LOAD_PLAN_SQL, (document_version_id,)).fetchone()
            if row is None:
                raise IndexingPersistenceConflict
            item_rows = connection.execute(
                LOAD_PLAN_ITEMS_SQL,
                (document_version_id,),
            ).fetchall()
        return self._plan_from_rows(row, item_rows)

    def lookup_cache(
        self,
        *,
        lease: JobLease,
        lookups: Sequence[EmbeddingCacheLookup],
    ) -> Mapping[str, CachedEmbedding]:
        if not lookups or len(lookups) > 10:
            raise IndexingPersistenceConflict
        with self._transaction() as connection:
            resource_id = self._fence(connection, lease)
            found: dict[str, CachedEmbedding] = {}
            for lookup in lookups:
                if (
                    embedding_cache_key(
                        lookup.embedding_text_sha256,
                        lookup.profile_hash,
                    )
                    != lookup.cache_key
                ):
                    raise IndexingPersistenceConflict
                row = connection.execute(
                    CACHE_LOOKUP_SQL,
                    (resource_id, lookup.profile_hash, lookup.embedding_text_sha256),
                ).fetchone()
                if row is not None:
                    record = self._cache_from_row(lookup.cache_key, row)
                    found[lookup.cache_key] = record
            return found

    def begin_model_invocation(
        self,
        *,
        lease: JobLease,
        trace_id: uuid.UUID,
        profile_hash: str,
        model: str,
        item_count: int,
    ) -> uuid.UUID:
        if (
            trace_id != lease.job_id
            or not _valid_hash(profile_hash)
            or not _valid_model(model)
            or isinstance(item_count, bool)
            or not 1 <= item_count <= 10
        ):
            raise IndexingPersistenceConflict
        with self._transaction() as connection:
            resource_id = self._fence(connection, lease)
            row = connection.execute(
                BEGIN_INVOCATION_SQL,
                (trace_id, item_count, resource_id, profile_hash, model),
            ).fetchone()
            if row is None or not isinstance(row.get("id"), uuid.UUID):
                raise IndexingPersistenceConflict
            return cast(uuid.UUID, row["id"])

    def finish_model_invocation(
        self,
        *,
        lease: JobLease,
        invocation_id: uuid.UUID,
        status: InvocationStatus,
        provider_request_id: str | None,
        usage: ProviderUsage,
        elapsed_ms: float,
        error_code: str | None,
    ) -> None:
        _validate_invocation_finish(
            status=status,
            provider_request_id=provider_request_id,
            elapsed_ms=elapsed_ms,
            error_code=error_code,
        )
        prompt_tokens, completion_tokens, total_tokens = _invocation_counts(usage)
        elapsed = _elapsed_integer(elapsed_ms)
        with self._transaction() as connection:
            resource_id = self._fence(connection, lease)
            row = connection.execute(
                FINISH_INVOCATION_SQL,
                (
                    status,
                    provider_request_id,
                    prompt_tokens,
                    completion_tokens,
                    total_tokens,
                    elapsed,
                    error_code,
                    invocation_id,
                    lease.job_id,
                    resource_id,
                ),
            ).fetchone()
            if row is None:
                raise IndexingPersistenceConflict

    def fenced_persist_batch(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
        profile_hash: str,
        writes: Sequence[EmbeddingWrite],
    ) -> AssignmentProgress:
        if not writes or len(writes) > 10 or not _valid_hash(profile_hash):
            raise IndexingPersistenceConflict
        with self._transaction() as connection:
            resource_id = self._fence(connection, lease)
            if resource_id != document_version_id:
                raise IndexingPersistenceConflict
            for write in writes:
                self._persist_write(
                    connection,
                    document_version_id=document_version_id,
                    profile_hash=profile_hash,
                    write=write,
                )
            return self._read_progress(
                connection,
                document_version_id=document_version_id,
                profile_hash=profile_hash,
            )

    def fenced_activate(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
        profile_hash: str,
        expected_count: int,
    ) -> bool:
        if not _valid_hash(profile_hash) or isinstance(expected_count, bool) or expected_count < 0:
            raise IndexingPersistenceConflict
        with self._transaction() as connection:
            resource_id = self._fence(connection, lease)
            if resource_id != document_version_id:
                raise IndexingPersistenceConflict
            progress, chunk_count, assignment_count = self._read_progress_details(
                connection,
                document_version_id=document_version_id,
                profile_hash=profile_hash,
            )
            if not (
                progress.expected_count
                == expected_count
                == progress.ready_count
                == chunk_count
                == assignment_count
            ):
                return False

            version_row = connection.execute(
                MARK_VERSION_READY_SQL,
                (document_version_id, profile_hash),
            ).fetchone()
            if version_row is None or not isinstance(version_row.get("document_id"), uuid.UUID):
                raise IndexingPersistenceConflict
            document_id = cast(uuid.UUID, version_row["document_id"])
            document_row = connection.execute(
                MARK_DOCUMENT_ACTIVE_SQL,
                (document_version_id, document_version_id),
            ).fetchone()
            if document_row is None:
                raise IndexingPersistenceConflict
            connection.execute(
                DEACTIVATE_OLD_CHUNKS_SQL,
                (document_id, document_version_id),
            )
            activated_rows = connection.execute(
                ACTIVATE_CURRENT_CHUNKS_SQL,
                (document_version_id, profile_hash),
            ).fetchall()
            if len(activated_rows) != expected_count:
                raise IndexingPersistenceConflict
            return True

    @staticmethod
    def _plan_from_rows(
        row: IndexingRow,
        item_rows: Sequence[IndexingRow],
    ) -> ApprovedIndexingPlan:
        expected_count = _required_integer(row, "expected_chunk_count")
        manifest_count = _required_integer(row, "manifest_count")
        eligible_count = _required_integer(row, "eligible_chunk_count")
        if expected_count < 0 or not expected_count == manifest_count == eligible_count == len(
            item_rows
        ):
            raise IndexingPersistenceConflict
        profile = ActiveEmbeddingProfile(
            profile_hash=_required_hash(row, "profile_hash"),
            model=_required_text(row, "model"),
            dimension=_required_integer(row, "dimension"),
            synthetic=_required_boolean(row, "synthetic"),
        )
        if profile.dimension != 1024:
            raise IndexingPersistenceConflict
        items = tuple(PostgresIndexingRepository._item_from_row(item) for item in item_rows)
        return ApprovedIndexingPlan(
            knowledge_base_id=_required_uuid(row, "knowledge_base_id"),
            document_version_id=_required_uuid(row, "document_version_id"),
            review_state=_required_text(row, "review_state"),
            outbound_manifest_sha256=_required_hash(row, "outbound_manifest_sha256"),
            expected_count=expected_count,
            profile=profile,
            items=items,
        )

    @staticmethod
    def _item_from_row(row: IndexingRow) -> IndexingItem:
        status = _required_text(row, "assignment_status")
        if status not in {"PENDING", "EMBEDDING", "READY", "FAILED", "STALE"}:
            raise IndexingPersistenceConflict
        return IndexingItem(
            chunk_id=_required_uuid(row, "chunk_id"),
            ordinal=_required_integer(row, "ordinal"),
            embedding_text=_required_text(row, "embedding_text", strip=False),
            embedding_text_sha256=_required_hash(row, "embedding_text_sha256"),
            assignment_status=cast(AssignmentStatus, status),
        )

    @staticmethod
    def _cache_from_row(cache_key: str, row: IndexingRow) -> CachedEmbedding:
        dimension = _required_integer(row, "dimension")
        if dimension != 1024:
            raise IndexingPersistenceConflict
        return CachedEmbedding(
            cache_key=cache_key,
            profile_hash=_required_hash(row, "profile_hash"),
            embedding_text_sha256=_required_hash(row, "embedding_text_sha256"),
            resolved_model=_required_text(row, "resolved_model"),
            dimension=dimension,
        )

    @staticmethod
    def _persist_write(
        connection: Connection[IndexingRow],
        *,
        document_version_id: uuid.UUID,
        profile_hash: str,
        write: EmbeddingWrite,
    ) -> None:
        _validate_write(write, profile_hash)
        if write.source == "provider":
            vector = Vector(list(cast(tuple[float, ...], write.embedding)))
            usage = Jsonb(_usage_json(write.usage))
            elapsed = _elapsed_integer(write.elapsed_ms)
            connection.execute(
                INSERT_CACHE_SQL,
                (
                    profile_hash,
                    write.embedding_text_sha256,
                    vector,
                    write.resolved_model,
                    write.provider_request_id,
                    usage,
                    elapsed,
                    document_version_id,
                    profile_hash,
                    write.resolved_model,
                    vector,
                    vector,
                ),
            )
        cache_row = connection.execute(
            VERIFY_CACHE_FOR_WRITE_SQL,
            (
                document_version_id,
                profile_hash,
                write.embedding_text_sha256,
                write.resolved_model,
            ),
        ).fetchone()
        if cache_row is None:
            raise IndexingPersistenceConflict
        assignment_row = connection.execute(
            UPSERT_ASSIGNMENT_SQL,
            (
                write.chunk_id,
                document_version_id,
                write.embedding_text_sha256,
                profile_hash,
            ),
        ).fetchone()
        if assignment_row is None:
            raise IndexingPersistenceConflict
        connection.execute(
            PREPARE_STALE_CHUNK_SQL,
            (
                write.chunk_id,
                document_version_id,
                profile_hash,
                write.embedding_text_sha256,
            ),
        )
        chunk_row = connection.execute(
            UPDATE_CHUNK_FROM_CACHE_SQL,
            (
                write.chunk_id,
                document_version_id,
                write.embedding_text_sha256,
            ),
        ).fetchone()
        if chunk_row is None:
            chunk_row = connection.execute(
                VERIFY_READY_CHUNK_SQL,
                (
                    write.chunk_id,
                    document_version_id,
                    profile_hash,
                    write.embedding_text_sha256,
                ),
            ).fetchone()
        if chunk_row is None:
            raise IndexingPersistenceConflict

    @staticmethod
    def _read_progress(
        connection: Connection[IndexingRow],
        *,
        document_version_id: uuid.UUID,
        profile_hash: str,
    ) -> AssignmentProgress:
        progress, _, _ = PostgresIndexingRepository._read_progress_details(
            connection,
            document_version_id=document_version_id,
            profile_hash=profile_hash,
        )
        return progress

    @staticmethod
    def _read_progress_details(
        connection: Connection[IndexingRow],
        *,
        document_version_id: uuid.UUID,
        profile_hash: str,
    ) -> tuple[AssignmentProgress, int, int]:
        row = connection.execute(
            PROGRESS_SQL,
            (document_version_id, profile_hash),
        ).fetchone()
        if row is None:
            raise IndexingPersistenceConflict
        expected_count = _required_integer(row, "expected_count")
        chunk_count = _required_integer(row, "chunk_count")
        assignment_count = _required_integer(row, "assignment_count")
        ready_count = _required_integer(row, "ready_count")
        if not 0 <= ready_count <= assignment_count <= chunk_count <= expected_count:
            raise IndexingPersistenceConflict
        return AssignmentProgress(expected_count, ready_count), chunk_count, assignment_count


def _validate_write(write: EmbeddingWrite, profile_hash: str) -> None:
    if (
        write.profile_hash != profile_hash
        or not _valid_hash(write.embedding_text_sha256)
        or embedding_cache_key(write.embedding_text_sha256, profile_hash) != write.cache_key
        or not _valid_model(write.resolved_model)
        or isinstance(write.batch_index, bool)
        or not 0 <= write.batch_index < 10
        or isinstance(write.elapsed_ms, bool)
        or not isinstance(write.elapsed_ms, (int, float))
        or not math.isfinite(float(write.elapsed_ms))
        or write.elapsed_ms < 0
        or (
            write.provider_request_id is not None
            and not _valid_identifier(write.provider_request_id)
        )
    ):
        raise IndexingPersistenceConflict
    _invocation_counts(write.usage)
    if write.source == "cache":
        if write.embedding is not None:
            raise IndexingPersistenceConflict
        return
    if write.source != "provider" or write.embedding is None:
        raise IndexingPersistenceConflict
    if len(write.embedding) != 1024:
        raise IndexingPersistenceConflict
    norm = 0.0
    for component in write.embedding:
        if (
            isinstance(component, bool)
            or not isinstance(component, (int, float))
            or not math.isfinite(float(component))
        ):
            raise IndexingPersistenceConflict
        norm += float(component) ** 2
    if norm <= 0:
        raise IndexingPersistenceConflict


def _validate_invocation_finish(
    *,
    status: InvocationStatus,
    provider_request_id: str | None,
    elapsed_ms: float,
    error_code: str | None,
) -> None:
    if status not in {"SUCCEEDED", "FAILED", "UNKNOWN"}:
        raise IndexingPersistenceConflict
    if (status == "SUCCEEDED") != (error_code is None):
        raise IndexingPersistenceConflict
    if error_code is not None and _ERROR_CODE.fullmatch(error_code) is None:
        raise IndexingPersistenceConflict
    if provider_request_id is not None and not _valid_identifier(provider_request_id):
        raise IndexingPersistenceConflict
    _elapsed_integer(elapsed_ms)


def _invocation_counts(usage: ProviderUsage) -> tuple[int, int, int]:
    values = (usage.input_tokens, usage.output_tokens, usage.total_tokens)
    if any(
        isinstance(value, bool) or (value is not None and (not isinstance(value, int) or value < 0))
        for value in values
    ):
        raise IndexingPersistenceConflict
    prompt = usage.input_tokens
    completion = usage.output_tokens or 0
    if prompt is None:
        prompt = max(0, (usage.total_tokens or 0) - completion)
    total = prompt + completion
    if usage.total_tokens is not None and usage.total_tokens != total:
        raise IndexingPersistenceConflict
    return prompt, completion, total


def _usage_json(usage: ProviderUsage) -> dict[str, int]:
    prompt, completion, total = _invocation_counts(usage)
    return {
        "input_tokens": prompt,
        "output_tokens": completion,
        "total_tokens": total,
    }


def _elapsed_integer(value: float) -> int:
    if (
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(float(value))
        or not 0 <= value <= 2_147_483_647
    ):
        raise IndexingPersistenceConflict
    return int(round(value))


def _valid_hash(value: str) -> bool:
    return isinstance(value, str) and _HASH.fullmatch(value) is not None


def _valid_model(value: str) -> bool:
    return isinstance(value, str) and value == value.strip() and _valid_identifier(value)


def _valid_identifier(value: str) -> bool:
    return isinstance(value, str) and _SAFE_IDENTIFIER.fullmatch(value) is not None


def _required_uuid(row: Mapping[str, object], key: str) -> uuid.UUID:
    value = row.get(key)
    if not isinstance(value, uuid.UUID):
        raise IndexingPersistenceConflict
    return value


def _required_text(row: Mapping[str, object], key: str, *, strip: bool = True) -> str:
    value = row.get(key)
    if not isinstance(value, str) or not value or (strip and value != value.strip()):
        raise IndexingPersistenceConflict
    return value


def _required_hash(row: Mapping[str, object], key: str) -> str:
    value = _required_text(row, key)
    if not _valid_hash(value):
        raise IndexingPersistenceConflict
    return value


def _required_integer(row: Mapping[str, object], key: str) -> int:
    value = row.get(key)
    if isinstance(value, bool) or not isinstance(value, int):
        raise IndexingPersistenceConflict
    return value


def _required_boolean(row: Mapping[str, object], key: str) -> bool:
    value = row.get(key)
    if not isinstance(value, bool):
        raise IndexingPersistenceConflict
    return value