"""
Search Service for RAG

Handles semantic search across legal documents using Qdrant and embeddings.
"""

import httpx
from typing import List, Optional, Dict, Any
from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance, VectorParams, PointStruct,
    Filter, FieldCondition, MatchValue
)
import structlog

logger = structlog.get_logger()


class SearchService:
    """Service for semantic search across legal documents."""

    def __init__(self, settings):
        self.settings = settings
        self.qdrant = QdrantClient(url=settings.qdrant_url)
        self.collection = settings.qdrant_collection
        self.regulations: Dict[str, Dict] = {}
        self.total_chunks = 0

    async def initialize(self):
        """Initialize the search service and load legal corpus."""
        # Ensure collection exists
        try:
            self.qdrant.get_collection(self.collection)
            logger.info("Using existing collection", collection=self.collection)
        except Exception:
            # Create collection
            self.qdrant.create_collection(
                collection_name=self.collection,
                vectors_config=VectorParams(
                    size=1024,  # bge-m3 dimension
                    distance=Distance.COSINE
                )
            )
            logger.info("Created collection", collection=self.collection)

        # Load built-in regulations metadata
        self._load_regulations_metadata()

        # Index legal corpus if empty
        info = self.qdrant.get_collection(self.collection)
        if info.points_count == 0:
            await self._index_legal_corpus()

        self.total_chunks = info.points_count

    def _load_regulations_metadata(self):
        """Load metadata for available regulations."""
        self.regulations = {
            "DSGVO": {
                "code": "DSGVO",
                "name": "Datenschutz-Grundverordnung",
                "full_name": "Verordnung (EU) 2016/679",
                "effective": "2018-05-25",
                "chunks": 99,
                "articles": list(range(1, 100))
            },
            "AI_ACT": {
                "code": "AI_ACT",
                "name": "EU AI Act",
                "full_name": "Verordnung über Künstliche Intelligenz",
                "effective": "2025-02-02",
                "chunks": 85,
                "articles": list(range(1, 114))
            },
            "NIS2": {
                "code": "NIS2",
                "name": "NIS 2 Directive",
                "full_name": "Richtlinie (EU) 2022/2555",
                "effective": "2024-10-17",
                "chunks": 46,
                "articles": list(range(1, 47))
            },
            "TDDDG": {
                "code": "TDDDG",
                "name": "TDDDG",
                "full_name": "Telekommunikation-Digitale-Dienste-Datenschutz-Gesetz",
                "effective": "2021-12-01",
                "chunks": 30,
                "articles": list(range(1, 31))
            },
            "BDSG": {
                "code": "BDSG",
                "name": "BDSG",
                "full_name": "Bundesdatenschutzgesetz",
                "effective": "2018-05-25",
                "chunks": 86,
                "articles": list(range(1, 87))
            }
        }

    async def _index_legal_corpus(self):
        """Index the legal corpus into Qdrant."""
        logger.info("Indexing legal corpus...")

        # Sample chunks for demonstration
        # In production, this would load actual legal documents
        sample_chunks = [
            {
                "content": "Art. 9 Abs. 1 DSGVO: Die Verarbeitung personenbezogener Daten, aus denen die rassische und ethnische Herkunft, politische Meinungen, religiöse oder weltanschauliche Überzeugungen oder die Gewerkschaftszugehörigkeit hervorgehen, sowie die Verarbeitung von genetischen Daten, biometrischen Daten zur eindeutigen Identifizierung einer natürlichen Person, Gesundheitsdaten oder Daten zum Sexualleben oder der sexuellen Orientierung einer natürlichen Person ist untersagt.",
                "regulation_code": "DSGVO",
                "article": "9",
                "paragraph": "1"
            },
            {
                "content": "Art. 6 Abs. 1 DSGVO: Die Verarbeitung ist nur rechtmäßig, wenn mindestens eine der nachstehenden Bedingungen erfüllt ist: a) Die betroffene Person hat ihre Einwilligung zu der Verarbeitung der sie betreffenden personenbezogenen Daten für einen oder mehrere bestimmte Zwecke gegeben.",
                "regulation_code": "DSGVO",
                "article": "6",
                "paragraph": "1"
            },
            {
                "content": "Art. 32 DSGVO: Unter Berücksichtigung des Stands der Technik, der Implementierungskosten und der Art, des Umfangs, der Umstände und der Zwecke der Verarbeitung sowie der unterschiedlichen Eintrittswahrscheinlichkeit und Schwere des Risikos für die Rechte und Freiheiten natürlicher Personen treffen der Verantwortliche und der Auftragsverarbeiter geeignete technische und organisatorische Maßnahmen.",
                "regulation_code": "DSGVO",
                "article": "32",
                "paragraph": "1"
            },
            {
                "content": "Art. 6 AI Act: Hochrisiko-KI-Systeme. Als Hochrisiko-KI-Systeme gelten KI-Systeme, die als Sicherheitskomponente eines Produkts oder selbst als Produkt bestimmungsgemäß verwendet werden sollen.",
                "regulation_code": "AI_ACT",
                "article": "6",
                "paragraph": "1"
            },
            {
                "content": "Art. 21 NIS2: Risikomanagementmaßnahmen im Bereich der Cybersicherheit. Die Mitgliedstaaten stellen sicher, dass wesentliche und wichtige Einrichtungen geeignete und verhältnismäßige technische, operative und organisatorische Maßnahmen ergreifen.",
                "regulation_code": "NIS2",
                "article": "21",
                "paragraph": "1"
            }
        ]

        # Generate embeddings and index
        points = []
        for i, chunk in enumerate(sample_chunks):
            embedding = await self._get_embedding(chunk["content"])
            points.append(PointStruct(
                id=i,
                vector=embedding,
                payload=chunk
            ))

        self.qdrant.upsert(
            collection_name=self.collection,
            points=points
        )

        logger.info("Indexed legal corpus", chunks=len(points))

    async def _get_embedding(self, text: str) -> List[float]:
        """Get embedding for text using Ollama."""
        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    f"{self.settings.ollama_url}/api/embeddings",
                    json={
                        "model": self.settings.embedding_model,
                        "prompt": text
                    },
                    timeout=30.0
                )
                response.raise_for_status()
                return response.json()["embedding"]
        except Exception as e:
            logger.error("Embedding failed", error=str(e))
            # Return zero vector as fallback
            return [0.0] * 1024

    async def search(
        self,
        query: str,
        regulation_codes: Optional[List[str]] = None,
        limit: int = 10,
        min_score: float = 0.7
    ) -> List[Dict[str, Any]]:
        """Perform semantic search."""
        # Get query embedding
        query_embedding = await self._get_embedding(query)

        # Build filter
        search_filter = None
        if regulation_codes:
            search_filter = Filter(
                should=[
                    FieldCondition(
                        key="regulation_code",
                        match=MatchValue(value=code)
                    )
                    for code in regulation_codes
                ]
            )

        # Search
        results = self.qdrant.search(
            collection_name=self.collection,
            query_vector=query_embedding,
            query_filter=search_filter,
            limit=limit,
            score_threshold=min_score
        )

        return [
            {
                "content": hit.payload.get("content", ""),
                "regulation_code": hit.payload.get("regulation_code", ""),
                "article": hit.payload.get("article"),
                "paragraph": hit.payload.get("paragraph"),
                "score": hit.score,
                "metadata": hit.payload
            }
            for hit in results
        ]

    def get_regulations(self) -> List[Dict]:
        """Get list of available regulations."""
        return [
            {
                "code": reg["code"],
                "name": reg["name"],
                "chunks": reg["chunks"],
                "last_updated": reg["effective"]
            }
            for reg in self.regulations.values()
        ]

    def get_regulation(self, code: str) -> Optional[Dict]:
        """Get details of a specific regulation."""
        return self.regulations.get(code)
