import os
import sys
import requests
from dotenv import load_dotenv
import chromadb
from sentence_transformers import SentenceTransformer

# Set offline mode agar sentence-transformers tidak mencoba menghubungi Hugging Face di jaringan on-premise
os.environ["HF_HUB_OFFLINE"] = "1"

# Load environment variables
load_dotenv()

CHROMA_DB_DIR = os.getenv("CHROMA_DB_DIR", "chroma_db")
EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL_NAME", "all-MiniLM-L6-v2")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "http://localhost:1234/v1")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "lm-studio")
LLM_MODEL_NAME = os.getenv("LLM_MODEL_NAME", "deepseek-r1-distill-qwen-7b")

def main():
    # 1. Inisialisasi Database Vektor
    if not os.path.exists(CHROMA_DB_DIR):
        print(f"Database Chroma DB di '{CHROMA_DB_DIR}' tidak ditemukan. Silakan jalankan extract_text.py dan populate_db.py terlebih dahulu.")
        return

    print("Menghubungkan ke Chroma DB...")
    chroma_client = chromadb.PersistentClient(path=CHROMA_DB_DIR)
    try:
        collection = chroma_client.get_collection(name="company_sop")
    except Exception as e:
        print(f"Koleksi 'company_sop' tidak ditemukan di database. Pastikan populate_db.py sudah dijalankan dengan sukses. Error: {e}")
        return

    # 2. Inisialisasi Model Embedding lokal
    print(f"Memuat model embedding lokal '{EMBEDDING_MODEL_NAME}' untuk kueri...")
    embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
    print("Model embedding berhasil dimuat.")

    # 3. Konfigurasi koneksi LM Studio native v1 API
    LM_STUDIO_API_URL = os.getenv("LM_STUDIO_API_URL")
    if not LM_STUDIO_API_URL:
        openai_base = os.getenv("OPENAI_BASE_URL", "http://localhost:1234/v1")
        if openai_base.endswith("/v1"):
            LM_STUDIO_API_URL = openai_base.replace("/v1", "/api/v1/chat")
        else:
            LM_STUDIO_API_URL = f"{openai_base.rstrip('/')}/api/v1/chat"
    
    print(f"Mengonfigurasi koneksi LLM ke native API: {LM_STUDIO_API_URL} (Model: {LLM_MODEL_NAME})...")

    print("\n" + "="*60)
    print(" PIPELINE RAG LOKAL - ASISTEN SOP PERUSAHAAN (DEEPSEEK)")
    print(" Ketik 'keluar' atau 'exit' untuk menyudahi percakapan.")
    print("="*60 + "\n")

    while True:
        try:
            query = input("\nPertanyaan Anda: ").strip()
            if not query:
                continue
            if query.lower() in ["keluar", "exit", "q", "quit"]:
                print("Sampai jumpa!")
                break

            print("\n[1/3] Mencari dokumen referensi relevan di database lokal...", end="", flush=True)
            # Buat embedding kueri
            query_embedding = embedding_model.encode([query])[0].tolist()
            
            # Cari kueri di Chroma DB (ambil 6 chunk teratas)
            results = collection.query(
                query_embeddings=[query_embedding],
                n_results=6
            )
            print(" Selesai!")

            retrieved_chunks = results['documents'][0]
            retrieved_metadatas = results['metadatas'][0]

            if not retrieved_chunks or len(retrieved_chunks) == 0:
                print("⚠️ Tidak ditemukan referensi dokumen yang cocok dengan pertanyaan Anda.")
                continue

            # Tampilkan referensi yang ditemukan
            print("\n[Referensi yang Ditemukan]:")
            for idx, meta in enumerate(retrieved_metadatas):
                print(f"  - [{idx+1}] File: {meta['source']} (Chunk: {meta['chunk_index']})")

            # 4. Susun Prompt dengan Konteks SOP
            context = "\n\n---\n\n".join(retrieved_chunks)
            
            system_prompt = (
                "Anda adalah asisten AI perusahaan yang profesional. Tugas Anda adalah memberikan jawaban "
                "yang valid, akurat, dan sesuai dengan Standar Operasional Prosedur (SOP) atau dokumen acuan perusahaan "
                "yang disediakan di bawah ini.\n"
                "Patuhi aturan berikut:\n"
                "1. Jawablah HANYA berdasarkan informasi yang ada dalam dokumen acuan di bawah.\n"
                "2. Jika jawaban tidak dapat ditemukan di dalam dokumen tersebut secara eksplisit atau logis, katakan dengan sopan "
                "bahwa 'Maaf, informasi tersebut tidak ditemukan dalam dokumen SOP/acuan perusahaan kami.' Jangan mengarang informasi.\n"
                "3. Sajikan data dengan valid dan rapi."
            )

            user_prompt = f"""Dokumen SOP / Acuan Perusahaan:
=========================================
{context}
=========================================

Pertanyaan Pengguna: {query}

Jawaban berdasarkan Dokumen Acuan:"""

            print(f"\n[2/3] Menghubungi LLM DeepSeek lokal di {LM_STUDIO_API_URL}...", end="", flush=True)
            
            # Panggil LM Studio native v1 API
            headers = {
                "Content-Type": "application/json"
            }
            if OPENAI_API_KEY and OPENAI_API_KEY != "lm-studio":
                headers["Authorization"] = f"Bearer {OPENAI_API_KEY}"

            payload = {
                "model": LLM_MODEL_NAME,
                "input": user_prompt,
                "system_prompt": system_prompt,
                "temperature": 0.1,
                "max_output_tokens": 32000
            }

            response = requests.post(LM_STUDIO_API_URL, json=payload, headers=headers)
            response.raise_for_status()
            print(" Selesai!")

            answer = response.json()["response"]
            
            # Pisahkan proses berpikir (<think>) jika ada (khusus model reasoning seperti DeepSeek R1)
            import re
            think_match = re.search(r'<think>(.*?)</think>', answer, re.DOTALL)
            clean_answer = re.sub(r'<think>.*?</think>', '', answer, flags=re.DOTALL).strip()
            
            if think_match and think_match.group(1).strip():
                print(" Selesai!")
                print("\n[Proses Berpikir DeepSeek]:")
                print("." * 50)
                print(think_match.group(1).strip())
                print("." * 50)
            else:
                print(" Selesai!")

            print("\n[3/3] Respon Asisten SOP:")
            print("-"*50)
            print(clean_answer)
            print("-"*50)

        except Exception as e:
            print(f"\n❌ Terjadi kesalahan: {e}")
            print("Harap pastikan server LLM lokal Anda (LM Studio/vLLM/llama.cpp) sedang berjalan dan dapat diakses.")

if __name__ == "__main__":
    main()