import os
import logging
from typing import List, Dict, Any
from dotenv import load_dotenv

import cohere

load_dotenv()
logger = logging.getLogger(__name__)

class ContextReranker:
    """
    Modul Reranking menggunakan Native Cohere Rerank v3.0 API.
    Sangat akurat, anti-halusinasi, dan tidak memerlukan format JSON yang rapuh.
    """

    def __init__(self, model_name: str = "rerank-multilingual-v3.0"):
        self.api_key = os.getenv("COHERE_API_KEY")
        self.model_name = model_name
        self.client = None
        
        if self.api_key:
            try:
                self.client = cohere.ClientV2(self.api_key)
            except Exception as e:
                logger.error(f"Gagal inisialisasi Cohere Client untuk Reranker: {e}")
        else:
            logger.warning("COHERE_API_KEY tidak ditemukan.")

    async def score_chunks(self, query: str, chunks_list: List[Dict[str, Any]], top_k: int = 5) -> List[Dict[str, Any]]:
        """
        Melakukan reranking menggunakan Cohere API native.
        """
        if not self.client or not chunks_list:
            logger.warning("Reranker tidak aktif atau chunks kosong. Mengembalikan data asli terbatas.")
            return chunks_list[:top_k]

        if len(chunks_list) <= top_k:
            return chunks_list

        try:
            # Ambil isi teks dokumen
            docs = [c.get('content', '') for c in chunks_list]
            
            response = self.client.rerank(
                model=self.model_name,
                query=query,
                documents=docs,
                top_n=top_k
            )
            
            filtered_chunks = []
            for result in response.results:
                filtered_chunks.append(chunks_list[result.index])
                
            return filtered_chunks
            
        except Exception as e:
            logger.error(f"Cohere Reranker API error: {e}. Fallback ke chunk asli.")
            return chunks_list[:top_k]
