"""Retrieval Service"""
import asyncio
import hashlib
import json
from typing import List, Dict, Any, Tuple
import uuid
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from normbase.models.norm import Norm
from normbase.services.qdrant_service import QdrantService
from normbase.services.cache_service import CacheService
from normbase.config import settings
from llm.providers import LLMProviderFactory


def get_llm_provider():
    """Get configured LLM provider from factory-driven config."""
    provider_name = settings.llm_provider
    provider_config = settings.get_llm_provider_config()
    LLMProviderFactory.configure(provider_config)
    provider = LLMProviderFactory.get(provider_name)
    if provider is None:
        raise RuntimeError(
            f"Failed to initialize LLM provider '{provider_name}'. "
            "Check NORMBASE_LLM_PROVIDER/NORMBASE_LLM_API_KEY/"
            "NORMBASE_LLM_MODEL/NORMBASE_LLM_EMBEDDING_MODEL."
        )
    return provider


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],
        agent_role: str,
        field_ids: List[uuid.UUID] = None,
        top_k: int = 5
    ) -> Tuple[List[Tuple[Norm, float]], bool]:
        """
        Retrieve applicable norms using hybrid search
        
        Args:
            db: Database session
            context: Task context
            agent_role: Agent role
            field_ids: Information field IDs to search in
            top_k: Number of results
        
        Returns:
            (List of (Norm, score) tuples, cache_hit)
        """
        cache_key = self._build_cache_key(
            context=context,
            agent_role=agent_role,
            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:
                        return cached_results[:top_k], True
            except Exception as e:
                print(f"Retrieve cache read error: {e}")

        # Use field_ids for SQL search
        if field_ids and len(field_ids) > 0:
            sql_results = await self._sql_search_by_field_ids(db, agent_role, field_ids)
        else:
            sql_results = await self._sql_search_all_active(db, agent_role)
        
        # Channel 2: Vector semantic search (if LLM provider available and context is text-rich)
        if self.llm and self._should_use_vector_search(context):
            vector_results = await self._vector_search(context, agent_role, top_k * 2)
        else:
            vector_results = []
        
        # RRF Fusion
        combined = self._rrf_fusion(sql_results, vector_results, top_k)

        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
    
    async def _sql_search_all_active(
        self,
        db: AsyncSession,
        agent_role: str = None
    ) -> List[Tuple[Norm, float]]:
        """SQL-based search for all active norms, filtered by role"""
        query = select(Norm).where(Norm.status == "active")
        
        if agent_role:
            query = query.where(Norm.scope_roles.contains([agent_role]))
        
        result = await db.execute(query)
        norms = result.scalars().all()
        
        return [(norm, 1.0) for norm in norms]
    
    async def _sql_search_by_field_ids(
        self,
        db: AsyncSession,
        agent_role: str,
        field_ids: List[uuid.UUID]
    ) -> List[Tuple[Norm, float]]:
        """SQL-based exact matching for multiple field IDs"""
        query = select(Norm).where(Norm.status == "active")
        
        if agent_role:
            query = query.where(Norm.scope_roles.contains([agent_role]))
        
        if field_ids:
            query = query.where(Norm.field_id.in_(field_ids))
        
        result = await db.execute(query)
        norms = result.scalars().all()
        
        return [(norm, 1.0) for norm in norms]
    
    async def _sql_search_by_domains(
        self,
        db: AsyncSession,
        agent_role: str = None,
        domains: List[str] = None
    ) -> List[Tuple[Norm, float]]:
        """SQL-based exact matching for one or multiple domains"""
        query = select(Norm).where(Norm.status == "active")
        
        if agent_role:
            query = query.where(Norm.scope_roles.contains([agent_role]))
        
        if domains:
            query = query.where(Norm.domain.in_(domains))
        
        result = await db.execute(query)
        norms = result.scalars().all()
        
        return [(norm, 1.0) for norm in norms]
    
    async def _vector_search(
        self,
        context: Dict[str, Any],
        agent_role: str,
        top_k: int
    ) -> List[Tuple[str, float]]:
        """Vector-based semantic search using Qdrant"""
        try:
            # Build query text
            context_text = self._build_context_text(context, agent_role)

            # 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(context_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 = [f"Agent role: {agent_role}"]
        
        if context:
            for key, value in context.items():
                parts.append(f"{key}: {value}")
        
        return " | ".join(parts)

    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

    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).where(Norm.id.in_(ids)))
        ).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],
        agent_role: str,
        field_ids: List[uuid.UUID] | None,
        top_k: int,
    ) -> str:
        normalized = {
            "agent_role": agent_role,
            "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
