"""Formal two-stage retrieval use case with server-owned authorization scope."""

from __future__ import annotations

import asyncio
import math
import re
import time
import uuid
from dataclasses import dataclass
from typing import Literal

from app.core.problems import ApiProblem
from app.persistence.retrieval import (
    ActiveEmbeddingProfile,
    RetrievalCandidate,
    RetrievalPersistenceError,
    RetrievalRepository,
)
from app.ports.model_providers import (
    EmbeddingProvider,
    ModelProviderError,
    RankedItem,
    Reranker,
    RerankResult,
)

QUERY_MAX_LENGTH = 500
VECTOR_TOP_K_DEFAULT = 50
VECTOR_TOP_K_MAX = 50
RERANK_TOP_N_DEFAULT = 10
RERANK_TOP_N_MAX = 10
RERANK_TEXT_MAX_BYTES = 4_000
RERANK_REQUEST_MAX_BYTES = 120_000
SNIPPET_MAX_LENGTH = 1_200
SOURCE_NAME_MAX_LENGTH = 240
RERANK_INSTRUCT = (
    "Given a geological exploration question, rank passages that directly support "
    "an evidence-grounded answer."
)
_SPACE_PATTERN = re.compile(r"\s+")


@dataclass(frozen=True, slots=True)
class RetrievalGrant:
    """A server-resolved knowledge-base grant; never built from request scope fields."""

    knowledge_base_id: uuid.UUID
    access_scope_ids: tuple[uuid.UUID, ...]


@dataclass(frozen=True, slots=True)
class RetrievalActor:
    """Authenticated identity projection consumed by the retrieval service."""

    subject: str
    grants: tuple[RetrievalGrant, ...]

    def scopes_for(self, knowledge_base_id: uuid.UUID) -> tuple[uuid.UUID, ...]:
        scopes: list[uuid.UUID] = []
        for grant in self.grants:
            if grant.knowledge_base_id == knowledge_base_id:
                scopes.extend(grant.access_scope_ids)
        return tuple(dict.fromkeys(scopes))


@dataclass(frozen=True, slots=True)
class EffectiveRetrievalParameters:
    vector_top_k: int
    rerank_top_n: int


@dataclass(frozen=True, slots=True)
class RetrievalTimings:
    embedding_ms: float
    database_ms: float
    rerank_ms: float
    total_ms: float


@dataclass(frozen=True, slots=True)
class RetrievalHit:
    rank: int
    vector_rank: int
    citation_id: uuid.UUID
    document_id: uuid.UUID
    source_name: str
    snippet: str
    section_path: tuple[str, ...]
    page_start: int | None
    page_end: int | None
    page_label: str
    vector_score: float
    rerank_score: float | None


@dataclass(frozen=True, slots=True)
class RetrievalResult:
    status: Literal["ok", "empty"]
    knowledge_base_id: uuid.UUID
    access_scope_count: int
    profile: ActiveEmbeddingProfile
    parameters: EffectiveRetrievalParameters
    rerank_status: Literal["applied", "degraded", "skipped_empty"]
    degradation_reason: Literal["rerank_unavailable"] | None
    embedding_request_id: str | None
    rerank_request_id: str | None
    embedding_model: str
    rerank_model: str | None
    timings: RetrievalTimings
    results: tuple[RetrievalHit, ...]


class RetrievalService:
    """Coordinate query embedding, authorized vector search, and bounded reranking."""

    def __init__(
        self,
        *,
        repository: RetrievalRepository,
        embedding_provider: EmbeddingProvider,
        reranker: Reranker,
        synthetic_embedding_provider: EmbeddingProvider | None = None,
        synthetic_reranker: Reranker | None = None,
    ) -> None:
        self._repository = repository
        self._embedding_provider = embedding_provider
        self._reranker = reranker
        self._synthetic_embedding_provider = synthetic_embedding_provider
        self._synthetic_reranker = synthetic_reranker

    async def search(
        self,
        *,
        actor: RetrievalActor,
        knowledge_base_id: uuid.UUID,
        query: str,
        vector_top_k: int = VECTOR_TOP_K_DEFAULT,
        rerank_top_n: int = RERANK_TOP_N_DEFAULT,
    ) -> RetrievalResult:
        started = time.perf_counter()
        normalized_query = _normalize_query(query)
        parameters = _effective_parameters(vector_top_k, rerank_top_n)
        allowed_scope_ids = actor.scopes_for(knowledge_base_id)
        if not allowed_scope_ids:
            raise ApiProblem(
                status=403,
                code="RETRIEVAL_SCOPE_FORBIDDEN",
                title="Knowledge base access denied",
                detail="The current identity cannot search this knowledge base.",
            )

        database_started = time.perf_counter()
        try:
            profile = await asyncio.to_thread(
                self._repository.resolve_active_profile,
                knowledge_base_id,
                allowed_scope_ids=allowed_scope_ids,
            )
        except RetrievalPersistenceError as exc:
            raise _storage_unavailable() from exc
        profile_database_ms = (time.perf_counter() - database_started) * 1_000
        if profile is None:
            raise ApiProblem(
                status=409,
                code="KNOWLEDGE_BASE_NOT_SEARCHABLE",
                title="Knowledge base is not searchable",
                detail="No enabled active embedding profile is available for this knowledge base.",
            )

        embedding_provider, reranker = self._providers(profile)
        try:
            embedding = await embedding_provider.embed_query(normalized_query)
        except ModelProviderError as exc:
            raise _embedding_problem(exc) from exc
        query_vector = _validated_query_vector(embedding.vectors, profile)
        if embedding.model != profile.model:
            raise ApiProblem(
                status=502,
                code="EMBEDDING_PROFILE_MISMATCH",
                title="Embedding response did not match the active profile",
                detail="The query embedding model did not match the knowledge base profile.",
            )

        search_started = time.perf_counter()
        try:
            candidates = await asyncio.to_thread(
                self._repository.search_candidates,
                knowledge_base_id,
                allowed_scope_ids=allowed_scope_ids,
                profile_hash=profile.profile_hash,
                query_vector=query_vector,
                limit=parameters.vector_top_k,
            )
        except RetrievalPersistenceError as exc:
            raise _storage_unavailable() from exc
        database_ms = profile_database_ms + (time.perf_counter() - search_started) * 1_000

        candidates = _unique_candidates(candidates)
        if not candidates:
            total_ms = (time.perf_counter() - started) * 1_000
            return RetrievalResult(
                status="empty",
                knowledge_base_id=knowledge_base_id,
                access_scope_count=len(allowed_scope_ids),
                profile=profile,
                parameters=parameters,
                rerank_status="skipped_empty",
                degradation_reason=None,
                embedding_request_id=embedding.request_id,
                rerank_request_id=None,
                embedding_model=embedding.model,
                rerank_model=None,
                timings=RetrievalTimings(
                    embedding_ms=_safe_elapsed(embedding.elapsed_ms),
                    database_ms=max(0.0, database_ms),
                    rerank_ms=0.0,
                    total_ms=max(0.0, total_ms),
                ),
                results=(),
            )

        effective_top_n = min(parameters.rerank_top_n, len(candidates))
        documents = _bounded_rerank_documents(normalized_query, candidates)
        rerank_result: RerankResult | None = None
        try:
            attempted = await reranker.rerank(
                normalized_query,
                documents,
                top_n=effective_top_n,
                instruct=RERANK_INSTRUCT,
            )
            if _valid_rerank(attempted, documents, effective_top_n):
                rerank_result = attempted
        except ModelProviderError:
            pass

        selected: tuple[tuple[int, float | None], ...]
        if rerank_result is None:
            selected = tuple((index, None) for index in range(effective_top_n))
            rerank_status: Literal["applied", "degraded"] = "degraded"
            degradation_reason: Literal["rerank_unavailable"] | None = "rerank_unavailable"
            rerank_request_id = None
            rerank_model = None
            rerank_ms = 0.0
        else:
            selected = tuple((item.index, item.relevance_score) for item in rerank_result.items)
            rerank_status = "applied"
            degradation_reason = None
            rerank_request_id = rerank_result.request_id
            rerank_model = rerank_result.model
            rerank_ms = _safe_elapsed(rerank_result.elapsed_ms)

        hits = tuple(
            _hit(
                candidate=candidates[candidate_index],
                rank=rank,
                vector_rank=candidate_index + 1,
                rerank_score=rerank_score,
            )
            for rank, (candidate_index, rerank_score) in enumerate(selected, start=1)
        )
        total_ms = (time.perf_counter() - started) * 1_000
        return RetrievalResult(
            status="ok",
            knowledge_base_id=knowledge_base_id,
            access_scope_count=len(allowed_scope_ids),
            profile=profile,
            parameters=parameters,
            rerank_status=rerank_status,
            degradation_reason=degradation_reason,
            embedding_request_id=embedding.request_id,
            rerank_request_id=rerank_request_id,
            embedding_model=embedding.model,
            rerank_model=rerank_model,
            timings=RetrievalTimings(
                embedding_ms=_safe_elapsed(embedding.elapsed_ms),
                database_ms=max(0.0, database_ms),
                rerank_ms=rerank_ms,
                total_ms=max(0.0, total_ms),
            ),
            results=hits,
        )

    def _providers(
        self,
        profile: ActiveEmbeddingProfile,
    ) -> tuple[EmbeddingProvider, Reranker]:
        if not profile.synthetic:
            return self._embedding_provider, self._reranker
        if self._synthetic_embedding_provider is None or self._synthetic_reranker is None:
            raise ApiProblem(
                status=503,
                code="SYNTHETIC_PROVIDER_UNAVAILABLE",
                title="Synthetic retrieval provider unavailable",
                detail="The active synthetic profile has no matching local provider.",
            )
        return self._synthetic_embedding_provider, self._synthetic_reranker


def _normalize_query(value: str) -> str:
    if not isinstance(value, str):
        raise ApiProblem(
            status=400,
            code="INVALID_RETRIEVAL_QUERY",
            title="Invalid retrieval query",
            detail="The query must be non-empty text.",
        )
    normalized = _SPACE_PATTERN.sub(" ", value).strip()
    if not normalized or len(normalized) > QUERY_MAX_LENGTH:
        raise ApiProblem(
            status=400,
            code="INVALID_RETRIEVAL_QUERY",
            title="Invalid retrieval query",
            detail=f"The query must contain between 1 and {QUERY_MAX_LENGTH} characters.",
        )
    return normalized


def _effective_parameters(vector_top_k: int, rerank_top_n: int) -> EffectiveRetrievalParameters:
    for value in (vector_top_k, rerank_top_n):
        if isinstance(value, bool) or not isinstance(value, int) or value < 1:
            raise ApiProblem(
                status=400,
                code="INVALID_RETRIEVAL_PARAMETERS",
                title="Invalid retrieval parameters",
                detail="Retrieval limits must be positive integers.",
            )
    bounded_vector_top_k = min(vector_top_k, VECTOR_TOP_K_MAX)
    bounded_rerank_top_n = min(rerank_top_n, RERANK_TOP_N_MAX, bounded_vector_top_k)
    return EffectiveRetrievalParameters(
        vector_top_k=bounded_vector_top_k,
        rerank_top_n=bounded_rerank_top_n,
    )


def _validated_query_vector(
    vectors: tuple[tuple[float, ...], ...],
    profile: ActiveEmbeddingProfile,
) -> tuple[float, ...]:
    if len(vectors) != 1 or len(vectors[0]) != profile.dimension:
        raise ApiProblem(
            status=502,
            code="INVALID_EMBEDDING_RESPONSE",
            title="Invalid embedding response",
            detail="The embedding provider returned an unexpected vector shape.",
        )
    vector = vectors[0]
    if any(
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(float(value))
        for value in vector
    ):
        raise ApiProblem(
            status=502,
            code="INVALID_EMBEDDING_RESPONSE",
            title="Invalid embedding response",
            detail="The embedding provider returned an invalid vector.",
        )
    normalized = tuple(float(value) for value in vector)
    if math.hypot(*normalized) <= 0:
        raise ApiProblem(
            status=502,
            code="INVALID_EMBEDDING_RESPONSE",
            title="Invalid embedding response",
            detail="The embedding provider returned a zero vector.",
        )
    return normalized


def _unique_candidates(candidates: list[RetrievalCandidate]) -> list[RetrievalCandidate]:
    seen: set[uuid.UUID] = set()
    unique: list[RetrievalCandidate] = []
    for candidate in candidates:
        if candidate.citation_id not in seen:
            seen.add(candidate.citation_id)
            unique.append(candidate)
    return unique


def _bounded_rerank_documents(
    query: str,
    candidates: list[RetrievalCandidate],
) -> tuple[str, ...]:
    query_bytes = len(query.encode("utf-8"))
    available = RERANK_REQUEST_MAX_BYTES - query_bytes * len(candidates)
    per_document = min(RERANK_TEXT_MAX_BYTES, available // len(candidates))
    if per_document < 1:
        # The public query and candidate limits make this unreachable. Keep a
        # fail-closed guard so future limit changes cannot exceed provider bounds.
        raise ApiProblem(
            status=400,
            code="RERANK_BUDGET_EXCEEDED",
            title="Rerank request is too large",
            detail="The effective retrieval request exceeds the rerank input budget.",
        )
    return tuple(_truncate_utf8(_rerank_text(candidate), per_document) for candidate in candidates)


def _rerank_text(candidate: RetrievalCandidate) -> str:
    section = " > ".join(candidate.section_path) if candidate.section_path else "章节未知"
    page = _page_label(candidate.page_start, candidate.page_end)
    return f"章节：{section}\n页码：{page}\n{candidate.cloud_text}"


def _truncate_utf8(value: str, maximum_bytes: int) -> str:
    encoded = value.encode("utf-8")
    if len(encoded) <= maximum_bytes:
        return value
    truncated = encoded[:maximum_bytes].decode("utf-8", errors="ignore").rstrip()
    return truncated or value[0]


def _valid_rerank(
    result: RerankResult,
    documents: tuple[str, ...],
    expected: int,
) -> bool:
    if len(result.items) != expected or not result.model.strip():
        return False
    seen: set[int] = set()
    previous = math.inf
    for item in result.items:
        if not _valid_ranked_item(item, documents, seen, previous):
            return False
        seen.add(item.index)
        previous = item.relevance_score
    return True


def _valid_ranked_item(
    item: RankedItem,
    documents: tuple[str, ...],
    seen: set[int],
    previous: float,
) -> bool:
    return (
        not isinstance(item.index, bool)
        and isinstance(item.index, int)
        and 0 <= item.index < len(documents)
        and item.index not in seen
        and item.document == documents[item.index]
        and isinstance(item.relevance_score, (int, float))
        and not isinstance(item.relevance_score, bool)
        and math.isfinite(float(item.relevance_score))
        and 0.0 <= item.relevance_score <= 1.0
        and item.relevance_score <= previous
    )


def _hit(
    *,
    candidate: RetrievalCandidate,
    rank: int,
    vector_rank: int,
    rerank_score: float | None,
) -> RetrievalHit:
    return RetrievalHit(
        rank=rank,
        vector_rank=vector_rank,
        citation_id=candidate.citation_id,
        document_id=candidate.document_id,
        source_name=_bounded_text(candidate.source_name, SOURCE_NAME_MAX_LENGTH),
        snippet=_bounded_text(candidate.cloud_text, SNIPPET_MAX_LENGTH),
        section_path=candidate.section_path,
        page_start=candidate.page_start,
        page_end=candidate.page_end,
        page_label=_page_label(candidate.page_start, candidate.page_end),
        vector_score=round(max(-1.0, min(1.0, candidate.vector_score)), 6),
        rerank_score=round(rerank_score, 6) if rerank_score is not None else None,
    )


def _bounded_text(value: str, maximum: int) -> str:
    normalized = _SPACE_PATTERN.sub(" ", value).strip()
    if len(normalized) <= maximum:
        return normalized
    return f"{normalized[: maximum - 1]}…"


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


def _page_label(page_start: int | None, page_end: int | None) -> str:
    if page_start is None or page_end is None:
        return "页码未知"
    if page_start == page_end:
        return f"第 {page_start} 页"
    return f"第 {page_start}-{page_end} 页"


def _storage_unavailable() -> ApiProblem:
    return ApiProblem(
        status=503,
        code="RETRIEVAL_STORAGE_UNAVAILABLE",
        title="Retrieval storage unavailable",
        detail="The retrieval index is temporarily unavailable.",
    )


def _embedding_problem(exc: ModelProviderError) -> ApiProblem:
    if exc.kind.value == "invalid_response":
        return ApiProblem(
            status=502,
            code="INVALID_EMBEDDING_RESPONSE",
            title="Invalid embedding response",
            detail="The embedding provider returned an invalid response.",
        )
    return ApiProblem(
        status=503,
        code="EMBEDDING_UNAVAILABLE",
        title="Embedding service unavailable",
        detail="The query embedding service is temporarily unavailable.",
    )