"""Hybrid fix search with LLM re-ranking (Phase 1).

Combines three search stages:
1. Fast vector similarity search (top 50)
2. Keyword/metadata filtering
3. LLM re-ranking for relevance (top 10)
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, List, Optional

from autodebug.llm_first.models import (
    FailureContext,
    FailureQuery,
    FixStrategy,
    FlakinessCategory,
    RankingResponse,
    SimilarFix,
)

if TYPE_CHECKING:
    from autodebug.llm_first.fix_embeddings import FixEmbeddingPipeline

logger = logging.getLogger(__name__)


class HybridFixSearch:
    """Hybrid search combining vector similarity, filtering, and LLM re-ranking.

    Three-stage search:
    1. Vector search: Fast approximate nearest neighbor (top 50)
    2. Keyword filter: Filter by error type, category, framework
    3. LLM re-rank: Use LLM to score relevance (top 10)
    """

    def __init__(
        self,
        db=None,
        embedding_pipeline: Optional[FixEmbeddingPipeline] = None,
        instructor_client=None,
    ):
        """Initialize hybrid search.

        Args:
            db: Database client for vector search.
            embedding_pipeline: Pipeline for generating query embeddings.
            instructor_client: LLM client for re-ranking.
        """
        self.db = db
        self.embedding_pipeline = embedding_pipeline
        self.instructor_client = instructor_client

    async def search(
        self,
        query: FailureQuery,
        top_k: int = 10,
        use_reranking: bool = True,
    ) -> List[SimilarFix]:
        """Search for similar fixes.

        Args:
            query: Query with error type, message, and optional embedding.
            top_k: Number of results to return.
            use_reranking: Whether to use LLM re-ranking.

        Returns:
            List of similar fixes ranked by relevance.
        """
        # Stage 1: Vector search (fast, returns top 50)
        vector_results = await self._vector_search(query, limit=50)

        if not vector_results:
            logger.info("No vector search results found")
            return []

        # Stage 2: Keyword/metadata filtering
        filtered_results = self._keyword_filter(vector_results, query)

        if not filtered_results:
            # Fall back to vector results if filter eliminates everything
            filtered_results = vector_results[:20]

        # Stage 3: LLM re-ranking (expensive, top 20 only)
        if use_reranking and len(filtered_results) > 1:
            reranked = await self._llm_rerank(
                filtered_results[:20],
                query,
            )
            return reranked[:top_k]

        return filtered_results[:top_k]

    async def search_from_context(
        self,
        context: FailureContext,
        top_k: int = 10,
    ) -> List[SimilarFix]:
        """Search for similar fixes from a failure context.

        Args:
            context: Full failure context.
            top_k: Number of results to return.

        Returns:
            List of similar fixes.
        """
        # Generate embedding if not present
        embedding = context.embedding
        if not embedding and self.embedding_pipeline:
            embedding = await self.embedding_pipeline.embed_failure_context(context)

        query = FailureQuery(
            error_type=context.error_type,
            error_message=context.error_message,
            category=context.category,
            embedding=embedding,
            framework=context.framework,
        )

        return await self.search(query, top_k=top_k)

    async def _vector_search(
        self,
        query: FailureQuery,
        limit: int = 50,
    ) -> List[SimilarFix]:
        """Fast vector similarity search.

        Args:
            query: Query with embedding.
            limit: Maximum results.

        Returns:
            List of fixes with similarity scores.
        """
        if not self.db:
            logger.warning("No database configured for vector search")
            return []

        # Generate embedding if not provided
        embedding = query.embedding
        if not embedding and self.embedding_pipeline:
            embedding = await self.embedding_pipeline.embed_error_query(
                query.error_type,
                query.error_message or "",
                query.category,
            )

        if not embedding:
            logger.warning("No embedding available for search")
            return []

        # Query database
        try:
            # Use pgvector cosine similarity search
            results = await self.db.fetch_all(
                """
                SELECT
                    fe.fix_id,
                    vf.description,
                    vf.diff,
                    fe.error_type,
                    fe.category,
                    fe.strategy,
                    fe.verification_pass_rate,
                    fe.user_acceptance_rate,
                    fe.times_led_to_success,
                    1 - (fe.combined_embedding <=> $1::halfvec) as similarity
                FROM fix_embeddings fe
                JOIN verified_fixes vf ON vf.id = fe.fix_id
                ORDER BY fe.combined_embedding <=> $1::halfvec
                LIMIT $2
                """,
                embedding,
                limit,
            )

            return [
                SimilarFix(
                    fix_id=str(r["fix_id"]),
                    similarity_score=float(r["similarity"]),
                    error_type=r["error_type"],
                    category=FlakinessCategory(r["category"]) if r["category"] else FlakinessCategory.UNKNOWN,
                    description=r["description"] or "",
                    diff=r["diff"] or "",
                    verification_pass_rate=float(r["verification_pass_rate"] or 1.0),
                    user_acceptance_rate=float(r["user_acceptance_rate"]) if r["user_acceptance_rate"] else None,
                    times_successful=int(r["times_led_to_success"] or 0),
                    strategy=FixStrategy(r["strategy"]) if r["strategy"] else None,
                )
                for r in results
            ]
        except Exception as e:
            logger.error(f"Vector search failed: {e}")
            return []

    def _keyword_filter(
        self,
        results: List[SimilarFix],
        query: FailureQuery,
    ) -> List[SimilarFix]:
        """Filter results by keywords and metadata.

        Args:
            results: Vector search results.
            query: Query with filter criteria.

        Returns:
            Filtered results.
        """
        filtered = []

        for fix in results:
            score_boost = 0.0

            # Exact error type match gets boost
            if fix.error_type.lower() == query.error_type.lower():
                score_boost += 0.2

            # Category match gets boost
            if query.category and fix.category == query.category:
                score_boost += 0.15

            # Framework match would get boost (if available)
            # if query.framework and fix.framework == query.framework:
            #     score_boost += 0.1

            # Error message keyword overlap
            if query.error_message:
                query_words = set(query.error_message.lower().split())
                desc_words = set(fix.description.lower().split()) if fix.description else set()
                overlap = len(query_words & desc_words)
                if overlap > 0:
                    score_boost += min(0.1 * overlap, 0.2)

            # Quality signals
            if fix.verification_pass_rate > 0.9:
                score_boost += 0.1
            if fix.user_acceptance_rate and fix.user_acceptance_rate > 0.7:
                score_boost += 0.1
            if fix.times_successful > 5:
                score_boost += 0.05

            # Apply boost to similarity score
            boosted_score = fix.similarity_score + score_boost

            filtered.append(SimilarFix(
                fix_id=fix.fix_id,
                similarity_score=min(boosted_score, 1.0),
                error_type=fix.error_type,
                category=fix.category,
                description=fix.description,
                diff=fix.diff,
                verification_pass_rate=fix.verification_pass_rate,
                user_acceptance_rate=fix.user_acceptance_rate,
                times_successful=fix.times_successful,
                strategy=fix.strategy,
            ))

        # Sort by boosted score
        filtered.sort(key=lambda x: x.similarity_score, reverse=True)
        return filtered

    async def _llm_rerank(
        self,
        candidates: List[SimilarFix],
        query: FailureQuery,
    ) -> List[SimilarFix]:
        """Use LLM to re-rank candidates by relevance.

        Args:
            candidates: Candidate fixes to re-rank.
            query: Original query.

        Returns:
            Re-ranked fixes.
        """
        if not self.instructor_client:
            logger.warning("No LLM client configured for re-ranking")
            return candidates

        # Format candidates for LLM
        candidates_text = self._format_candidates(candidates)

        prompt = f"""
You are evaluating fix candidates for relevance to a specific error.

Current failure:
- Error type: {query.error_type}
- Error message: {query.error_message or 'N/A'}
- Category: {query.category.value if query.category else 'unknown'}

Candidate fixes (each with ID, category, description):
{candidates_text}

Rank these candidates by how relevant they are to the current failure.
Consider:
1. Error type match (exact match is best)
2. Error message similarity
3. Category alignment
4. Fix strategy appropriateness
5. Historical success rate

Return a JSON object with rankings (list of fix IDs in order of relevance).
"""

        try:
            response = await self.instructor_client.generate(
                prompt,
                response_model=RankingResponse,
            )

            return self._reorder_by_ranking(candidates, response.rankings)

        except Exception as e:
            logger.error(f"LLM re-ranking failed: {e}")
            return candidates

    def _format_candidates(self, candidates: List[SimilarFix]) -> str:
        """Format candidates for LLM prompt.

        Args:
            candidates: List of candidate fixes.

        Returns:
            Formatted text for prompt.
        """
        parts = []
        for i, fix in enumerate(candidates, 1):
            # Truncate diff for prompt
            diff_preview = fix.diff[:500] + "..." if len(fix.diff) > 500 else fix.diff

            parts.append(f"""
[{i}] ID: {fix.fix_id}
    Error type: {fix.error_type}
    Category: {fix.category.value}
    Strategy: {fix.strategy.value if fix.strategy else 'unknown'}
    Description: {fix.description[:200]}
    Pass rate: {fix.verification_pass_rate:.0%}
    Diff preview: {diff_preview[:200]}
""")
        return "\n".join(parts)

    def _reorder_by_ranking(
        self,
        candidates: List[SimilarFix],
        rankings: List[str],
    ) -> List[SimilarFix]:
        """Reorder candidates based on LLM rankings.

        Args:
            candidates: Original candidate list.
            rankings: List of fix IDs in ranked order.

        Returns:
            Reordered candidate list.
        """
        # Create lookup
        fix_by_id = {fix.fix_id: fix for fix in candidates}

        # Reorder based on rankings
        reordered = []
        for fix_id in rankings:
            if fix_id in fix_by_id:
                reordered.append(fix_by_id[fix_id])
                del fix_by_id[fix_id]

        # Add any remaining (not in rankings)
        reordered.extend(fix_by_id.values())

        return reordered


async def search_similar_fixes(
    error_type: str,
    error_message: str,
    category: Optional[FlakinessCategory] = None,
    db=None,
    top_k: int = 5,
) -> List[SimilarFix]:
    """Convenience function to search for similar fixes.

    Args:
        error_type: Error/exception type.
        error_message: Error message.
        category: Optional category filter.
        db: Database client.
        top_k: Number of results.

    Returns:
        List of similar fixes.
    """
    from autodebug.llm_first.fix_embeddings import FixEmbeddingPipeline

    pipeline = FixEmbeddingPipeline()
    search = HybridFixSearch(db=db, embedding_pipeline=pipeline)

    query = FailureQuery(
        error_type=error_type,
        error_message=error_message,
        category=category,
    )

    return await search.search(query, top_k=top_k, use_reranking=False)
