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

from app.services.ai.ai_gateway import create_ai_gateway

from app.services.chatbot.research_policy import ResearchPolicy
from app.services.chatbot.graph_retriever import GraphRetriever
from app.services.chatbot.vector_retriever import VectorRetriever
from app.services.chatbot.chat_history import ChatHistoryManager

load_dotenv()
logger = logging.getLogger(__name__)

class PRALOrchestrator:
    """
    Otak utama Chatbot BRAIN MySRE.
    Menggunakan AI Gateway untuk Penalaran & Jawaban.
    """

    def __init__(self):
        try:
            self.gateway = create_ai_gateway()
            logger.info("Berhasil inisialisasi AI Gateway untuk Orchestrator.")
        except Exception as e:
            logger.error(f"Gagal inisialisasi AI Gateway: {e}")
            self.gateway = None

        self.graph_retriever = GraphRetriever()
        self.vector_retriever = VectorRetriever()
        self.history_manager = ChatHistoryManager()
        
        self.policy = ResearchPolicy()

    async def connect(self):
        """Memastikan koneksi database aktif."""
        await self.graph_retriever.connect()

    async def close(self):
        """Menutup koneksi database."""
        await self.graph_retriever.close()

    async def _perceive_and_reason(self, query: str) -> Dict[str, Any]:
        """
        Fase Perceive & Reason menggunakan AI Gateway.
        """
        if not self.gateway:
            return {"layer": 1, "intent": "Fallback to specific (No Gateway)"}

        prompt = f"""
        Analisis pertanyaan pengguna berikut dan tentukan strategi pencarian yang paling tepat.
        Pilih salah satu dari Layer berikut:
        1 - Layer Spesifik: Pertanyaan sangat fokus pada detail dari SATU dokumen. Tidak memerlukan perbandingan.
        2 - Layer Komparatif: Pertanyaan meminta perbandingan, mencari kesamaan/perbedaan, atau mengeksplorasi hubungan antar dokumen (neighbors).
        3 - Layer General: Pertanyaan bersifat umum, tidak terkait dengan dokumen atau literatur manapun, atau pertanyaan yang sangat "Out of the Box".

        Pertanyaan: "{query}"

        Kembalikan HANYA format JSON valid berikut tanpa tambahan teks apapun:
        {{"layer": <angka 1, 2, atau 3>, "intent": "<penjelasan singkat intent>"}}
        """
            
        try:
            # Use the task-aware async gateway path so reasoning follows the same routing policy.
            response = await self.gateway.agenerate_llm_response(
                messages=[{"role": "user", "content": prompt}],
                task_name="chatbot",
                temperature=0.0
            )
            raw_text = response.choices[0].message.content
            
            # Cari JSON di dalam teks jika ada kata-kata ekstra
            json_match = re.search(r'\{.*\}', raw_text, re.DOTALL)
            if json_match:
                result = json.loads(json_match.group())
            else:
                result = json.loads(raw_text)
                
            layer = int(result.get("layer", 1))
            intent = result.get("intent", "Unknown")
            logger.info(f"[Reasoning] Terpilih Layer {layer} (Intent: {intent})")
            return {"layer": layer, "intent": intent}
        except Exception as e:
            logger.error(f"[Reasoning] Gagal mem-parse intent: {e}. Fallback ke Layer 1.")
            return {"layer": 1, "intent": "Fallback to specific"}

    async def process_chat(
        self,
        query: str,
        active_hashes: List[str],
        project_id: str,
        user_id: str,
        mode: str = "STRICT",
        has_used_research: bool = False,
    ) -> AsyncGenerator[str, None]:
        """
        Siklus lengkap PRAL dengan output StreamingResponse via AIGateway.
        Dilengkapi dengan Mode-Aware UX Guardrails.
        """
        start_time = time.time()
        metrics = {"project_id": project_id, "query_intent": "", "layer_strategy": 1, "source_count": 0}

        try:
            is_document_empty = len(active_hashes) == 0

            # ==========================================
            # LOGIKA MODE STRICT
            # ==========================================
            if mode == "STRICT":
                
                # 1. Handling Jika Dokumen Kosong (Skenario A & B)
                if is_document_empty:
                    if not has_used_research:
                        # Skenario A: Mode STRICT (Awal Sesi, Tanpa Dokumen)
                        yield "Untuk memulai analisis, silakan pilih dokumen melalui *Canvas Knowledge Graph* atau gunakan ikon tambah (+). Jika Anda belum memiliki literatur dalam sistem, silakan unggah dokumen PDF Anda terlebih dahulu. Setelah dokumen aktif, Anda dapat mengajukan pertanyaan terkait konten literatur tersebut."
                        return
                    else:
                        # Skenario B: Mode STRICT (Tanpa Dokumen, Pernah Pakai RESEARCH)
                        yield "Saat ini Anda berada di mode **STRICT**. Mode ini membutuhkan setidaknya satu dokumen aktif untuk dianalisis. Silakan pilih dokumen melalui *Canvas Knowledge Graph* terlebih dahulu.\n\nAtau, jika Anda ingin mencari hubungan antar-literatur secara bebas, silakan alihkan kembali ke mode **RESEARCH**."
                        return

                # 2. Jika dokumen ADA, kita cek niat pertanyaannya (Intent Classification)
                reasoning = await self._perceive_and_reason(query)
                layer = reasoning["layer"]
                metrics["layer_strategy"] = layer
                metrics["query_intent"] = reasoning["intent"]

                if layer == 2:
                    # Skenario C: Mode STRICT (Ada Dokumen, Tapi Mencari Relasi)
                    yield "Sistem mendeteksi bahwa Anda sedang mencari relasi atau literatur lain yang terhubung. Untuk keperluan tersebut, silakan alihkan ke mode **RESEARCH** agar sistem dapat mengeksplorasi *Knowledge Graph* secara menyeluruh. Di mode **STRICT** ini, analisis dibatasi secara eksklusif pada isi teks dokumen yang sedang Anda pilih."
                    return
                    
                elif layer == 3:
                    # Skenario E: Mode STRICT (Ada Dokumen, Pertanyaan Di Luar Konteks)
                    yield "Pertanyaan yang Anda berikan tidak berkaitan langsung dengan dokumen yang sedang aktif. Di mode **STRICT**, sistem memfokuskan analisis secara murni pada literatur yang Anda pilih. Jika Anda ingin membahas topik umum atau melakukan pencarian di luar dokumen ini, silakan alihkan ke mode **RESEARCH**."
                    return

            # ==========================================
            # LOGIKA MODE RESEARCH
            # ==========================================
            elif mode == "RESEARCH":
                
                if is_document_empty:
                    # Skenario D: Mode RESEARCH (Tanpa Dokumen)
                    yield "Saat ini Anda berada di mode **RESEARCH**. Anda dapat mengajukan pertanyaan diskusi secara umum. Namun, untuk mendapatkan analisis yang spesifik beserta referensinya, silakan pilih dokumen melalui *Canvas Knowledge Graph* agar sistem dapat memetakan relasi antar-literatur tersebut."
                    return
                else:
                    # Mode Research Multi-Paper
                    layer = 2
                    metrics["layer_strategy"] = layer
                    metrics["query_intent"] = "Explicit RESEARCH mode (Multi-Paper Comparative)"

            # ==========================================
            # EKSEKUSI RAG NORMAL (Layer 1 atau 2)
            # ==========================================
            target_hashes = list(active_hashes)
            graph_context_text = ""

            # A: Act - Data Retrieval (Graf Relasi hanya di Layer 2)
            if layer == 2:
                t0 = time.time()
                all_neighbors = []
                for a_hash in active_hashes:
                    neighbors = await self.graph_retriever.get_relevant_neighbors(
                        a_hash,
                        project_id=project_id,
                        user_id=user_id,
                        max_nodes=3,
                    )
                    all_neighbors.extend(neighbors)
                neo4j_ms = (time.time() - t0) * 1000
                metrics["neo4j_ms"] = neo4j_ms
                
                if all_neighbors:
                    # Deduplikasi neighbor hash
                    existing = set(target_hashes)
                    graph_context_text = "\n=== HUBUNGAN LITERATUR (DARI KNOWLEDGE GRAPH) ===\n"
                    
                    added_neighbors = 0
                    for n in all_neighbors:
                        if n["hash"] not in existing:
                            target_hashes.append(n["hash"])
                            existing.add(n["hash"])
                            
                        # Catat reasoning hubungan graf untuk diberikan ke LLM
                        graph_context_text += f"Terdapat koneksi dengan literatur '{n['title']}'. "
                        graph_context_text += f"Alasan (Reasoning): {n['reasoning']}\n"
                        added_neighbors += 1
                        
                    if added_neighbors == 0:
                         logger.info("Tetangga sudah ada di target hashes, tidak ada node baru dari Graf.")
                else:
                    logger.info("Tidak ada artikel tetangga yang relevan ditemukan di Graf.")

            final_chunks = []
            if layer in [1, 2]:
                t0 = time.time()
                
                # Menggunakan fungsi Multi-Query Expansion sebagai pengganti Reranker
                final_chunks = self.vector_retriever.similarity_search_multi_query(
                    query=query, 
                    hash_list=target_hashes, 
                    top_k=25
                )
                
                supabase_ms = (time.time() - t0) * 1000
                metrics["supabase_ms"] = supabase_ms
                
                if not final_chunks:
                    logger.warning(f"[Orchestrator] Refusal triggered: final_chunks kosong setelah multi-query search. target_hashes={target_hashes}")
                    missing_doc_msg = "Maaf, saya tidak menemukan informasi yang relevan di dalam dokumen terpilih terkait pertanyaan Anda."
                    self.history_manager.save_turn(
                        project_id=project_id,
                        user_id=user_id,
                        user_query=query,
                        ai_response=missing_doc_msg,
                        references=[],
                        metadata=metrics,
                    )
                    yield missing_doc_msg
                    return

            
            # 3. A: Act - Build Prompt & Generate
            chat_history = self.history_manager.get_history(
                project_id=project_id,
                user_id=user_id,
                limit=3,
            )
            
            doc_context_text = ""
            unique_sources = set()
            
            if layer in [1, 2] and final_chunks:
                doc_context_text = "\n=== KUTIPAN TEKS DARI ARTIKEL ===\n"
                title_cache: Dict[str, str] = {}
                
                for chunk in final_chunks:
                    h = chunk.get("file_hash", "UNKNOWN")
                    unique_sources.add(h)
                    
                    doc_context_text += f"--- SUMBER ARTIKEL ---\n"
                    raw_title = chunk.get('title', '')
                    
                    # Jika judul kosong, 'unknown', atau berupa file .pdf, coba ambil dari Knowledge Graph
                    if not raw_title or raw_title.lower().endswith('.pdf') or raw_title.lower() == "unknown":
                        if h not in title_cache:
                            try:
                                neo4j_title = await self.graph_retriever.get_node_title_by_hash(h)
                                if neo4j_title and not neo4j_title.lower().endswith('.pdf'):
                                    title_cache[h] = neo4j_title
                                else:
                                    title_cache[h] = "Dokumen Akademis"
                            except Exception as e:
                                logger.error(f"Gagal memuat judul dari Neo4j untuk {h}: {e}")
                                title_cache[h] = "Dokumen Akademis"
                        if title_cache[h]:
                            final_title = title_cache[h]
                    else:
                        final_title = raw_title
                    
                    # Format: Propercase
                    final_title = final_title.title()
                    # Format: Truncation (Max 40 chars)
                    if len(final_title) > 40:
                        final_title = final_title[:40] + "..."
                    
                    # Format: HTML Hyperlink
                    html_link = f'<a href="graph://node/{h}">{final_title}</a>'
                    
                    doc_context_text += f"Judul: {html_link}\n"
                    doc_context_text += f"Bagian: {chunk.get('metadata_header', '')}\n"
                    doc_context_text += f"Isi Teks:\n{chunk.get('content', '')}\n\n"
            
            metrics["source_count"] = len(unique_sources)

            system_instruction = self.policy.get_system_prompt()
            layer_instruction = self.policy.format_layer_instruction(layer)
            
            full_system_prompt = f"{system_instruction}\n\n{layer_instruction}\n\n"
            if layer in [1, 2]:
                full_system_prompt += f"{graph_context_text}\n{doc_context_text}"

            # Format Messages untuk AIGateway
            formatted_contents = [{"role": "system", "content": full_system_prompt}]
            
            for msg in chat_history:
                role = "user" if msg["role"] == "user" else "assistant"
                formatted_contents.append({"role": role, "content": msg["content"]})
                
            formatted_contents.append({"role": "user", "content": query})

            logger.info(f"Memulai LLM Streaming... (Layer: {layer}, Sources: {metrics['source_count']})")
            
            if not self.gateway:
                raise ValueError("Tidak ada AI Gateway yang tersedia.")

            response_stream = await self.gateway.agenerate_llm_response(
                messages=formatted_contents,
                task_name="chatbot",
                temperature=0.3,
                stream=True
            )

            full_response = ""
            input_tokens = len(full_system_prompt) // 4 + sum(len(m["content"]) // 4 for m in chat_history) + len(query) // 4
            output_tokens = 0

            async for chunk in response_stream:
                if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
                    text_chunk = chunk.choices[0].delta.content
                    full_response += text_chunk
                    yield text_chunk

            if layer in [1, 2] and unique_sources and 'title_cache' in locals() and title_cache:
                sources_str = "\n\n**Sumber Bacaan:**\n"
                for h, t in title_cache.items():
                    final_title = t.title()
                    if len(final_title) > 40:
                        final_title = final_title[:40] + "..."
                    html_link = f'<a href="graph://node/{h}">{final_title}</a>'
                    sources_str += f"- {html_link}\n"
                
                full_response += sources_str
                yield sources_str

            if output_tokens == 0:
                output_tokens = len(full_response) // 4

            # 4. L: Learn (Save State & Metrics)
            metrics["latency_ms"] = (time.time() - start_time) * 1000

            references = []
            if layer in [1, 2] and unique_sources:
                for source_hash in unique_sources:
                    title = title_cache.get(source_hash, "Dokumen Akademis") if 'title_cache' in locals() else "Dokumen Akademis"
                    references.append(
                        {
                            "title": title,
                            "url": f"graph://node/{source_hash}",
                            "snippet": "",
                        }
                    )

            self.history_manager.save_turn(
                project_id=project_id,
                user_id=user_id,
                user_query=query,
                ai_response=full_response,
                references=references,
                metadata=metrics,
            )
            
            logger.info(f"[Metrics] Sesi {project_id} selesai: {metrics}")

        except Exception as e:
            logger.error(f"Error pada Orchestrator: {e}", exc_info=True)
            error_msg = f"Maaf, Asisten Peneliti mengalami kendala internal: {str(e)}"
            self.history_manager.save_turn(
                project_id=project_id,
                user_id=user_id,
                user_query=query,
                ai_response=error_msg,
                references=[],
                metadata=metrics,
            )
            yield error_msg
