"""Embedding service for fix similarity search.

Generates embeddings using OpenAI's text-embedding-3-small model
for storing and searching similar fixes.
"""

from __future__ import annotations

import hashlib
import logging
import os
from typing import Any, Dict, List, Optional

import httpx

from autodebug.retry import (
    RateLimitError,
    RetryableHTTPError,
    llm_retry,
)
from autodebug.retry import (
    is_available as retry_available,
)

logger = logging.getLogger(__name__)


class EmbeddingError(Exception):
    """Base class for embedding errors."""
    pass


class EmbeddingRateLimitError(EmbeddingError):
    """Rate limit exceeded."""
    pass


class EmbeddingProvider:
    """Embedding provider configuration."""

    OPENAI = "openai"
    VOYAGE = "voyage"


# Model configurations
EMBEDDING_MODELS = {
    # OpenAI models
    "text-embedding-3-small": {"provider": EmbeddingProvider.OPENAI, "dimension": 1536},
    "text-embedding-3-large": {"provider": EmbeddingProvider.OPENAI, "dimension": 3072},
    "text-embedding-ada-002": {"provider": EmbeddingProvider.OPENAI, "dimension": 1536},
    # Voyage AI models (recommended by Anthropic)
    "voyage-3": {"provider": EmbeddingProvider.VOYAGE, "dimension": 1024},
    "voyage-3-lite": {"provider": EmbeddingProvider.VOYAGE, "dimension": 512},
    "voyage-code-3": {"provider": EmbeddingProvider.VOYAGE, "dimension": 1024},
}


class EmbeddingService:
    """Generate embeddings using OpenAI or Voyage AI API."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "text-embedding-3-small",
        cache_enabled: bool = True,
        provider: Optional[str] = None,
    ):
        """Initialize embedding service.

        Args:
            api_key: API key. Defaults to OPENAI_API_KEY or VOYAGE_API_KEY env var.
            model: Embedding model to use.
            cache_enabled: Whether to cache embeddings in memory.
            provider: Force provider (openai or voyage). Auto-detected from model if not set.
        """
        # Determine provider from model or explicit setting
        model_config = EMBEDDING_MODELS.get(model, {})
        self.provider = provider or model_config.get("provider", EmbeddingProvider.OPENAI)

        # Get appropriate API key
        if api_key:
            self.api_key = api_key
        elif self.provider == EmbeddingProvider.VOYAGE:
            self.api_key = os.environ.get("VOYAGE_API_KEY")
        else:
            self.api_key = os.environ.get("OPENAI_API_KEY")

        # Resolve embedding base URL (stays OpenAI unless overridden)
        self._base_url = os.environ.get("HIKAFLOW_EMBEDDING_BASE_URL") or "https://api.openai.com/v1"

        self.model = model
        self.dimension = model_config.get("dimension", 1536)
        self._cache_enabled = cache_enabled
        if cache_enabled:
            from functools import lru_cache
            self._cache: Dict[str, List[float]] = {}
            self._cache_max_size = 10_000
        else:
            self._cache = None
            self._cache_max_size = 0

    def _cache_put(self, key: str, value: List[float]) -> None:
        """Add to cache with bounded size (LRU-like eviction)."""
        if self._cache is None:
            return
        if len(self._cache) >= self._cache_max_size:
            # Evict oldest entries (first 10% of keys)
            evict_count = self._cache_max_size // 10
            keys_to_evict = list(self._cache.keys())[:evict_count]
            for k in keys_to_evict:
                del self._cache[k]
        self._cache[key] = value

    def _get_cache_key(self, text: str) -> str:
        """Generate cache key for text."""
        return hashlib.sha256(text.encode()).hexdigest()[:32]

    async def embed_text(self, text: str, max_retries: int = 3) -> List[float]:
        """Generate embedding for text.

        Args:
            text: Text to embed.
            max_retries: Maximum retry attempts (used if tenacity not available).

        Returns:
            List of floats representing the embedding.
        """
        if not self.api_key:
            logger.warning("No API key configured, returning zero vector")
            return [0.0] * self.dimension

        # Check cache
        cache_key = None
        if self._cache_enabled:
            cache_key = self._get_cache_key(text)
            if cache_key in self._cache:
                return self._cache[cache_key]

        # Truncate text if too long (max ~8000 tokens)
        if len(text) > 30000:
            text = text[:30000]

        try:
            embedding = await self._call_embedding_api(text)

            # Cache result
            if self._cache_enabled and cache_key:
                self._cache_put(cache_key, embedding)

            return embedding

        except Exception as e:
            logger.error(f"Failed to generate embedding: {e}")
            return [0.0] * self.dimension

    async def generate(self, text: str) -> List[float]:
        """Compatibility wrapper used by worker tasks."""
        return await self.embed_text(text)

    async def generate_batch(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for a batch of texts."""
        results: List[List[float]] = []
        for text in texts:
            results.append(await self.embed_text(text))
        return results

    async def _call_embedding_api(self, text: str) -> List[float]:
        """Call the embedding API with retry logic.

        Args:
            text: Text to embed.

        Returns:
            Embedding vector.

        Raises:
            RateLimitError: If rate limited.
            RetryableHTTPError: If server error.
            httpx.HTTPError: If other HTTP error.
        """
        if self.provider == EmbeddingProvider.VOYAGE:
            return await _voyage_embedding_call(self.api_key, self.model, text)
        else:
            return await _openai_embedding_call(self.api_key, self.model, text)

    def format_fix_for_embedding(self, fix_record: Dict[str, Any]) -> str:
        """Format fix record for embedding generation.

        Enhanced to include stack trace, code context, and I/O summary.

        Args:
            fix_record: Dictionary containing fix information.

        Returns:
            Formatted text for embedding.
        """
        parts = [
            f"Exception: {fix_record.get('exception_type', '')}",
            f"Message: {fix_record.get('exception_message', '')}",
            f"File: {fix_record.get('failing_file', '')}",
            f"Test: {fix_record.get('test_name', '')}",
            f"Root cause: {fix_record.get('root_cause', '')}",
            f"Strategy: {fix_record.get('strategy', '')}",
        ]

        # Add stack trace (top 3 frames, anonymized)
        stack = fix_record.get("stack_trace", [])
        if isinstance(stack, list) and stack:
            top_frames = stack[:3]
            parts.append(f"Stack: {' -> '.join(str(f) for f in top_frames)}")

        # Add code context snippet (first 500 chars)
        code_ctx = fix_record.get("code_context", "")
        if code_ctx:
            parts.append(f"Code: {code_ctx[:500]}")

        # Add fix diff summary
        diff = fix_record.get("unified_diff", "")
        if diff:
            added = diff.count("\n+")
            removed = diff.count("\n-")
            parts.append(f"Diff: +{added}/-{removed} lines")

        return "\n".join(p for p in parts if not p.endswith(": "))

    async def embed_query(self, text: str) -> List[float]:
        """Generate embedding for a search query.

        For Voyage models, uses input_type="query" for asymmetric search.
        For other models, delegates to embed_text.

        Args:
            text: Query text to embed.

        Returns:
            Embedding vector.
        """
        if self.provider == EmbeddingProvider.VOYAGE and self.api_key:
            # Check cache
            cache_key = None
            if self._cache_enabled:
                cache_key = self._get_cache_key(f"query:{text}")
                if cache_key in self._cache:
                    return self._cache[cache_key]

            try:
                embedding = await _voyage_embedding_call(
                    self.api_key, self.model, text, input_type="query"
                )
                if self._cache_enabled and cache_key:
                    self._cache_put(cache_key, embedding)
                return embedding
            except Exception as e:
                logger.error(f"Failed to generate query embedding: {e}")
                return [0.0] * self.dimension

        return await self.embed_text(text)

    async def embed_fix_record(self, fix_record: Dict[str, Any]) -> List[float]:
        """Generate embedding for a fix record.

        For Voyage models, uses input_type="document" for asymmetric search.

        Args:
            fix_record: Dictionary containing fix information.

        Returns:
            Embedding vector.
        """
        text = self.format_fix_for_embedding(fix_record)

        if self.provider == EmbeddingProvider.VOYAGE and self.api_key:
            cache_key = None
            if self._cache_enabled:
                cache_key = self._get_cache_key(f"doc:{text}")
                if cache_key in self._cache:
                    return self._cache[cache_key]

            try:
                embedding = await _voyage_embedding_call(
                    self.api_key, self.model, text, input_type="document"
                )
                if self._cache_enabled and cache_key:
                    self._cache_put(cache_key, embedding)
                return embedding
            except Exception as e:
                logger.error(f"Failed to generate document embedding: {e}")
                return [0.0] * self.dimension

        return await self.embed_text(text)

    async def embed_batch(
        self,
        texts: List[str],
        batch_size: int = 100,
    ) -> List[List[float]]:
        """Batch embed multiple texts.

        Args:
            texts: List of texts to embed.
            batch_size: Number of texts per API call.

        Returns:
            List of embedding vectors.
        """
        if not self.api_key:
            return [[0.0] * self.dimension for _ in texts]

        results = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]

            # Check cache for each
            batch_results = []
            uncached_texts = []
            uncached_indices = []

            for j, text in enumerate(batch):
                if self._cache_enabled:
                    cache_key = self._get_cache_key(text)
                    if cache_key in self._cache:
                        batch_results.append((j, self._cache[cache_key]))
                        continue
                uncached_texts.append(text)
                uncached_indices.append(j)

            # Fetch uncached embeddings
            if uncached_texts:
                try:
                    async with httpx.AsyncClient(timeout=60.0) as client:
                        response = await client.post(
                            f"{self._base_url}/embeddings",
                            headers={"Authorization": f"Bearer {self.api_key}"},
                            json={"model": self.model, "input": uncached_texts},
                        )
                        response.raise_for_status()
                        data = response.json()

                        for k, item in enumerate(data["data"]):
                            idx = uncached_indices[k]
                            embedding = item["embedding"]
                            batch_results.append((idx, embedding))

                            # Cache
                            if self._cache_enabled:
                                cache_key = self._get_cache_key(uncached_texts[k])
                                self._cache_put(cache_key, embedding)

                except Exception as e:
                    logger.error(f"Batch embedding error: {e}")
                    # Fill with zero vectors
                    for k, idx in enumerate(uncached_indices):
                        batch_results.append((idx, [0.0] * self.dimension))

            # Sort by original index and extract embeddings
            batch_results.sort(key=lambda x: x[0])
            results.extend([emb for _, emb in batch_results])

        return results


# Standalone functions for embedding API calls (allows decorator to work)

async def _openai_embedding_call(api_key: str, model: str, text: str) -> List[float]:
    """Call OpenAI embedding API.

    Args:
        api_key: OpenAI API key.
        model: Model name.
        text: Text to embed.

    Returns:
        Embedding vector.
    """
    _emb_base = os.environ.get("HIKAFLOW_EMBEDDING_BASE_URL") or "https://api.openai.com/v1"
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(
            f"{_emb_base}/embeddings",
            headers={"Authorization": f"Bearer {api_key}"},
            json={"model": model, "input": text},
        )

        if response.status_code == 429:
            retry_after = int(response.headers.get("Retry-After", 60))
            raise RateLimitError(retry_after=retry_after)

        if response.status_code >= 500:
            raise RetryableHTTPError(
                status_code=response.status_code,
                message=response.text[:200],
            )

        response.raise_for_status()
        data = response.json()
        return data["data"][0]["embedding"]


async def _voyage_embedding_call(
    api_key: str,
    model: str,
    text: str,
    input_type: str = "document",
) -> List[float]:
    """Call Voyage AI embedding API.

    Voyage AI is recommended by Anthropic for embeddings.

    Args:
        api_key: Voyage API key.
        model: Model name (voyage-3, voyage-3-lite, voyage-code-3).
        text: Text to embed.
        input_type: "document" for stored content, "query" for search queries.

    Returns:
        Embedding vector.
    """
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(
            "https://api.voyageai.com/v1/embeddings",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": model,
                "input": text,
                "input_type": input_type,
            },
        )

        if response.status_code == 429:
            retry_after = int(response.headers.get("Retry-After", 60))
            raise RateLimitError(retry_after=retry_after)

        if response.status_code >= 500:
            raise RetryableHTTPError(
                status_code=response.status_code,
                message=response.text[:200],
            )

        response.raise_for_status()
        data = response.json()
        return data["data"][0]["embedding"]


# Apply retry decorator if available
if retry_available():
    _openai_embedding_call = llm_retry(max_attempts=3)(_openai_embedding_call)
    _voyage_embedding_call = llm_retry(max_attempts=3)(_voyage_embedding_call)
