"""Lease-fenced document embedding orchestration.

This module deliberately owns no database transaction.  Every repository call
is a complete short operation, while provider I/O happens only after that call
has returned.  Persistence implementations must atomically validate the full
``JobLease`` on every mutation.
"""

from __future__ import annotations

import asyncio
import hashlib
import math
import re
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Literal, Protocol

from app.persistence.job_queue import JobLease
from app.persistence.retrieval import ActiveEmbeddingProfile
from app.ports.model_providers import (
    EmbeddingProvider,
    EmbeddingResult,
    ModelProviderError,
    ProviderErrorKind,
    ProviderUsage,
)

EMBEDDING_BATCH_SIZE = 10
EMBEDDING_DIMENSION = 1024
_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")

type AssignmentStatus = Literal["PENDING", "EMBEDDING", "READY", "FAILED", "STALE"]
type InvocationStatus = Literal["SUCCEEDED", "FAILED", "UNKNOWN"]
type WriteSource = Literal["cache", "provider"]


class IndexingError(RuntimeError):
    """Base class for safe indexing failures."""


class InvalidIndexingPlanError(IndexingError):
    """Approved indexing inputs violated their immutable contract."""


class InvalidEmbeddingResponseError(IndexingError):
    """The provider returned an unusable or profile-incompatible embedding."""


class IndexingNotReadyError(IndexingError):
    """Activation was refused because not every expected assignment is READY."""


class IndexingProviderError(IndexingError):
    """An unexpected provider failure was converted to a safe worker error."""


@dataclass(frozen=True, slots=True)
class IndexingItem:
    chunk_id: uuid.UUID
    ordinal: int
    embedding_text: str
    embedding_text_sha256: str
    assignment_status: AssignmentStatus


@dataclass(frozen=True, slots=True)
class ApprovedIndexingPlan:
    knowledge_base_id: uuid.UUID
    document_version_id: uuid.UUID
    review_state: str
    outbound_manifest_sha256: str
    expected_count: int
    profile: ActiveEmbeddingProfile
    items: tuple[IndexingItem, ...]


@dataclass(frozen=True, slots=True)
class CachedEmbedding:
    cache_key: str
    profile_hash: str
    embedding_text_sha256: str
    resolved_model: str
    dimension: int


@dataclass(frozen=True, slots=True)
class EmbeddingCacheLookup:
    """Derived key plus the indexed database key needed for an efficient lookup."""

    cache_key: str
    profile_hash: str
    embedding_text_sha256: str


@dataclass(frozen=True, slots=True)
class EmbeddingWrite:
    chunk_id: uuid.UUID
    batch_index: int
    cache_key: str
    profile_hash: str
    embedding_text_sha256: str
    source: WriteSource
    embedding: tuple[float, ...] | None
    resolved_model: str
    provider_request_id: str | None
    usage: ProviderUsage
    elapsed_ms: float


@dataclass(frozen=True, slots=True)
class AssignmentProgress:
    expected_count: int
    ready_count: int


@dataclass(frozen=True, slots=True)
class IndexingResult:
    document_version_id: uuid.UUID
    profile_hash: str
    expected_count: int
    ready_count: int
    cache_hit_count: int
    newly_embedded_count: int
    provider_call_count: int
    activated: bool


class IndexingRepository(Protocol):
    """Short-operation persistence port; implementations must not retain transactions."""

    def load_approved_plan(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
    ) -> ApprovedIndexingPlan: ...

    def lookup_cache(
        self,
        *,
        lease: JobLease,
        lookups: Sequence[EmbeddingCacheLookup],
    ) -> Mapping[str, CachedEmbedding]: ...

    def begin_model_invocation(
        self,
        *,
        lease: JobLease,
        trace_id: uuid.UUID,
        profile_hash: str,
        model: str,
        item_count: int,
    ) -> uuid.UUID: ...

    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: ...

    def fenced_persist_batch(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
        profile_hash: str,
        writes: Sequence[EmbeddingWrite],
    ) -> AssignmentProgress: ...

    def fenced_activate(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
        profile_hash: str,
        expected_count: int,
    ) -> bool: ...


class DocumentIndexingService:
    """Resolve cache hits, embed misses, and activate only a complete projection."""

    def __init__(
        self,
        *,
        repository: IndexingRepository,
        embedding_provider: EmbeddingProvider,
        synthetic_embedding_provider: EmbeddingProvider | None = None,
    ) -> None:
        self._repository = repository
        self._embedding_provider = embedding_provider
        self._synthetic_embedding_provider = synthetic_embedding_provider

    async def index_document_version(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
        trace_id: uuid.UUID,
    ) -> IndexingResult:
        plan = await asyncio.to_thread(
            self._repository.load_approved_plan,
            lease=lease,
            document_version_id=document_version_id,
        )
        _validate_plan(plan, document_version_id)
        provider = self._provider_for(plan.profile)

        initial_ready = sum(item.assignment_status == "READY" for item in plan.items)
        progress = AssignmentProgress(plan.expected_count, initial_ready)
        pending = tuple(item for item in plan.items if item.assignment_status != "READY")
        items_by_cache_key = _group_by_cache_key(pending, plan.profile.profile_hash)

        cached_by_key: dict[str, CachedEmbedding] = {}
        cache_keys = tuple(items_by_cache_key)
        for key_batch in _batches(cache_keys, EMBEDDING_BATCH_SIZE):
            lookups = tuple(
                EmbeddingCacheLookup(
                    cache_key=key,
                    profile_hash=plan.profile.profile_hash,
                    embedding_text_sha256=items_by_cache_key[key][0].embedding_text_sha256,
                )
                for key in key_batch
            )
            found = await asyncio.to_thread(
                self._repository.lookup_cache,
                lease=lease,
                lookups=lookups,
            )
            cached_by_key.update(_validated_cache(found, lookups, plan.profile))

        cache_writes = tuple(
            _cache_write(item, cached_by_key[key], plan.profile)
            for key, items in items_by_cache_key.items()
            if key in cached_by_key
            for item in items
        )
        for write_batch in _batches(cache_writes, EMBEDDING_BATCH_SIZE):
            progress = await self._persist(
                lease=lease,
                plan=plan,
                writes=write_batch,
                previous=progress,
            )

        missing_keys = tuple(key for key in items_by_cache_key if key not in cached_by_key)
        provider_call_count = 0
        newly_embedded_count = 0
        for key_batch in _batches(missing_keys, EMBEDDING_BATCH_SIZE):
            batch_items = tuple(items_by_cache_key[key][0] for key in key_batch)
            invocation_id = await asyncio.to_thread(
                self._repository.begin_model_invocation,
                lease=lease,
                trace_id=trace_id,
                profile_hash=plan.profile.profile_hash,
                model=plan.profile.model,
                item_count=len(batch_items),
            )
            provider_call_count += 1
            result, validated = await self._call_provider(
                provider=provider,
                items=batch_items,
                profile=plan.profile,
                lease=lease,
                invocation_id=invocation_id,
            )
            newly_embedded_count += len(validated)

            provider_writes = tuple(
                _provider_write(
                    item,
                    cache_key=cache_key,
                    batch_index=validated_index,
                    vector=validated[validated_index],
                    result=result,
                    profile=plan.profile,
                )
                for validated_index, cache_key in enumerate(key_batch)
                for item in items_by_cache_key[cache_key]
            )
            for write_batch in _batches(provider_writes, EMBEDDING_BATCH_SIZE):
                progress = await self._persist(
                    lease=lease,
                    plan=plan,
                    writes=write_batch,
                    previous=progress,
                )

        if (
            progress.expected_count != plan.expected_count
            or progress.ready_count != plan.expected_count
        ):
            raise IndexingNotReadyError("not every expected embedding assignment is READY")

        activated = await asyncio.to_thread(
            self._repository.fenced_activate,
            lease=lease,
            document_version_id=plan.document_version_id,
            profile_hash=plan.profile.profile_hash,
            expected_count=plan.expected_count,
        )
        if not activated:
            raise IndexingNotReadyError("fenced activation rejected an incomplete projection")
        return IndexingResult(
            document_version_id=plan.document_version_id,
            profile_hash=plan.profile.profile_hash,
            expected_count=plan.expected_count,
            ready_count=progress.ready_count,
            cache_hit_count=len(cache_writes),
            newly_embedded_count=newly_embedded_count,
            provider_call_count=provider_call_count,
            activated=True,
        )

    def _provider_for(self, profile: ActiveEmbeddingProfile) -> EmbeddingProvider:
        if not profile.synthetic:
            return self._embedding_provider
        if self._synthetic_embedding_provider is None:
            raise InvalidIndexingPlanError("synthetic profile has no local embedding provider")
        return self._synthetic_embedding_provider

    async def _call_provider(
        self,
        *,
        provider: EmbeddingProvider,
        items: tuple[IndexingItem, ...],
        profile: ActiveEmbeddingProfile,
        lease: JobLease,
        invocation_id: uuid.UUID,
    ) -> tuple[EmbeddingResult, tuple[tuple[float, ...], ...]]:
        try:
            result = await provider.embed_documents(tuple(item.embedding_text for item in items))
        except ModelProviderError as exc:
            await asyncio.to_thread(
                self._repository.finish_model_invocation,
                lease=lease,
                invocation_id=invocation_id,
                status=_provider_failure_status(exc.kind),
                provider_request_id=_safe_request_id(exc.request_id),
                usage=ProviderUsage(),
                elapsed_ms=0.0,
                error_code=f"EMBEDDING_{exc.kind.value.upper()}",
            )
            raise
        except Exception:
            await asyncio.to_thread(
                self._repository.finish_model_invocation,
                lease=lease,
                invocation_id=invocation_id,
                status="UNKNOWN",
                provider_request_id=None,
                usage=ProviderUsage(),
                elapsed_ms=0.0,
                error_code="EMBEDDING_PROVIDER_UNEXPECTED",
            )
            raise IndexingProviderError("embedding provider failed unexpectedly") from None

        try:
            validated = _validated_result(result, items, profile)
        except InvalidEmbeddingResponseError:
            await asyncio.to_thread(
                self._repository.finish_model_invocation,
                lease=lease,
                invocation_id=invocation_id,
                status="FAILED",
                provider_request_id=_safe_request_id(result.request_id),
                usage=_safe_usage(result.usage),
                elapsed_ms=_safe_elapsed(result.elapsed_ms),
                error_code="INVALID_EMBEDDING_RESPONSE",
            )
            raise

        await asyncio.to_thread(
            self._repository.finish_model_invocation,
            lease=lease,
            invocation_id=invocation_id,
            status="SUCCEEDED",
            provider_request_id=result.request_id,
            usage=result.usage,
            elapsed_ms=result.elapsed_ms,
            error_code=None,
        )
        return result, validated

    async def _persist(
        self,
        *,
        lease: JobLease,
        plan: ApprovedIndexingPlan,
        writes: tuple[EmbeddingWrite, ...],
        previous: AssignmentProgress,
    ) -> AssignmentProgress:
        if not writes or len(writes) > EMBEDDING_BATCH_SIZE:
            raise InvalidIndexingPlanError(
                "persistence batches must contain between 1 and 10 items"
            )
        progress = await asyncio.to_thread(
            self._repository.fenced_persist_batch,
            lease=lease,
            document_version_id=plan.document_version_id,
            profile_hash=plan.profile.profile_hash,
            writes=writes,
        )
        if (
            progress.expected_count != plan.expected_count
            or not previous.ready_count <= progress.ready_count <= plan.expected_count
        ):
            raise IndexingNotReadyError("assignment progress violated the expected READY count")
        return progress


def embedding_cache_key(embedding_text_sha256: str, profile_hash: str) -> str:
    """Return the deterministic application cache key required by the indexing contract."""

    if not _valid_sha256(embedding_text_sha256) or not _valid_sha256(profile_hash):
        raise InvalidIndexingPlanError("cache key inputs must be lowercase SHA-256 values")
    return hashlib.sha256(f"{embedding_text_sha256}{profile_hash}".encode()).hexdigest()


def _validate_plan(plan: ApprovedIndexingPlan, requested_version_id: uuid.UUID) -> None:
    profile = plan.profile
    if plan.document_version_id != requested_version_id:
        raise InvalidIndexingPlanError("repository returned a different document version")
    if plan.review_state != "CLOUD_APPROVED":
        raise InvalidIndexingPlanError("document version is not cloud approved")
    if not _valid_sha256(plan.outbound_manifest_sha256):
        raise InvalidIndexingPlanError("approved manifest hash is invalid")
    if (
        not _valid_sha256(profile.profile_hash)
        or not profile.model.strip()
        or profile.dimension != EMBEDDING_DIMENSION
    ):
        raise InvalidIndexingPlanError("embedding profile is invalid or unsupported")
    if (
        isinstance(plan.expected_count, bool)
        or plan.expected_count < 0
        or plan.expected_count != len(plan.items)
    ):
        raise InvalidIndexingPlanError("expected chunk count does not match the approved plan")

    chunk_ids: set[uuid.UUID] = set()
    ordinals: set[int] = set()
    valid_statuses = {"PENDING", "EMBEDDING", "READY", "FAILED", "STALE"}
    for item in plan.items:
        if item.chunk_id in chunk_ids or item.ordinal in ordinals:
            raise InvalidIndexingPlanError("approved indexing items must be unique")
        if isinstance(item.ordinal, bool) or item.ordinal < 0:
            raise InvalidIndexingPlanError("chunk ordinal is invalid")
        if not item.embedding_text:
            raise InvalidIndexingPlanError("approved embedding text must not be empty")
        if item.assignment_status not in valid_statuses:
            raise InvalidIndexingPlanError("embedding assignment status is invalid")
        actual_text_hash = hashlib.sha256(item.embedding_text.encode()).hexdigest()
        if item.embedding_text_sha256 != actual_text_hash:
            raise InvalidIndexingPlanError("approved embedding text hash does not match its text")
        chunk_ids.add(item.chunk_id)
        ordinals.add(item.ordinal)
    if ordinals != set(range(plan.expected_count)):
        raise InvalidIndexingPlanError("approved chunk ordinals must be contiguous")


def _group_by_cache_key(
    items: tuple[IndexingItem, ...],
    profile_hash: str,
) -> dict[str, list[IndexingItem]]:
    grouped: dict[str, list[IndexingItem]] = {}
    for item in items:
        key = embedding_cache_key(item.embedding_text_sha256, profile_hash)
        grouped.setdefault(key, []).append(item)
    return grouped


def _validated_cache(
    found: Mapping[str, CachedEmbedding],
    lookups: tuple[EmbeddingCacheLookup, ...],
    profile: ActiveEmbeddingProfile,
) -> dict[str, CachedEmbedding]:
    requested = {lookup.cache_key: lookup for lookup in lookups}
    if any(key not in requested for key in found):
        raise InvalidIndexingPlanError("cache lookup returned an unrequested key")
    validated: dict[str, CachedEmbedding] = {}
    for key, record in found.items():
        expected_key = embedding_cache_key(record.embedding_text_sha256, record.profile_hash)
        if (
            key != record.cache_key
            or key != expected_key
            or record.profile_hash != profile.profile_hash
            or record.profile_hash != requested[key].profile_hash
            or record.embedding_text_sha256 != requested[key].embedding_text_sha256
            or record.resolved_model != profile.model
            or record.dimension != profile.dimension
        ):
            raise InvalidIndexingPlanError("cache record does not match the active profile")
        validated[key] = record
    return validated


def _validated_result(
    result: EmbeddingResult,
    items: tuple[IndexingItem, ...],
    profile: ActiveEmbeddingProfile,
) -> tuple[tuple[float, ...], ...]:
    if result.model != profile.model:
        raise InvalidEmbeddingResponseError("embedding model did not match the active profile")
    if len(result.vectors) != len(items):
        raise InvalidEmbeddingResponseError("embedding result count did not match the batch")
    if (
        isinstance(result.elapsed_ms, bool)
        or not isinstance(result.elapsed_ms, (int, float))
        or not math.isfinite(float(result.elapsed_ms))
        or result.elapsed_ms < 0
    ):
        raise InvalidEmbeddingResponseError("embedding elapsed time is invalid")
    if result.request_id is not None and _safe_request_id(result.request_id) is None:
        raise InvalidEmbeddingResponseError("embedding request identifier is invalid")
    if _safe_usage(result.usage) != result.usage:
        raise InvalidEmbeddingResponseError("embedding usage metadata is invalid")

    vectors: list[tuple[float, ...]] = []
    for index, vector in enumerate(result.vectors):
        if index >= len(items):
            raise InvalidEmbeddingResponseError("embedding index exceeded the requested batch")
        if len(vector) != profile.dimension:
            raise InvalidEmbeddingResponseError("embedding dimension did not match the profile")
        normalized: list[float] = []
        for component in vector:
            if (
                isinstance(component, bool)
                or not isinstance(component, (int, float))
                or not math.isfinite(float(component))
            ):
                raise InvalidEmbeddingResponseError("embedding contains a non-finite component")
            normalized.append(float(component))
        if math.hypot(*normalized) <= 0:
            raise InvalidEmbeddingResponseError("embedding vector must have a nonzero norm")
        vectors.append(tuple(normalized))
    return tuple(vectors)


def _cache_write(
    item: IndexingItem,
    cached: CachedEmbedding,
    profile: ActiveEmbeddingProfile,
) -> EmbeddingWrite:
    return EmbeddingWrite(
        chunk_id=item.chunk_id,
        batch_index=0,
        cache_key=cached.cache_key,
        profile_hash=profile.profile_hash,
        embedding_text_sha256=item.embedding_text_sha256,
        source="cache",
        embedding=None,
        resolved_model=cached.resolved_model,
        provider_request_id=None,
        usage=ProviderUsage(),
        elapsed_ms=0.0,
    )


def _provider_write(
    item: IndexingItem,
    *,
    cache_key: str,
    batch_index: int,
    vector: tuple[float, ...],
    result: EmbeddingResult,
    profile: ActiveEmbeddingProfile,
) -> EmbeddingWrite:
    return EmbeddingWrite(
        chunk_id=item.chunk_id,
        batch_index=batch_index,
        cache_key=cache_key,
        profile_hash=profile.profile_hash,
        embedding_text_sha256=item.embedding_text_sha256,
        source="provider",
        embedding=vector,
        resolved_model=result.model,
        provider_request_id=result.request_id,
        usage=result.usage,
        elapsed_ms=result.elapsed_ms,
    )


def _provider_failure_status(kind: ProviderErrorKind) -> InvocationStatus:
    if kind in {ProviderErrorKind.TIMEOUT, ProviderErrorKind.TRANSPORT}:
        return "UNKNOWN"
    return "FAILED"


def _safe_request_id(value: str | None) -> str | None:
    if value is None:
        return None
    if not value.strip() or len(value) > 512 or any(character.isspace() for character in value):
        return None
    return value


def _safe_usage(value: ProviderUsage) -> ProviderUsage:
    def valid(component: int | None) -> int | None:
        if component is not None and (
            isinstance(component, bool) or not isinstance(component, int) or component < 0
        ):
            return None
        return component

    return ProviderUsage(
        input_tokens=valid(value.input_tokens),
        output_tokens=valid(value.output_tokens),
        total_tokens=valid(value.total_tokens),
    )


def _safe_elapsed(value: float) -> float:
    if (
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(float(value))
        or value < 0
    ):
        return 0.0
    return float(value)


def _valid_sha256(value: str) -> bool:
    return bool(_SHA256_PATTERN.fullmatch(value))


def _batches[T](values: Sequence[T], size: int) -> tuple[tuple[T, ...], ...]:
    return tuple(tuple(values[index : index + size]) for index in range(0, len(values), size))