import os
import re
from typing import Any, Union, List
from langchain_core.runnables import RunnableLambda
from src.config import MODEL_SERVER_URL, MODEL_SERVER_API_KEY, MODEL_SERVER_PATH
from dataclasses import dataclass
import uuid
from dotenv import load_dotenv
from sqlalchemy import Column, Float, Integer, String, ForeignKey, Text, UUID, ARRAY, DateTime, Boolean, JSON
from sqlalchemy.orm import relationship
from src.database.core import Base
from src.models import DefaultMixin, IdentityMixin
from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.llms import VLLM, VLLMOpenAI
from langchain_ollama import ChatOllama
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import PydanticOutputParser, JsonOutputParser
from src.ircfa_interaction.schema import IntentDetection


load_dotenv()

@dataclass
class ClassifierConfig:
    llm = ChatOllama(base_url=os.getenv("OLLAMA_BASE_URL", "http://192.168.1.100:11434"), model="iRCFA-1:30b", temperature=0, top_k=10, top_p=0.9, reasoning=False, format="json")
    parser = PydanticOutputParser(pydantic_object=IntentDetection)
    prompt = ChatPromptTemplate.from_template(
        """Berikan jawaban hanya dalam format JSON dengan field berikut: intent, equipment_name, function_call, confidence.

    Pertanyaan: {question}

    {format_instructions}

    Jawaban:"""
    ).partial(format_instructions=parser.get_format_instructions())
    chain = prompt | llm | parser

class VLLMClassifierConfig:
    """
    Configuration class for VLLM Classifier Agent
    """
    llm = ChatOpenAI(
        openai_api_key="EMPTY",
        openai_api_base=os.getenv("VLLM_BASE_URL", "http://192.168.1.100:8000/v1"),
        model_name="/models/qwen_6_std",
        # model_kwargs={"stop": ["."]},
        temperature=0,
        top_p=0.9,
        # max_tokens=40_000,
    )
    parser = PydanticOutputParser(pydantic_object=IntentDetection)
    prompt = ChatPromptTemplate.from_template(
        """Berikan jawaban hanya dalam format JSON dengan field berikut: intent, equipment_name, function_call, confidence.

    Pertanyaan: {question}

    {format_instructions}
    You must always return valid JSON fenced by a markdown code block. Do not return any additional text.
    Jawaban:"""
    ).partial(format_instructions=parser.get_format_instructions())
    chain = prompt | llm | parser

class LLMConfig:
    temperature = 0
    top_k = 10
    top_p = 0.9
    max_tokens = 2048
    num_predict = 1024
    num_ctx = 16384
    num_gpu = 2
    reasoning = False
    llm = ChatOllama(base_url=os.getenv("OLLAMA_BASE_URL", "http://192.168.1.100:11434"), model="qwen3:8b", temperature=temperature, top_k=top_k, top_p=top_p,  reasoning=reasoning, num_ctx=num_ctx, num_gpu=num_gpu, num_predict=num_predict, max_tokens=max_tokens, stop=["Laporan Selesai", "Done", "Finished"])
    prompt = ChatPromptTemplate([
    ("system", "Anda adalah pakar dan teknisi ahli sebuah pembangkit listrik. Berikan respons yang akurat sesuai dengan permintaan. Jika permintaan pengguna tidak berkaitan dengan peralatan, laporan, Root Cause Failure Analysis (RCFA) atau apapun yang tidak berkaitan dengan topik pembangkit listrik, tolak permintaan pengguna. JANGAN mengulangi data atau tabel yang sudah ditampilkan pada riwayat percakapan sebelumnya kecuali diminta secara eksplisit oleh pengguna. Berikan pertanyaan tindak lanjut yang berkaitan dengan pertanyaan pengguna pada akhir respons. Berikan jawaban dalam tag HTML dan dalam Bahasa Indonesia."),
    ("human", "{question}"),
    ])

class VLLMConfig:
    llm = ChatOpenAI(
        openai_api_key="EMPTY",
        openai_api_base=os.getenv("VLLM_BASE_URL", "http://192.168.1.100:8000/v1"),
        model_name="/models/qwen_6_std",
        # model_kwargs={"stop": ["."]},
        temperature=LLMConfig.temperature,
        top_p=LLMConfig.top_p,
        # max_tokens=LLMConfig.max_tokens,
        streaming=True,
    )

class RCFAContextExtraction(BaseModel):
    symptoms: Any = Field(description="Gejala kegagalan yang diamati")
    characteristics: Any = Field(description="Karakteristik teknis kegagalan")
    failure_modes: Any = Field(description="Mode kegagalan dan analisis 5 Whys")
    cause_action: Any = Field(description="Akar penyebab dan tindakan konfirmasi")
    suggested_questions: List[str] = Field(description="3 pertanyaan tindak lanjut yang relevan untuk membantu user melanjutkan investigasi")

class RCFAContextExtractorConfig:
    llm = ChatOpenAI(
        openai_api_key=MODEL_SERVER_API_KEY,
        openai_api_base=MODEL_SERVER_URL,
        model_name=MODEL_SERVER_PATH,
        temperature=0.3,
        top_p=0.9,
    )
    parser = JsonOutputParser(pydantic_object=RCFAContextExtraction)
    prompt = ChatPromptTemplate.from_template(
        """Anda adalah analis RCFA pakar. Tugas Anda adalah memantau percakapan dan memperbarui state investigasi RCFA.
        
        Berdasarkan riwayat percakapan dan respons terbaru, ekstrak informasi terbaru untuk state berikut:
        1. symptoms: Gejala yang dilaporkan user atau terdeteksi.
        2. characteristics: Karakteristik teknis (vibrasi, suhu, suara, dll).
        3. failure_modes: Mode kegagalan yang mungkin dan ringkasan 5 Whys jika sudah dibahas.
        4. cause_action: Penyebab utama yang dicurigai dan langkah verifikasi/tindakan.
        5. suggested_questions: 1-3 pertanyaan yang MENGARAHKAN user untuk menggunakan KAPABILITAS SISTEM berikut:
        
           **KAPABILITAS SISTEM YANG TERSEDIA:**
           - Prediksi kegagalan equipment (PFI) → contoh: "Cek prediksi kegagalan untuk 3FW-H050"
           - Data sensor real-time & historis → contoh: "Tampilkan data sensor historis 3AH-H501A"  
           - Data reliability (AHI, MTBF) → contoh: "Cek data reliability Boiler Feed Pump saat ini"
           - Pencarian root cause dari database SPLN → contoh: "Apa kemungkinan root cause tube leak pada HPH?"
           - Analisis RCFA spesifik equipment → contoh: "Mulai analisis RCFA untuk Condensate Extraction Pump"
           
           **ATURAN SUGGESTED QUESTIONS:**
           - WAJIB berupa kalimat perintah/pertanyaan yang bisa langsung diketik user (minimal 20 karakter).
           - WAJIB mengarahkan user untuk MENGEKSPLORASI kapabilitas sistem di atas.
           - Jika sudah ada equipment/tag spesifik dalam percakapan, GUNAKAN tag tersebut dalam saran.
           - Jika BELUM ada equipment spesifik, berikan saran yang membantu user memilih equipment, misal: ["Saya ingin analisis RCFA untuk High Pressure Heater", "Cek prediksi kegagalan equipment saat ini", "Tampilkan status reliability Boiler Feed Pump"]
           - DILARANG memparafrase ulang apa yang sudah dikatakan asisten dalam responnya.
           - DILARANG menggunakan heading, label, atau instruksi dari respons asisten (misal: "Deskripsi Masalah:", "Mohon sebutkan equipment").

        PENTING: 
        - JAWAB HANYA DENGAN JSON.
        - DILARANG menyertakan tag <think> atau penjelasan apapun di luar JSON.
        - Gunakan Bahasa Indonesia yang profesional dan padat.
        - Pastikan symptoms, characteristics, failure_modes, dan cause_action adalah STRING.
        - suggested_questions harus berupa ARRAY OF STRINGS (1-3 items, minimal 20 karakter per item).
        - Jika informasi belum tersedia untuk field 1-4, gunakan nilai dari State Sebelumnya.
        - Jika benar-benar tidak ada info, gunakan "Waiting for identification..." atau sejenisnya.

        State Sebelumnya:
        {current_state}

        Riwayat Percakapan & Respons Terbaru:
        {history}

        Jawaban (JSON):"""
    )

    @classmethod
    def get_chain(cls, llm_config: dict = None):
        """
        Returns a chain for RCFA context extraction.
        If llm_config is provided, it uses the specified model.
        Otherwise, it falls back to the default cloud model.
        """
        if not llm_config or llm_config.get("provider") == "n8n":
            return cls.chain
            
        provider = llm_config.get("provider", "openai")
        model_name = llm_config.get("model")
        server = llm_config.get("model_server")
        api_key = llm_config.get("api_key", "not-needed")
        
        # Build LLM based on provider
        if provider == "ollama":
            target_llm = ChatOllama(
                base_url=server, 
                model=model_name, 
                temperature=0.3
            )
        else:
            # Handle openai, vllm, groq, etc. using ChatOpenAI (OpenAI-compatible)
            target_llm = ChatOpenAI(
                openai_api_key=api_key if api_key != "not-needed" else "EMPTY",
                openai_api_base=server,
                model_name=model_name,
                temperature=0.3,
            )
            
        return cls.prompt | target_llm | RunnableLambda(cls.strip_think_blocks) | cls.parser

    @staticmethod
    def strip_think_blocks(text_input) -> str:
        text = text_input.content if hasattr(text_input, 'content') else str(text_input)
        
        # 1. Remove <think>...</think>
        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
        text = re.sub(r'<think>.*', '', text, flags=re.DOTALL)
        
        # 2. Try to find JSON block specifically (looking for { ... })
        # This handles cases where the LLM wraps JSON in markdown ```json or adds extra text
        json_match = re.search(r'(\{[\s\S]*\})', text)
        if json_match:
            candidate = json_match.group(1).strip()
            # Basic validation: check if it has balanced braces roughly
            if candidate.count('{') == candidate.count('}'):
                return candidate
        
        return text.strip()

    chain = prompt | llm | RunnableLambda(strip_think_blocks) | parser

class ChatHistory(Base, DefaultMixin, IdentityMixin):
    __tablename__ = "rcfa_chat_history"

    conversation_id = Column(UUID(as_uuid=True), ForeignKey('rcfa_conversation.id'), nullable=False)
    user_id = Column(UUID(as_uuid=True), nullable=False)
    message = Column(Text, nullable=False)
    attachments = Column(ARRAY(Text), nullable=True)  # Store as comma-separated text
    attachments_filepath = Column(ARRAY(Text), nullable=True)  # Store as comma-separated text array
    response = Column(Text, nullable=False)
    context = Column(ARRAY(Text), nullable=True)
    sources = Column(ARRAY(Text), nullable=True)  # Store as comma-separated text array
    sources_filepath = Column(ARRAY(Text), nullable=True)  # Store as comma-separated text array
    faithfullness = Column(Float, nullable=True)
    answer_relevancy = Column(Float, nullable=True)
    context_recall = Column(Float, nullable=True)
    context_relevancy = Column(Float, nullable=True)
    chat_model = Column(String(100), nullable=False, default="qwen3")
    embedding_model = Column(String(100), nullable=False, default="qwen3")
    ocr_model = Column(String(100), nullable=False, default="marker")
    message_timestamp = Column(DateTime, nullable=False)  # Store as datetime
    response_timestamp = Column(DateTime, nullable=False)  # Store as datetime
    time_taken = Column(Float, nullable=True)  # Store as float for seconds
    prompt_tokens = Column(Float, nullable=True)
    response_tokens = Column(Float, nullable=True)
    tokens_used = Column(Float, nullable=True)
    feedback = Column(String, nullable=True, default=None)
    feedback_comment = Column(Text, nullable=True, default=None)
    feedback_screenshot = Column(Text, nullable=True, default=None)  # base64 compressed image
    last_location_tag = Column(String, nullable=True, default=None)  # New column to store last location_tag
    document_context = Column(Text, nullable=True)  # Store as comma-separated text array
    context_info = Column(JSON, nullable=True)  # Store minimalist context state (symptoms, etc.)
    # uploaded_files = relationship("UploadedFile", back_populates="rag_history", cascade="all, delete-orphan")

    conversation = relationship("Conversation")

class RCFAMemory(Base, DefaultMixin, IdentityMixin):
    __tablename__ = "rcfa_memory"

    # New records are always private and bound to their authenticated owner.
    # The column remains nullable for compatibility with already-created tables.
    user_id = Column(UUID(as_uuid=True), nullable=True)
    conversation_id = Column(UUID(as_uuid=True), nullable=True)
    chat_history_id = Column(UUID(as_uuid=True), nullable=True)
    scope = Column(String(20), nullable=False, default="user")  # user
    source_type = Column(String(50), nullable=False, default="chat_text")
    title = Column(Text, nullable=True)
    summary = Column(Text, nullable=False)
    excerpt = Column(Text, nullable=True)
    filepath = Column(Text, nullable=True)
    location_tag = Column(String, nullable=True)
    equipment = Column(String, nullable=True)
    confidence = Column(Float, nullable=True, default=0.7)
    metadata_json = Column(JSON, nullable=True)

class Conversation(Base, DefaultMixin, IdentityMixin):
    __tablename__ = "rcfa_conversation"

    user_id = Column(UUID(as_uuid=True), nullable=False)
    conversation_id = Column(UUID(as_uuid=True), nullable=False)
    conversation_title = Column(Text, nullable=False)
    is_exported = Column(Boolean, nullable=False, default=False)
    exported_at = Column(DateTime, nullable=True)
    report_name = Column(String, nullable=True)
    report_content = Column(Text, nullable=True)
    rcfa_phase = Column(String, nullable=True, default="1. Failure Prediction")  # Track RCFA Phase
    rcfa_state = Column(JSON, nullable=True)  # Store state checkpoint

    rag_histories = relationship("ChatHistory", back_populates="conversation", cascade="all, delete-orphan")


class ModelConfig(Base, DefaultMixin, IdentityMixin):
    __tablename__ = "rcfa_model_config"

    name = Column(String(100), unique=True, nullable=False)
    display_name = Column(String(200), nullable=False)
    model = Column(String(500), nullable=False)
    model_server = Column(String(500), nullable=False)
    api_key = Column(String(500), nullable=False, default="not-needed")
    provider = Column(String(100), nullable=False, default="openai")
    description = Column(Text, nullable=True)
    stream = Column(Boolean, nullable=False, default=True)
    enabled = Column(Boolean, nullable=False, default=True)
    is_default = Column(Boolean, nullable=False, default=False)
    generate_cfg = Column(JSON, nullable=True)