"""HyDE (Hypothetical Document Embeddings) query expansion implementation.

This module implements HyDE query expansion using Ollama LLMs to generate
hypothetical answer documents that improve semantic search quality.

Reference:
    - HyDE: Precise Zero-Shot Dense Retrieval without Relevance Labels
      https://arxiv.org/abs/2212.10496
    
HyDE works by:
    1. Taking the user's query
    2. Using an LLM to generate a hypothetical answer/document
    3. Embedding that hypothetical document
    4. Using the hypothetical embedding for semantic search
    
This is particularly effective for:
    - Complex queries that need interpretation
    - Vague or underspecified queries
    - Queries where the answer format differs from the question format
    
Example:
    >>> from hyde import HyDEQueryExpander
    >>> 
    >>> expander = HyDEQueryExpander(model_name="qwen3:0.6b")
    >>> 
    >>> # Simple expansion
    >>> result = expander.expand_query("how to implement auth")
    >>> print(result.hypothetical_document)
    >>> 
    >>> # Multiple hypotheses for better coverage
    >>> results = expander.expand_query_multi("async patterns", num_hypotheses=3)
    >>> for r in results:
    ...     print(f"Hypothesis: {r.hypothetical_document[:100]}...")
"""

import asyncio
import logging
from dataclasses import dataclass
from typing import List, Optional, Tuple

import httpx

logger = logging.getLogger(__name__)


@dataclass
class HyDEResult:
    """Result of HyDE query expansion.
    
    Attributes:
        original_query: The original user query
        hypothetical_document: Generated hypothetical answer document
        latency_ms: Time taken to generate the document (in milliseconds)
        model_used: Name of the LLM model used for generation
    """
    original_query: str
    hypothetical_document: str
    latency_ms: float
    model_used: str


class HyDEQueryExpander:
    """HyDE query expander using Ollama LLMs.
    
    HyDE (Hypothetical Document Embeddings) improves retrieval by generating
    hypothetical answer documents and using their embeddings for search.
    
    The key insight is that while the query might be vague or complex, the
    hypothetical answer (generated by an LLM) is often semantically closer
    to the actual relevant documents in the embedding space.
    
    Performance Characteristics:
        - Latency: ~100-500ms per query (depends on model & prompt)
        - Memory: Minimal (only stores prompt and generated text)
        - Quality: +15-30% improvement on complex/vague queries
        
    Configuration Options:
        - prompt_template: Customize the hypothetical document generation
        - max_tokens: Control the length of generated documents
        - temperature: Control creativity vs consistency
        - num_hypotheses: Generate multiple hypotheses for better coverage
        
    Example:
        >>> expander = HyDEQueryExpander(model_name="qwen3:0.6b")
        >>> 
        >>> # Single hypothesis (default)
        >>> result = expander.expand_query("error handling patterns")
        >>> print(result.hypothetical_document)
        
        >>> # Multiple hypotheses for ensemble search
        >>> results = expander.expand_query_multi("database connection", num_hypotheses=3)
        >>> embeddings = [embed(r.hypothetical_document) for r in results]
    """
    
    # Default prompt template for code/technical queries
    DEFAULT_PROMPT_TEMPLATE = """You are a helpful coding assistant. Given a user query about code, generate a hypothetical code snippet or documentation excerpt that would be a perfect answer to the query.

The hypothetical document should:
- Be realistic and detailed (like actual code/documentation)
- Include specific function names, parameters, and implementation details
- Match the style of actual source code or technical documentation
- Directly address the query with concrete examples

User Query: {query}

Generate a hypothetical code snippet or documentation excerpt that answers this query:"""

    # Alternative prompt for explanation-style queries
    EXPLANATION_PROMPT_TEMPLATE = """Given the following query about code or software development, write a detailed explanation that would be found in high-quality technical documentation or code comments.

Query: {query}

Write a comprehensive answer as if it were from expert documentation:"""

    def __init__(
        self,
        model_name: str = "qwen3:0.6b",
        ollama_base_url: str = "http://localhost:11434",
        prompt_template: Optional[str] = None,
        max_tokens: int = 300,
        temperature: float = 0.7,
        timeout: float = 60.0,
        max_concurrent: int = 3
    ):
        """Initialize the HyDE query expander.
        
        Args:
            model_name: Name of the Ollama LLM to use for generation.
                       Recommended: "qwen3:0.6b" (fast, capable) or
                       "phi3:mini" (very fast, lightweight)
            ollama_base_url: URL for the Ollama API server
            prompt_template: Custom prompt template for generation.
                           Must contain {query} placeholder.
                           If None, uses the default code-focused template.
            max_tokens: Maximum tokens to generate for hypothetical document.
                       Longer documents may capture more context but increase
                       latency and may dilute the key semantic signal.
            temperature: Sampling temperature (0.0-1.0).
                        Lower = more deterministic, higher = more creative.
                        0.3-0.5 recommended for factual queries.
            timeout: HTTP request timeout in seconds
            max_concurrent: Maximum concurrent generation requests.
                           Relevant when generating multiple hypotheses.
        
        Raises:
            ValueError: If model_name is empty or prompt template is invalid
        """
        if not model_name or not model_name.strip():
            raise ValueError("model_name cannot be empty")
        
        self.model_name = model_name
        self.ollama_base_url = ollama_base_url.rstrip("/")
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.timeout = timeout
        self.max_concurrent = max_concurrent
        
        # Set prompt template
        self.prompt_template = prompt_template or self.DEFAULT_PROMPT_TEMPLATE
        
        # Validate prompt template
        if "{query}" not in self.prompt_template:
            raise ValueError("prompt_template must contain {query} placeholder")
        
        # Ollama API endpoint
        self.api_url = f"{self.ollama_base_url}/api/generate"
        
        logger.info(
            "Initialized HyDEQueryExpander with model=%s, max_tokens=%d, temp=%.2f",
            model_name, max_tokens, temperature
        )
    
    def expand_query(self, query: str) -> HyDEResult:
        """Expand a query into a hypothetical document.
        
        This is the main HyDE operation - generate a single hypothetical
        document that represents the ideal answer to the query.
        
        Args:
            query: The user's search query
            
        Returns:
            HyDEResult containing the hypothetical document and metadata
            
        Raises:
            ValueError: If query is empty
            
        Example:
            >>> expander = HyDEQueryExpander()
            >>> result = expander.expand_query("authentication middleware")
            >>> print(result.hypothetical_document)
        """
        if not query or not query.strip():
            raise ValueError("query cannot be empty")
        
        import time
        start_time = time.time()
        
        try:
            hypothetical_doc = asyncio.run(self._generate_async(query))
        except RuntimeError as e:
            if "already running" in str(e):
                logger.debug("Event loop already running, using sync fallback")
                hypothetical_doc = self._generate_sync(query)
            else:
                raise
        
        latency_ms = (time.time() - start_time) * 1000
        
        logger.debug(
            "HyDE expansion completed in %.2fms (model=%s, doc_length=%d)",
            latency_ms, self.model_name, len(hypothetical_doc)
        )
        
        return HyDEResult(
            original_query=query,
            hypothetical_document=hypothetical_doc,
            latency_ms=latency_ms,
            model_used=self.model_name
        )
    
    def expand_query_multi(
        self, 
        query: str, 
        num_hypotheses: int = 3
    ) -> List[HyDEResult]:
        """Generate multiple hypothetical documents for ensemble retrieval.
        
        Multiple hypotheses can improve recall by capturing different
        interpretations of the query. The embeddings of all hypothetical
        documents can be combined (e.g., averaged) for search.
        
        Args:
            query: The user's search query
            num_hypotheses: Number of hypothetical documents to generate.
                           More hypotheses = better coverage but higher latency.
                           Default: 3 (good balance of quality vs speed)
        
        Returns:
            List of HyDEResult objects, one per hypothesis
            
        Raises:
            ValueError: If query is empty or num_hypotheses < 1
            
        Example:
            >>> expander = HyDEQueryExpander()
            >>> results = expander.expand_query_multi("async patterns", num_hypotheses=3)
            >>> 
            >>> # Average embeddings for ensemble search
            >>> embeddings = [embed(r.hypothetical_document) for r in results]
            >>> avg_embedding = np.mean(embeddings, axis=0)
        """
        if not query or not query.strip():
            raise ValueError("query cannot be empty")
        if num_hypotheses < 1:
            raise ValueError("num_hypotheses must be at least 1")
        
        import time
        start_time = time.time()
        
        try:
            documents = asyncio.run(self._generate_multi_async(query, num_hypotheses))
        except RuntimeError as e:
            if "already running" in str(e):
                logger.debug("Event loop already running, using sync fallback")
                documents = self._generate_multi_sync(query, num_hypotheses)
            else:
                raise
        
        total_latency_ms = (time.time() - start_time) * 1000
        per_doc_latency = total_latency_ms / len(documents) if documents else 0
        
        logger.debug(
            "HyDE multi-expansion completed: %d hypotheses in %.2fms (%.2fms each)",
            len(documents), total_latency_ms, per_doc_latency
        )
        
        return [
            HyDEResult(
                original_query=query,
                hypothetical_document=doc,
                latency_ms=per_doc_latency,
                model_used=self.model_name
            )
            for doc in documents
        ]
    
    async def _generate_async(self, query: str) -> str:
        """Generate a single hypothetical document asynchronously."""
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            return await self._generate_single(client, query)
    
    def _generate_sync(self, query: str) -> str:
        """Generate a single hypothetical document synchronously."""
        with httpx.Client(timeout=self.timeout) as client:
            return self._generate_single_sync(client, query)
    
    async def _generate_multi_async(
        self, 
        query: str, 
        num_hypotheses: int
    ) -> List[str]:
        """Generate multiple hypothetical documents asynchronously."""
        semaphore = asyncio.Semaphore(self.max_concurrent)
        
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            tasks = [
                self._generate_single_with_semaphore(client, query, i, semaphore)
                for i in range(num_hypotheses)
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)
        
        documents: List[str] = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.error("Hypothesis %d generation failed: %s", i, result)
                # Generate a fallback hypothesis
                documents.append(f"Information about: {query}")
            else:
                documents.append(result)
        
        return documents
    
    def _generate_multi_sync(self, query: str, num_hypotheses: int) -> List[str]:
        """Generate multiple hypothetical documents synchronously."""
        documents: List[str] = []
        
        with httpx.Client(timeout=self.timeout) as client:
            for i in range(num_hypotheses):
                try:
                    doc = self._generate_single_sync(client, query)
                    documents.append(doc)
                except Exception as e:
                    logger.error("Hypothesis %d generation failed: %s", i, e)
                    documents.append(f"Information about: {query}")
        
        return documents
    
    async def _generate_single_with_semaphore(
        self,
        client: httpx.AsyncClient,
        query: str,
        idx: int,
        semaphore: asyncio.Semaphore
    ) -> str:
        """Generate with semaphore-controlled concurrency."""
        async with semaphore:
            # Add slight variation for diversity in multi-hypothesis mode
            if idx > 0:
                varied_query = self._add_query_variation(query, idx)
                return await self._generate_single(client, varied_query)
            return await self._generate_single(client, query)
    
    async def _generate_single(
        self,
        client: httpx.AsyncClient,
        query: str
    ) -> str:
        """Generate a single hypothetical document."""
        prompt = self.prompt_template.format(query=query)
        
        response = await client.post(
            self.api_url,
            json={
                "model": self.model_name,
                "prompt": prompt,
                "stream": False,
                "options": {
                    "temperature": self.temperature,
                    "num_predict": self.max_tokens
                }
            }
        )
        response.raise_for_status()
        data = response.json()
        
        return data.get("response", "").strip()
    
    def _generate_single_sync(
        self,
        client: httpx.Client,
        query: str
    ) -> str:
        """Generate a single hypothetical document synchronously."""
        prompt = self.prompt_template.format(query=query)
        
        response = client.post(
            self.api_url,
            json={
                "model": self.model_name,
                "prompt": prompt,
                "stream": False,
                "options": {
                    "temperature": self.temperature,
                    "num_predict": self.max_tokens
                }
            }
        )
        response.raise_for_status()
        data = response.json()
        
        return data.get("response", "").strip()
    
    def _add_query_variation(self, query: str, idx: int) -> str:
        """Add slight variation to query for diverse hypotheses.
        
        This encourages the model to generate different perspectives
        on the same query, improving coverage.
        """
        variations = [
            query,
            f"Detailed implementation of: {query}",
            f"Example code showing: {query}",
            f"Best practices for: {query}",
            f"How to implement {query} with examples",
            f"Documentation for: {query}",
        ]
        return variations[idx % len(variations)]
    
    def check_model_available(self) -> Tuple[bool, Optional[str]]:
        """Check if the configured model is available in Ollama.
        
        Returns:
            Tuple of (is_available, error_message)
        """
        try:
            with httpx.Client(timeout=10.0) as client:
                response = client.get(f"{self.ollama_base_url}/api/tags")
                response.raise_for_status()
                data = response.json()
                
                models = data.get("models", [])
                model_names = [m.get("name", "") for m in models]
                
                # Check for exact match or tag match
                if self.model_name in model_names:
                    return True, None
                
                # Check without tag
                base_name = self.model_name.split(":")[0]
                if any(m.startswith(base_name) for m in model_names):
                    return True, None
                
                return False, f"Model '{self.model_name}' not found. Available: {model_names}"
                
        except Exception as e:
            return False, f"Failed to check model availability: {e}"
