"""Enhanced fix embedding pipeline for RAG-based fix retrieval (Phase 1).

Creates multi-signal embeddings combining:
- Error context (type + message)
- Code before/after
- Fix description
- Strategy and category metadata

Uses learned weights to combine signals into a unified embedding.
"""

from __future__ import annotations

import logging
import math
from typing import Any, Dict, List, Optional

from autodebug.learning.embeddings import EmbeddingService
from autodebug.llm_first.models import (
    FailureContext,
    FixEmbedding,
    FixStrategy,
    FlakinessCategory,
)

logger = logging.getLogger(__name__)


class FixEmbeddingPipeline:
    """Pipeline for creating rich fix embeddings.

    Combines multiple embedding signals (error, code, description) into
    a unified representation for similarity search.
    """

    # Default weights for combining embeddings
    DEFAULT_WEIGHTS = {
        "error": 0.35,
        "code_before": 0.15,
        "code_after": 0.15,
        "description": 0.35,
    }

    def __init__(
        self,
        embedding_service: Optional[EmbeddingService] = None,
        weights: Optional[Dict[str, float]] = None,
        dimension: int = 512,
    ):
        """Initialize embedding pipeline.

        Args:
            embedding_service: Service for generating embeddings.
            weights: Custom weights for combining signals.
            dimension: Target embedding dimension.
        """
        self.embedding_service = embedding_service or EmbeddingService(
            model="voyage-code-3"  # Use code-specialized model
        )
        self.weights = weights or self.DEFAULT_WEIGHTS
        self.dimension = dimension

    async def embed_fix(
        self,
        fix_id: str,
        error_type: str,
        error_message: str,
        code_before: Optional[str] = None,
        code_after: Optional[str] = None,
        description: str = "",
        category: Optional[FlakinessCategory] = None,
        strategy: Optional[FixStrategy] = None,
        diff: Optional[str] = None,
        verification_pass_rate: float = 1.0,
        user_acceptance_rate: Optional[float] = None,
    ) -> FixEmbedding:
        """Create rich embedding combining multiple signals.

        Args:
            fix_id: Unique identifier for the fix.
            error_type: Exception/error type.
            error_message: Error message.
            code_before: Original code that was fixed.
            code_after: Fixed code.
            description: Human-readable description of the fix.
            category: Flakiness category (timing, ordering, etc.).
            strategy: Fix generation strategy used.
            diff: The diff if available.
            verification_pass_rate: Rate of passing verification runs.
            user_acceptance_rate: Rate of user acceptance.

        Returns:
            FixEmbedding with all signals and combined embedding.
        """
        # Format error context
        error_context = f"{error_type}: {error_message}"
        if category:
            error_context += f"\nCategory: {category.value}"

        # Generate individual embeddings
        error_embedding = await self.embedding_service.embed_text(error_context)

        code_before_embedding = None
        if code_before:
            code_before_embedding = await self.embedding_service.embed_text(
                f"Original code:\n{code_before[:3000]}"
            )

        code_after_embedding = None
        if code_after:
            code_after_embedding = await self.embedding_service.embed_text(
                f"Fixed code:\n{code_after[:3000]}"
            )
        elif diff:
            # Use diff if code_after not available
            code_after_embedding = await self.embedding_service.embed_text(
                f"Fix diff:\n{diff[:3000]}"
            )

        # Format description with strategy context
        full_description = description
        if strategy:
            full_description = f"Strategy: {strategy.value}\n{description}"
        description_embedding = await self.embedding_service.embed_text(full_description)

        # Combine embeddings with learned weights
        combined = self._fuse_embeddings(
            error_embedding=error_embedding,
            code_before_embedding=code_before_embedding,
            code_after_embedding=code_after_embedding,
            description_embedding=description_embedding,
        )

        return FixEmbedding(
            fix_id=fix_id,
            error_embedding=error_embedding,
            code_before_embedding=code_before_embedding,
            code_after_embedding=code_after_embedding,
            description_embedding=description_embedding,
            combined_embedding=combined,
            metadata={
                "error_type": error_type,
                "category": category.value if category else None,
                "strategy": strategy.value if strategy else None,
                "verification_pass_rate": verification_pass_rate,
                "user_acceptance_rate": user_acceptance_rate,
            },
        )

    def _fuse_embeddings(
        self,
        error_embedding: List[float],
        code_before_embedding: Optional[List[float]],
        code_after_embedding: Optional[List[float]],
        description_embedding: List[float],
    ) -> List[float]:
        """Fuse multiple embeddings into a combined representation.

        Uses weighted average with normalization.

        Args:
            error_embedding: Error context embedding.
            code_before_embedding: Original code embedding.
            code_after_embedding: Fixed code embedding.
            description_embedding: Description embedding.

        Returns:
            Combined embedding vector.
        """
        embeddings = []
        weights = []

        # Error embedding (always present)
        embeddings.append(error_embedding)
        weights.append(self.weights["error"])

        # Code before (optional)
        if code_before_embedding:
            embeddings.append(code_before_embedding)
            weights.append(self.weights["code_before"])

        # Code after (optional)
        if code_after_embedding:
            embeddings.append(code_after_embedding)
            weights.append(self.weights["code_after"])

        # Description (always present)
        embeddings.append(description_embedding)
        weights.append(self.weights["description"])

        # Normalize weights
        total_weight = sum(weights)
        normalized_weights = [w / total_weight for w in weights]

        # Weighted average (pure Python, no numpy)
        dim = len(embeddings[0])
        combined = [0.0] * dim

        for emb, weight in zip(embeddings, normalized_weights):
            for i in range(dim):
                combined[i] += weight * emb[i]

        # Normalize to unit vector
        norm = math.sqrt(sum(x * x for x in combined))
        if norm > 0:
            combined = [x / norm for x in combined]

        return combined

    async def embed_failure_context(
        self,
        context: FailureContext,
    ) -> List[float]:
        """Create embedding for a failure context (for querying).

        Args:
            context: Failure context to embed.

        Returns:
            Query embedding.
        """
        # Format query text
        query_parts = [
            f"{context.error_type}: {context.error_message}",
        ]

        if context.category:
            query_parts.append(f"Category: {context.category.value}")

        if context.code_snippet:
            query_parts.append(f"Code:\n{context.code_snippet[:1500]}")

        if context.stack_trace:
            # Just include the most relevant part of stack trace
            stack_lines = context.stack_trace.strip().split("\n")
            relevant_stack = "\n".join(stack_lines[-10:])  # Last 10 lines
            query_parts.append(f"Stack trace:\n{relevant_stack}")

        query_text = "\n\n".join(query_parts)
        return await self.embedding_service.embed_text(query_text)

    async def embed_error_query(
        self,
        error_type: str,
        error_message: str,
        category: Optional[FlakinessCategory] = None,
    ) -> List[float]:
        """Create embedding for an error-based query.

        Args:
            error_type: Exception type.
            error_message: Error message.
            category: Optional category filter.

        Returns:
            Query embedding.
        """
        query_text = f"{error_type}: {error_message}"
        if category:
            query_text += f"\nCategory: {category.value}"

        return await self.embedding_service.embed_text(query_text)

    async def batch_embed_fixes(
        self,
        fixes: List[Dict[str, Any]],
    ) -> List[FixEmbedding]:
        """Batch embed multiple fixes.

        Args:
            fixes: List of fix records with required fields.

        Returns:
            List of FixEmbedding objects.
        """
        results = []
        for fix in fixes:
            try:
                embedding = await self.embed_fix(
                    fix_id=fix["id"],
                    error_type=fix.get("error_type", "Unknown"),
                    error_message=fix.get("error_message", ""),
                    code_before=fix.get("code_before"),
                    code_after=fix.get("code_after"),
                    description=fix.get("description", ""),
                    category=FlakinessCategory(fix["category"]) if fix.get("category") else None,
                    strategy=FixStrategy(fix["strategy"]) if fix.get("strategy") else None,
                    diff=fix.get("diff"),
                    verification_pass_rate=fix.get("verification_pass_rate", 1.0),
                    user_acceptance_rate=fix.get("user_acceptance_rate"),
                )
                results.append(embedding)
            except Exception as e:
                logger.error(f"Failed to embed fix {fix.get('id')}: {e}")

        return results

    def update_weights(self, new_weights: Dict[str, float]) -> None:
        """Update fusion weights (for learning).

        Args:
            new_weights: New weight values.
        """
        for key in new_weights:
            if key in self.weights:
                self.weights[key] = new_weights[key]

        # Ensure weights sum to 1
        total = sum(self.weights.values())
        self.weights = {k: v / total for k, v in self.weights.items()}
