import os
import json
import logging
import re
from typing import List, Dict, Any, Optional
from dotenv import load_dotenv
import requests

try:
    from supabase import create_client, Client
except ImportError:
    create_client = None
    Client = None

from app.services.ai.ai_gateway import create_ai_gateway

load_dotenv(override=True)
logger = logging.getLogger(__name__)

class VectorRetriever:
    """
    Modul untuk mengambil teks chunks dari Supabase Vector Database.
    Mendukung Metadata-Filtered Vector Search menggunakan OpenAI Embeddings
    beserta Multi-Query Expansion.
    """

    def __init__(self, table_name: str = "document_chunk"):
        self.table_name = table_name
        self.url = os.getenv("SUPABASE_URL")
        self.key = os.getenv("SUPABASE_SERVICE_ROLE_KEY")

        self.client: Optional[Client] = None
        if create_client and self.url and self.key:
            try:
                self.client = create_client(self.url, self.key)
                logger.info("Berhasil inisialisasi Supabase Client untuk VectorRetriever.")
            except Exception as e:
                logger.error(f"Gagal inisialisasi Supabase Client: {e}")
        else:
            logger.warning("Kredensial Supabase tidak lengkap atau library supabase belum diinstal.")

        self.gateway = None
        self.embedding_engine = None
        try:
            self.gateway = create_ai_gateway()
            self.embedding_engine = self.gateway.get_embedding_engine()
        except Exception as exc:
            logger.warning("Gateway atau embedding engine tidak siap: %s", exc)

    def get_openai_embedding(self, text: str) -> List[float]:
        if not self.embedding_engine:
            raise ValueError("Gateway embedding engine tidak tersedia")

        return self.embedding_engine.embed(text)

    def _generate_multi_queries(self, original_query: str) -> List[str]:
        """
        Menghasilkan variasi query menggunakan LLM untuk meningkatkan recall pencarian vektor.
        """
        if not self.gateway:
            return [original_query]

        prompt = f"""
Anda adalah ahli sistem RAG. Tugas Anda adalah mengembangkan pertanyaan pengguna berikut menjadi 3 pertanyaan berbeda.
Tujuannya adalah menangkap sinonim, singkatan, dan konteks semantik yang lebih luas untuk pencarian Vector Database.
Jika pertanyaannya sangat pendek atau tidak logis untuk diperluas, kembalikan saja pertanyaan aslinya berulang kali dengan variasi minimal.

Pertanyaan Asli: "{original_query}"

Kembalikan HANYA format JSON list of strings yang valid, berisi pertanyaan asli ditambah 2 variasinya (total maksimal 3 query).
Contoh keluaran yang diharapkan:
["Pertanyaan asli", "Variasi makna 1", "Variasi makna 2"]
"""
        try:
            response = self.gateway.generate_llm_response(
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0
            )
            raw_text = response.choices[0].message.content

            # Cari JSON di dalam teks jika LLM bandel dan menambahkan teks ekstra
            json_match = re.search(r'\[.*\]', raw_text, re.DOTALL)
            if json_match:
                queries = json.loads(json_match.group())
            else:
                queries = json.loads(raw_text)

            if isinstance(queries, list) and len(queries) > 0:
                logger.info(f"Multi-query berhasil di-generate: {queries}")
                # Batasi maksimal 3 query untuk menghemat kuota RPC
                return queries[:3]

            return [original_query]
        except Exception as e:
            logger.error(f"Gagal generate multi-query (fallback ke original): {e}")
            return [original_query]

    def similarity_search_multi_query(self, query: str, hash_list: List[str], top_k: int = 25) -> List[Dict[str, Any]]:
        """
        Pencarian Semantic Vector yang menggabungkan hasil dari multi-query dan beberapa hashes.
        (Metode standar industri sebagai pengganti Cross-Encoder Reranker).
        """
        if not self.client or not self.embedding_engine or not hash_list:
            logger.warning("Supabase, gateway embedding, atau hash_list kosong!")
            return []

        try:
            # 1. Expand query menjadi beberapa variasi
            queries = self._generate_multi_queries(query)

            # 2. Embed semua variasi query
            query_vectors = [self.get_openai_embedding(q) for q in queries]

            unique_chunks = {}

            # 3. Looping pemanggilan pencarian vektor ke setiap hash DAN setiap vektor
            for active_hash in hash_list:
                for q_vector in query_vectors:
                    result = self.client.rpc(
                        "match_document_chunks",
                        {
                            "query_embedding": q_vector,
                            "match_count": top_k,  # Ambil top_k kandidat per kombinasi
                            "filter_hash": active_hash
                        }
                    ).execute()

                    for row in result.data:
                        chunk_id = row.get("id")
                        similarity = row.get("similarity", 0.0)

                        # 4. Deduplikasi: Simpan jika belum ada, atau jika similarity baru lebih tinggi
                        if chunk_id not in unique_chunks or similarity > unique_chunks[chunk_id].get("similarity", 0.0):
                            unique_chunks[chunk_id] = {
                                "id": chunk_id,
                                "file_hash": row.get("file_hash"),
                                "title": row.get("title", ""),
                                "content": row.get("content"),
                                "metadata_header": row.get("metadata_header", ""),
                                "similarity": similarity
                            }

            # 5. Urutkan secara menurun berdasarkan cosine similarity tertinggi
            sorted_chunks = sorted(unique_chunks.values(), key=lambda x: x["similarity"], reverse=True)

            final_results = sorted_chunks[:top_k]
            logger.info(f"Multi-Query Search berhasil menemukan {len(final_results)} top chunks unik dari {len(queries)} query dan {len(hash_list)} dokumen.")
            return final_results

        except Exception as e:
            logger.error(f"Multi-Query Search gagal: {e}", exc_info=True)
            return []

    def similarity_search_by_hash(self, query: str, active_hash: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        Pencarian Semantic Vector dengan strict filter pada active_hash (Legacy).
        """
        if not self.client or not self.embedding_engine:
            logger.warning("Supabase atau gateway embedding belum siap!")
            return []

        try:
            # 1. Embed query menjadi vektor 1536D
            query_vector = self.get_openai_embedding(query)

            # 2. Panggil RPC di Supabase
            result = self.client.rpc(
                "match_document_chunks",
                {
                    "query_embedding": query_vector,
                    "match_count": top_k,
                    "filter_hash": active_hash
                }
            ).execute()

            chunks = []
            for row in result.data:
                chunks.append({
                    "id": row.get("id"),
                    "file_hash": row.get("file_hash"),
                    "title": row.get("title", ""),
                    "content": row.get("content"),
                    "metadata_header": row.get("metadata_header", ""),
                    "similarity": row.get("similarity", 0.0)
                })

            logger.info(f"Vector Search berhasil menemukan {len(chunks)} chunks relevan.")
            return chunks
        except Exception as e:
            logger.error(f"Vector Search gagal: {e}")
            return []

    def fetch_chunks_by_hashes(self, hash_list: List[str], filter_pillar: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        Mengambil teks (chunks) dari Supabase berdasarkan daftar hash.
        Ini bukan vector search, melainkan batch retrieval terarah.
        
        Args:
            hash_list: Daftar file_hash dari dokumen yang relevan (dari frontend atau Neo4j).
            filter_pillar: (Opsional) jika ingin mengambil bagian tertentu saja, misalnya "Methodology".
            
        Returns:
            List of chunks (dictionaries) yang berisi teks dan metadatanya.
        """
        if not self.client or not hash_list:
            return []

        try:
            query = self.client.table(self.table_name).select("id, content, metadata_header, title, file_hash").in_("file_hash", hash_list)
            
            if filter_pillar:
                query = query.ilike("metadata_header", f"%{filter_pillar}%")
                
            response = query.execute()
            
            chunks = []
            for row in response.data:
                chunks.append({
                    "id": row.get("id"),
                    "file_hash": row.get("file_hash"),
                    "title": row.get("title", ""),
                    "content": row.get("content"),
                    "metadata_header": row.get("metadata_header", "")
                })
                
            logger.info(f"Berhasil mengambil {len(chunks)} chunks untuk {len(hash_list)} hash.")
            return chunks
            
        except Exception as e:
            logger.error(f"Gagal mengambil chunks dari Supabase secara batch: {e}")
            return []
