"""Retrieval Service"""
import asyncio
import hashlib
import json
from datetime import datetime, timedelta, timezone
from typing import List, Dict, Any, Tuple, Set, Optional, Literal
import uuid
from uuid import uuid4
from sqlalchemy import select, or_
from sqlalchemy.orm import selectinload
from sqlalchemy.ext.asyncio import AsyncSession
from normbase.models.norm import Norm, NormTag, NormUsageToken
from normbase.services.qdrant_service import QdrantService
from normbase.services.cache_service import CacheService
from normbase.config import settings
from normbase.infra.llm import get_llm_provider


def _utc_now() -> datetime:
    return datetime.now(timezone.utc).replace(tzinfo=None)


class RetrievalService:
    """Norm retrieval service with hybrid search"""
    
    def __init__(self, llm_provider=None, qdrant_url: str = None):
        if llm_provider is None:
            llm_provider = get_llm_provider()
        
        self.llm = llm_provider
        self.qdrant = QdrantService(url=qdrant_url or settings.qdrant_url)
        self.cache = CacheService()
    
    async def retrieve(
        self,
        db: AsyncSession,
        context: Dict[str, Any],
        query_mode: Literal["structured", "semantic", "hybrid"] = "hybrid",
        nl_query: Optional[str] = None,
        agent_role: Optional[str] = None,
        agent_ids: List[uuid.UUID] | None = None,
        tag_filters: Optional[Dict[str, List[str]]] = None,
        field_ids: List[uuid.UUID] = None,
        top_k: int = 5
    ) -> Tuple[List[Tuple[Norm, float]], bool, Dict[str, Any]]:
        """
        Retrieve applicable norms using hybrid search
        
        Args:
            db: Database session
            context: Task context
            agent_role: Optional agent role
            tag_filters: Optional tag filters by category
            field_ids: Information field IDs to search in
            top_k: Number of results
        
        Returns:
            (List of (Norm, score) tuples, cache_hit, meta)
        """
        effective_tag_filters = self._merge_tag_filters(agent_role, tag_filters)
        cache_key = self._build_cache_key(
            context=context,
            query_mode=query_mode,
            nl_query=nl_query,
            agent_role=agent_role,
            agent_ids=agent_ids,
            tag_filters=effective_tag_filters,
            field_ids=field_ids,
            top_k=top_k,
        )
        if settings.retrieve_cache_enabled:
            try:
                cached_payload = await self.cache.get_json(cache_key)
                if cached_payload is not None:
                    cached_results = await self._hydrate_cached_results(db, cached_payload)
                    if cached_results:
                        cached_query_text = self._resolve_semantic_query(
                            query_mode=query_mode,
                            nl_query=nl_query,
                            context=context,
                            agent_role=agent_role,
                        )
                        cached_semantic_attempted = bool(cached_query_text)
                        return (
                            cached_results[:top_k],
                            True,
                            {
                                "query_mode_effective": query_mode,
                                "semantic_used": False,
                                "semantic_attempted": cached_semantic_attempted,
                                "semantic_query_source": self._semantic_query_source(
                                    query_mode=query_mode,
                                    nl_query=nl_query,
                                    context=context,
                                ),
                                "sql_candidates": len(cached_results),
                                "vector_hits": 0,
                                "semantic_hits": 0,
                                "fusion_strategy": "cache",
                            },
                        )
            except Exception as e:
                print(f"Retrieve cache read error: {e}")

        context_terms = self._extract_context_terms(context or {})

        # Use field_ids for SQL search
        if field_ids and len(field_ids) > 0:
            sql_results = await self._sql_search_by_field_ids(
                db=db,
                tag_filters=effective_tag_filters,
                field_ids=field_ids,
                context_terms=context_terms,
                agent_ids=agent_ids,
            )
        else:
            sql_results = await self._sql_search_all_active(
                db=db,
                tag_filters=effective_tag_filters,
                context_terms=context_terms,
                agent_ids=agent_ids,
            )

        semantic_query_text = self._resolve_semantic_query(
            query_mode=query_mode,
            nl_query=nl_query,
            context=context,
            agent_role=agent_role,
        )
        semantic_query_source = self._semantic_query_source(
            query_mode=query_mode,
            nl_query=nl_query,
            context=context,
        )
        semantic_attempted = bool(semantic_query_text)
        vector_hits = 0
        fusion_strategy = "structured_only"

        if query_mode == "structured":
            combined = sql_results[:top_k]
        elif query_mode == "semantic":
            if semantic_query_text:
                vector_results = await self._vector_search(
                    query_text=semantic_query_text,
                    top_k=top_k * 2,
                )
            else:
                vector_results = []
            vector_hits = len(vector_results)
            combined = self._semantic_rank(sql_results, vector_results, top_k)
            fusion_strategy = "semantic_rank"
        else:
            if semantic_query_text:
                vector_results = await self._vector_search(
                    query_text=semantic_query_text,
                    top_k=top_k * 2,
                )
            else:
                vector_results = []
            vector_hits = len(vector_results)
            combined = self._rrf_fusion(sql_results, vector_results, top_k)
            fusion_strategy = "rrf"

        if settings.retrieve_cache_enabled:
            try:
                cache_payload = [
                    {"norm_id": str(norm.id), "score": score} for norm, score in combined
                ]
                await self.cache.set_json(
                    cache_key,
                    cache_payload,
                    ttl_seconds=settings.retrieve_cache_ttl_seconds,
                )
            except Exception as e:
                print(f"Retrieve cache write error: {e}")

        return (
            combined,
            False,
            {
                "query_mode_effective": query_mode,
                "semantic_used": semantic_attempted and vector_hits > 0,
                "semantic_attempted": semantic_attempted,
                "semantic_query_source": semantic_query_source,
                "sql_candidates": len(sql_results),
                "vector_hits": vector_hits,
                "semantic_hits": vector_hits,
                "fusion_strategy": fusion_strategy,
            },
        )

    async def retrieve_with_usage_token(
        self,
        db: AsyncSession,
        context: Dict[str, Any],
        query_mode: Literal["structured", "semantic", "hybrid"] = "hybrid",
        nl_query: Optional[str] = None,
        agent_role: Optional[str] = None,
        agent_ids: List[uuid.UUID] | None = None,
        tag_filters: Optional[Dict[str, List[str]]] = None,
        field_ids: List[uuid.UUID] = None,
        top_k: int = 5,
        agent_id: Optional[str] = None,
        key_id: Optional[str] = None,
    ) -> Tuple[List[Tuple[Norm, float]], str, bool, Dict[str, Any]]:
        """
        Retrieve norms and issue a bound usage token in one transaction path.

        Returns:
            (results, usage_token, cache_hit, meta)
        """
        results, cache_hit, meta = await self.retrieve(
            db=db,
            context=context,
            query_mode=query_mode,
            nl_query=nl_query,
            agent_role=agent_role,
            agent_ids=agent_ids,
            tag_filters=tag_filters,
            field_ids=field_ids,
            top_k=top_k,
        )

        usage_token = f"ut-{uuid4()}"
        token_row = NormUsageToken(
            usage_token=usage_token,
            norm_ids=[str(norm.id) for norm, _ in results],
            reported_norm_ids=[],
            agent_id=(agent_id or "").strip() or None,
            key_id=(key_id or "").strip() or None,
            expires_at=_utc_now() + timedelta(hours=24),
        )
        db.add(token_row)
        await db.commit()
        return results, usage_token, cache_hit, meta
    
    async def _sql_search_all_active(
        self,
        db: AsyncSession,
        tag_filters: Optional[Dict[str, List[str]]] = None,
        context_terms: Set[str] | None = None,
        agent_ids: List[uuid.UUID] | None = None,
    ) -> List[Tuple[Norm, float]]:
        """SQL-based search for all active norms with optional tag filters."""
        query = select(Norm).options(selectinload(Norm.tags)).where(Norm.status == "active")

        query = self._apply_tag_filters(query, tag_filters)
        query = self._apply_agent_id_filters(query, agent_ids)
        
        result = await db.execute(query)
        norms = result.scalars().all()
        
        terms = context_terms or set()
        return [(norm, self._score_norm_with_context(norm, terms)) for norm in norms]
    
    async def _sql_search_by_field_ids(
        self,
        db: AsyncSession,
        tag_filters: Optional[Dict[str, List[str]]],
        field_ids: List[uuid.UUID],
        context_terms: Set[str] | None = None,
        agent_ids: List[uuid.UUID] | None = None,
    ) -> List[Tuple[Norm, float]]:
        """SQL-based exact matching for multiple field IDs."""
        query = select(Norm).options(selectinload(Norm.tags)).where(Norm.status == "active")

        query = self._apply_tag_filters(query, tag_filters)
        query = self._apply_agent_id_filters(query, agent_ids)
        
        if field_ids:
            query = query.where(Norm.field_id.in_(field_ids))
        
        result = await db.execute(query)
        norms = result.scalars().all()
        
        terms = context_terms or set()
        return [(norm, self._score_norm_with_context(norm, terms)) for norm in norms]
    
    def _extract_context_terms(self, context: Dict[str, Any]) -> Set[str]:
        """Flatten context into lowercase lexical terms for lightweight SQL-side scoring."""
        terms: Set[str] = set()

        def _walk(value: Any) -> None:
            if isinstance(value, dict):
                for k, v in value.items():
                    terms.add(str(k).strip().lower())
                    _walk(v)
                return
            if isinstance(value, list):
                for item in value:
                    _walk(item)
                return
            if value is None:
                return
            terms.add(str(value).strip().lower())

        _walk(context or {})
        return {term for term in terms if term}

    def _score_norm_with_context(self, norm: Norm, context_terms: Set[str]) -> float:
        """
        Assign lightweight SQL-channel score from tags and norm.context keyword overlap.
        Keeps deterministic ranking even before vector fusion.
        """
        score = 1.0

        if not context_terms:
            return score

        # Boost on business tag categories matched by context terms.
        tags = getattr(norm, "tags", []) or []
        matched_tag_terms = 0
        for tag in tags:
            if tag.tag_category in {"scenario", "domain", "tool", "skill", "capability"}:
                value = str(tag.tag).strip().lower()
                if value and value in context_terms:
                    matched_tag_terms += 1
        if matched_tag_terms:
            score += 0.2 * matched_tag_terms

        # Boost on flattened context payload overlap.
        norm_terms = self._extract_context_terms(norm.context or {})
        overlap = len(norm_terms.intersection(context_terms))
        if overlap:
            score += min(0.5, 0.05 * overlap)

        return score
    
    async def _vector_search(
        self,
        query_text: str,
        top_k: int
    ) -> List[Tuple[str, float]]:
        """Vector-based semantic search using Qdrant"""
        try:
            # Get embedding from LLM provider (async)
            if not self.llm:
                # No provider means no semantic channel.
                return []

            # Bound semantic channel latency to avoid blocking the whole request.
            query_vector = await asyncio.wait_for(
                self.llm.async_get_embedding(query_text),
                timeout=12.0,
            )

            # Search Qdrant (sync client call), keep a timeout guard too.
            results = await asyncio.wait_for(
                asyncio.to_thread(
                    self.qdrant.search,
                    query_vector=query_vector,
                    limit=top_k,
                ),
                timeout=8.0,
            )
            return [(r["id"], r["score"]) for r in results]
        except asyncio.TimeoutError:
            print("Vector search timeout, fallback to SQL channel.")
            return []
        except Exception as e:
            print(f"Vector search error: {e}")
            return []
    
    def _build_context_text(self, context: Dict[str, Any], agent_role: str) -> str:
        """Build text for embedding"""
        parts: List[str] = []
        if agent_role:
            parts.append(f"Agent role: {agent_role}")
        
        if context:
            for key, value in context.items():
                parts.append(f"{key}: {value}")

        if not parts:
            return "retrieve norms"
        return " | ".join(parts)

    def _resolve_semantic_query(
        self,
        query_mode: Literal["structured", "semantic", "hybrid"],
        nl_query: Optional[str],
        context: Dict[str, Any],
        agent_role: Optional[str],
    ) -> str:
        """Resolve semantic query text based on query mode."""
        if query_mode == "structured":
            return ""
        text = str(nl_query or "").strip()
        if text:
            return text
        if query_mode == "semantic":
            return ""
        if self._should_use_vector_search(context):
            return self._build_context_text(context, agent_role or "")
        return ""

    def _semantic_query_source(
        self,
        query_mode: Literal["structured", "semantic", "hybrid"],
        nl_query: Optional[str],
        context: Dict[str, Any],
    ) -> str:
        if query_mode == "structured":
            return "none"
        if str(nl_query or "").strip():
            return "nl_query"
        if query_mode == "hybrid" and self._should_use_vector_search(context):
            return "context"
        return "none"

    def _should_use_vector_search(self, context: Dict[str, Any]) -> bool:
        """Skip semantic channel for purely structured/boolean contexts."""
        if not context:
            return False
        text_like_keys = {"query", "text", "message", "description", "content"}
        for key, value in context.items():
            key_lower = str(key).lower()
            if key_lower in text_like_keys and isinstance(value, str) and value.strip():
                return True
            if isinstance(value, str) and len(value.strip()) >= 12:
                return True
            if isinstance(value, dict) and self._should_use_vector_search(value):
                return True
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, str) and len(item.strip()) >= 12:
                        return True
                    if isinstance(item, dict) and self._should_use_vector_search(item):
                        return True
        return False
    
    def _rrf_fusion(
        self,
        sql_results: List[Tuple[Norm, float]],
        vector_results: List[Tuple[str, float]],
        top_k: int,
        k: int = 60
    ) -> List[Tuple[Norm, float]]:
        """
        Reciprocal Rank Fusion for combining search results
        
        Args:
            sql_results: SQL search results (Norm, rank)
            vector_results: Vector search results (norm_id, rank)
            top_k: Number of results to return
            k: RRF parameter
        
        Returns:
            Fused and ranked results
        """
        rrf_scores: Dict[str, float] = {}
        norm_map: Dict[str, Norm] = {}
        
        # RRF for SQL results
        for rank, (norm, _) in enumerate(sql_results, 1):
            norm_id = str(norm.id)
            rrf_scores[norm_id] = rrf_scores.get(norm_id, 0) + 1.0 / (k + rank)
            norm_map[norm_id] = norm
        
        # RRF for vector results
        for rank, (norm_id, score) in enumerate(vector_results, 1):
            rrf_scores[norm_id] = rrf_scores.get(norm_id, 0) + 1.0 / (k + rank)
        
        # Sort by RRF score
        sorted_ids = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
        
        # Return top-k with scores
        results = []
        for norm_id, score in sorted_ids[:top_k]:
            if norm_id in norm_map:
                results.append((norm_map[norm_id], score))
        
        return results

    def _semantic_rank(
        self,
        sql_results: List[Tuple[Norm, float]],
        vector_results: List[Tuple[str, float]],
        top_k: int,
    ) -> List[Tuple[Norm, float]]:
        """Rank SQL-filtered norms by semantic channel order only."""
        if not sql_results or not vector_results:
            return []
        sql_norm_map = {str(norm.id): norm for norm, _ in sql_results}
        ranked: List[Tuple[Norm, float]] = []
        for norm_id, score in vector_results:
            norm = sql_norm_map.get(str(norm_id))
            if norm is None:
                continue
            ranked.append((norm, float(score)))
            if len(ranked) >= top_k:
                break
        return ranked

    async def _hydrate_cached_results(
        self, db: AsyncSession, payload: List[Dict[str, Any]]
    ) -> List[Tuple[Norm, float]]:
        """Convert cached ids/scores back to ORM objects in ranked order."""
        if not payload:
            return []
        ids = []
        for item in payload:
            norm_id = item.get("norm_id")
            if not norm_id:
                continue
            try:
                ids.append(uuid.UUID(norm_id))
            except ValueError:
                continue
        if not ids:
            return []
        rows = (
            await db.execute(
                select(Norm)
                .options(selectinload(Norm.tags))
                .where(Norm.id.in_(ids), Norm.status == "active")
            )
        ).scalars().all()
        norm_map = {str(norm.id): norm for norm in rows}
        results: List[Tuple[Norm, float]] = []
        for item in payload:
            norm_id = item.get("norm_id")
            if norm_id in norm_map:
                results.append((norm_map[norm_id], float(item.get("score", 0.0))))
        return results

    def _build_cache_key(
        self,
        context: Dict[str, Any],
        query_mode: Literal["structured", "semantic", "hybrid"],
        nl_query: Optional[str],
        agent_role: Optional[str],
        agent_ids: List[uuid.UUID] | None,
        tag_filters: Optional[Dict[str, List[str]]],
        field_ids: List[uuid.UUID] | None,
        top_k: int,
    ) -> str:
        normalized = {
            "query_mode": query_mode,
            "nl_query": str(nl_query or "").strip(),
            "agent_role": agent_role,
            "agent_ids": self._normalize_agent_ids(agent_ids),
            "tag_filters": self._normalize_tag_filters(tag_filters),
            "context": self._normalize_context(context or {}),
            "field_ids": sorted(str(fid) for fid in (field_ids or [])),
            "top_k": top_k,
        }
        digest_source = json.dumps(
            normalized, ensure_ascii=False, sort_keys=True, separators=(",", ":")
        )
        digest = hashlib.sha256(digest_source.encode("utf-8")).hexdigest()
        return f"{settings.retrieve_cache_prefix}{digest}"

    def _normalize_context(self, value: Any) -> Any:
        if isinstance(value, dict):
            return {k: self._normalize_context(value[k]) for k in sorted(value)}
        if isinstance(value, list):
            return [self._normalize_context(item) for item in value]
        return value

    def _normalize_agent_ids(
        self, agent_ids: List[uuid.UUID] | None
    ) -> List[str]:
        if not isinstance(agent_ids, list):
            return []
        normalized = {
            str(item).strip()
            for item in agent_ids
            if str(item).strip()
        }
        return sorted(normalized)

    def _normalize_tag_filters(
        self, tag_filters: Optional[Dict[str, List[str]]]
    ) -> Dict[str, List[str]]:
        if not isinstance(tag_filters, dict):
            return {}
        normalized: Dict[str, List[str]] = {}
        for key, values in tag_filters.items():
            category = str(key).strip().lower()
            if not category:
                continue
            if not isinstance(values, list):
                continue
            cleaned = []
            for item in values:
                text = str(item).strip()
                if text and text not in cleaned:
                    cleaned.append(text)
            if cleaned:
                normalized[category] = sorted(cleaned)
        return normalized

    def _merge_tag_filters(
        self,
        agent_role: Optional[str],
        tag_filters: Optional[Dict[str, List[str]]],
    ) -> Dict[str, List[str]]:
        merged = self._normalize_tag_filters(tag_filters)
        role = (agent_role or "").strip()
        if role:
            current = merged.get("role", [])
            if role not in current:
                merged["role"] = [*current, role]
        return merged

    def _apply_tag_filters(
        self,
        query,
        tag_filters: Optional[Dict[str, List[str]]],
    ):
        normalized = self._normalize_tag_filters(tag_filters)
        for category, values in normalized.items():
            effective_values = list(values)
            if category == "role" and "*" not in effective_values:
                # role='*' means global role rule (applies to everyone)
                effective_values.append("*")
            tag_exists = (
                select(NormTag.id)
                .where(
                    NormTag.norm_id == Norm.id,
                    NormTag.tag_category == category,
                    NormTag.tag.in_(effective_values),
                )
                .exists()
            )
            query = query.where(tag_exists)
        return query

    def _apply_agent_id_filters(
        self,
        query,
        agent_ids: List[uuid.UUID] | None,
    ):
        normalized_ids = self._normalize_agent_ids(agent_ids)
        if not normalized_ids:
            return query
        conditions = [
            Norm.subject.contains({"canonical": {"agents": [agent_id]}})
            for agent_id in normalized_ids
        ]
        return query.where(or_(*conditions))
