"""
reranker.py
Reranks retrieved chunks using the Cohere Rerank API.

Cohere rerank takes the query and a list of texts and returns
relevance scores — no local model, no GPU needed.

Sign up: https://dashboard.cohere.com
Free tier: 1000 calls/month. Paid: $1 per 1000 calls.

Dependencies:
    pip install cohere python-dotenv
"""

import logging
import os
from typing import List

import cohere
from dotenv import load_dotenv

from .models import RankedResult, SearchResult

load_dotenv()

logger = logging.getLogger(__name__)

RERANK_MODEL = "rerank-english-v3.0"


def create_cohere_client() -> cohere.Client:
    """
    Create a Cohere client from environment variables.

    Requires in .env:
        COHERE_API_KEY = your-key
    """
    api_key = os.getenv("COHERE_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "COHERE_API_KEY not found. Add it to your .env:\n"
            "  COHERE_API_KEY=your-key\n"
            "Get one at: https://dashboard.cohere.com"
        )
    return cohere.Client(api_key=api_key)


def rerank(
    query: str,
    results: List[SearchResult],
    cohere_client: cohere.Client,
    top_n: int = 10,
) -> List[RankedResult]:
    """
    Rerank retrieved chunks using Cohere Rerank.

    Args:
        query:          User query string.
        results:        Chunks from the retrieval step.
        cohere_client:  Cohere client.
        top_n:          Number of top results to return after reranking.

    Returns:
        List of RankedResult ordered by Cohere relevance score, best first.
    """
    if not results:
        return []

    logger.info("Reranking %d chunks (top_n=%d)...", len(results), top_n)

    response = cohere_client.rerank(
        model=RERANK_MODEL,
        query=query,
        documents=[r.text for r in results],
        top_n=top_n,
    )

    ranked = []
    for hit in response.results:
        original = results[hit.index]
        ranked.append(RankedResult(
            chunk_id=     original.chunk_id,
            text=         original.text,
            source=       original.source,
            file_path=    original.file_path,
            rerank_score= hit.relevance_score,
        ))

    logger.info(
        "Reranking complete. Top score: %.4f  Bottom score: %.4f",
        ranked[0].rerank_score if ranked else 0,
        ranked[-1].rerank_score if ranked else 0,
    )
    return ranked