"""Hybrid retrieval for historical anomaly cases."""

from __future__ import annotations

import asyncio
import logging
import re
from dataclasses import dataclass
from typing import Any

import httpx
from sqlalchemy import select
from sqlalchemy import text as sql_text
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import get_settings
from app.models.entities import AnomalyCase
from app.schemas.risk_prediction import SimilarCase
from app.services.supplier_identity import resolve_supplier

logger = logging.getLogger("dieya_risk_prediction.rag")
_embedding_semaphore = asyncio.Semaphore(1)


@dataclass(frozen=True, slots=True)
class RetrievedCase:
    case: SimilarCase
    score: float
    case_id: int


class AnomalyCaseRAG:
    """Retrieve historical cases with business, text, and optional vector signals."""

    def __init__(
        self,
        session: AsyncSession,
        *,
        top_k: int | None = None,
        transport: httpx.AsyncBaseTransport | None = None,
    ) -> None:
        self.session = session
        self.top_k = top_k or get_settings().rag_top_k
        self.transport = transport

    async def search(
        self,
        *,
        query: str,
        category: str | None = None,
        material: str | None = None,
        supplier: str | None = None,
        factory: str | None = None,
        stage: str | None = None,
        anomaly_type: str | None = None,
        rows: list[dict[str, Any]] | None = None,
    ) -> list[RetrievedCase]:
        identity = await resolve_supplier(self.session, supplier)
        supplier_match = (
            identity.short_name or identity.aliases[0]
            if identity
            else supplier
        )
        terms = self._terms(query, category, material, supplier, factory, stage, anomaly_type)
        rows = rows if rows is not None else await self._load_rows(
            query,
            category=category,
            material=material,
            supplier=supplier,
            factory=factory,
            stage=stage,
            anomaly_type=anomaly_type,
        )
        ranked = [
            RetrievedCase(
                case=self._to_case(row),
                score=self._score(
                    row,
                    query=query,
                    category=category,
                    material=material,
                    supplier=supplier_match,
                    factory=factory,
                    stage=stage,
                    anomaly_type=anomaly_type,
                    terms=terms,
                ),
                case_id=int(row["id"]),
            )
            for row in rows
        ]
        ranked = [item for item in ranked if item.score > 0]
        ranked.sort(key=lambda item: (item.score, item.case_id), reverse=True)
        for item in ranked[: self.top_k]:
            vector_similarity = next(
                (
                    float(row["vector_score"])
                    for row in rows
                    if int(row["id"]) == item.case_id and row.get("vector_score") is not None
                ),
                None,
            )
            item.case.similarity = round(
                max(0.0, min(1.0, vector_similarity if vector_similarity is not None else item.score / 100)),
                3,
            )
        return ranked[: self.top_k]

    async def _load_rows(
        self,
        query: str,
        *,
        category: str | None,
        material: str | None,
        supplier: str | None,
        factory: str | None,
        stage: str | None,
        anomaly_type: str | None,
    ) -> list[dict[str, Any]]:
        fields = (
            AnomalyCase.id,
            AnomalyCase.date,
            AnomalyCase.product_name,
            AnomalyCase.product_category,
            AnomalyCase.material_type,
            AnomalyCase.supplier,
            AnomalyCase.factory,
            AnomalyCase.stage_new,
            AnomalyCase.anomaly_type,
            AnomalyCase.description,
            AnomalyCase.solution,
            AnomalyCase.root_cause,
            AnomalyCase.defect_level,
            AnomalyCase.result,
        )
        vector_ready = await self._vector_column_ready()
        has_vectors = await self._has_stored_vectors() if vector_ready else False
        vector = await self._query_embedding(query) if has_vectors else None
        if vector is None:
            return [dict(row) for row in (await self.session.execute(select(*fields).limit(5000))).mappings().all()]

        vector_literal = "[" + ",".join(str(float(value)) for value in vector) + "]"
        statement = sql_text(
            """
            SELECT id, date, product_name, product_category, material_type, supplier,
                   factory, stage_new, anomaly_type, description, solution, root_cause,
                   defect_level, result, 1 - (embedding_vector <=> CAST(:embedding AS halfvec)) AS vector_score
            FROM anomaly_cases
            WHERE embedding_vector IS NOT NULL
            ORDER BY embedding_vector <=> CAST(:embedding AS halfvec)
            LIMIT 5000
            """
        )
        result = await self.session.execute(statement, {"embedding": vector_literal})
        rows = [dict(row) for row in result.mappings().all()]
        if rows:
            return rows
        return [dict(row) for row in (await self.session.execute(select(*fields).limit(5000))).mappings().all()]

    async def _vector_column_ready(self) -> bool:
        try:
            result = await self.session.execute(
                sql_text(
                    """
                    SELECT EXISTS (
                        SELECT 1
                        FROM information_schema.columns
                        WHERE table_name = 'anomaly_cases'
                          AND column_name = 'embedding_vector'
                          AND udt_name IN ('vector', 'halfvec')
                    )
                    """
                )
            )
            return bool(result.scalar())
        except Exception:  # noqa: BLE001 - a missing optional extension uses lexical retrieval
            await self.session.rollback()
            return False

    async def _query_embedding(self, query: str) -> list[float] | None:
        settings = get_settings()
        api_key = (
            settings.rag_embedding_api_key.get_secret_value()
            if settings.rag_embedding_api_key
            else settings.llm_api_key.get_secret_value() if settings.llm_api_key else ""
        )
        base_url = settings.rag_embedding_base_url or settings.llm_base_url
        if not api_key or not base_url:
            return None
        endpoint = f"{base_url.rstrip('/')}/embeddings"
        if not base_url.rstrip("/").endswith("/v1"):
            endpoint = f"{base_url.rstrip('/')}/v1/embeddings"
        timeout = httpx.Timeout(
            settings.rag_embedding_timeout_seconds,
            connect=min(5, settings.rag_embedding_timeout_seconds),
        )
        try:
            async with _embedding_semaphore:
                async with httpx.AsyncClient(timeout=timeout, transport=self.transport) as client:
                    response = await client.post(
                        endpoint,
                        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
                        json={"model": settings.rag_embedding_model, "input": query[:2000]},
                    )
                    response.raise_for_status()
                    body = response.json()
            embedding = (body.get("data") or [{}])[0].get("embedding")
            if not isinstance(embedding, list) or len(embedding) != settings.rag_embedding_dimensions:
                return None
            return [float(value) for value in embedding]
        except (httpx.HTTPError, ValueError, TypeError, KeyError) as exc:
            logger.warning(
                "RAG embedding request failed error=%s; using lexical retrieval",
                exc.__class__.__name__,
            )
            return None

    async def _has_stored_vectors(self) -> bool:
        """Avoid a paid/network embedding call until case vectors exist."""

        try:
            result = await self.session.execute(
                sql_text(
                    "SELECT EXISTS ("
                    "SELECT 1 FROM anomaly_cases WHERE embedding_vector IS NOT NULL LIMIT 1"
                    ")"
                )
            )
            return bool(result.scalar())
        except Exception:  # noqa: BLE001 - lexical retrieval is the safe fallback
            await self.session.rollback()
            return False

    @staticmethod
    def _terms(*values: str | None) -> list[str]:
        terms: list[str] = []
        for value in values:
            if not value:
                continue
            cleaned = " ".join(str(value).strip().lower().split())
            runs = re.findall(r"[a-z0-9]+|[\u4e00-\u9fff]{2,}", cleaned)
            terms.extend(runs)
            for run in runs:
                if re.fullmatch(r"[\u4e00-\u9fff]+", run) and len(run) > 3:
                    terms.extend(
                        run[index : index + size]
                        for size in (2, 3, 4)
                        for index in range(len(run) - size + 1)
                    )
        return list(dict.fromkeys(term for term in terms if len(term) > 1))

    @classmethod
    def _score(
        cls,
        row: dict[str, Any],
        *,
        query: str,
        category: str | None,
        material: str | None,
        supplier: str | None,
        factory: str | None,
        stage: str | None,
        anomaly_type: str | None,
        terms: list[str],
    ) -> float:
        score = 0.0
        score += cls._field_score(row.get("product_category"), category, 7)
        score += cls._field_score(row.get("product_name"), category, 5)
        score += cls._field_score(row.get("material_type"), material, 7)
        score += cls._field_score(row.get("supplier"), supplier, 6)
        score += cls._field_score(row.get("factory"), factory, 6)
        score += cls._field_score(row.get("stage_new"), stage, 3)
        score += cls._field_score(row.get("anomaly_type"), anomaly_type, 7)
        vector_score = row.get("vector_score")
        if vector_score is not None:
            score += max(0.0, float(vector_score)) * 40
        searchable = " ".join(str(row.get(key) or "") for key in (
            "product_name",
            "product_category",
            "material_type",
            "anomaly_type",
            "description",
            "root_cause",
            "solution",
        )).lower()
        query_lower = query.lower().strip()
        if query_lower and query_lower in searchable:
            score += 8
        score += min(12.0, sum(1 for term in terms if term in searchable) * 1.5)
        if str(row.get("defect_level") or "").upper() == "A":
            score += 1.0
        return round(score, 2)

    @staticmethod
    def _field_score(actual: Any, expected: str | None, weight: float) -> float:
        if not actual or not expected:
            return 0.0
        actual_text = str(actual).lower()
        expected_text = str(expected).lower()
        if actual_text == expected_text:
            return weight
        if expected_text in actual_text or actual_text in expected_text:
            return weight * 0.7
        return 0.0

    @staticmethod
    def _to_case(row: dict[str, Any]) -> SimilarCase:
        return SimilarCase(
            case_id=int(row["id"]) if row.get("id") is not None else None,
            date=row.get("date").isoformat() if row.get("date") else None,
            product=row.get("product_name") or row.get("product_category") or "未命名产品",
            category=row.get("product_category"),
            material=row.get("material_type"),
            anomaly_type=row.get("anomaly_type"),
            description=row.get("description"),
            root_cause=row.get("root_cause"),
            defect_level=row.get("defect_level"),
            result=row.get("result"),
            problem=(row.get("anomaly_type") or row.get("description") or "历史异常未填写"),
            solution=(row.get("solution") or row.get("root_cause") or "请查看原始案例记录"),
        )