"""
RAG Microservice untuk CP 707 Knowledge Base
Berjalan di port 5002, dipanggil oleh backend Node.js
Menggunakan ChromaDB + SentenceTransformers (offline mode)
"""

import os
import sys
from pathlib import Path
from dotenv import load_dotenv

# Set offline agar tidak download dari HuggingFace
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

# Load .env dari direktori yang sama dengan script ini
script_dir = Path(__file__).parent
load_dotenv(script_dir / ".env")

CHROMA_DB_DIR = str(script_dir / os.getenv("CHROMA_DB_DIR", "chroma_db"))
EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL_NAME", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
RAG_PORT = int(os.getenv("RAG_PORT", "5002"))
COLLECTION_NAME = "company_sop"

print(f"[RAG Service] ChromaDB path: {CHROMA_DB_DIR}")
print(f"[RAG Service] Embedding model: {EMBEDDING_MODEL_NAME}")

# Import setelah env set
import chromadb
from sentence_transformers import SentenceTransformer
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# ─── Init ChromaDB & Embedding Model ─────────────────────────────────────────
print("[RAG Service] Memuat ChromaDB...")
try:
    chroma_client = chromadb.PersistentClient(path=CHROMA_DB_DIR)
    collection = chroma_client.get_collection(name=COLLECTION_NAME)
    total_chunks = collection.count()
    print(f"[RAG Service] ChromaDB loaded. Total chunks: {total_chunks}")
except Exception as e:
    print(f"[RAG Service] ERROR: Gagal load ChromaDB: {e}")
    sys.exit(1)

print(f"[RAG Service] Memuat embedding model '{EMBEDDING_MODEL_NAME}'...")
try:
    embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
    print("[RAG Service] Embedding model berhasil dimuat.")
except Exception as e:
    print(f"[RAG Service] ERROR: Gagal load embedding model: {e}")
    print("[RAG Service] Pastikan model sudah didownload. Jalankan populate_db.py terlebih dahulu.")
    sys.exit(1)

# ─── FastAPI App ──────────────────────────────────────────────────────────────
app = FastAPI(
    title="CP707 RAG Service",
    description="Retrieval-Augmented Generation service untuk buku Manajemen Broiler CP 707",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000", "http://localhost:5001", "http://127.0.0.1:5001"],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

# ─── Request/Response Models ─────────────────────────────────────────────────
class QueryRequest(BaseModel):
    query: str = Field(..., min_length=1, description="Query text untuk mencari chunk CP707 relevan")
    n_results: int = Field(default=4, ge=1, le=10, description="Jumlah chunk yang dikembalikan")
    topic: str = Field(default="", description="Topic insight: berat_ayam, fcr, iot_panel, dll")

class QueryResponse(BaseModel):
    success: bool
    chunks: list
    sources: list
    total_found: int
    query_used: str

class HealthResponse(BaseModel):
    status: str
    total_chunks: int
    embedding_model: str

# ─── Topic → Query enhancement mapping ──────────────────────────────────────
# Tambahkan keyword relevan per topic agar embedding search lebih tepat sasaran
TOPIC_QUERY_HINTS = {
    "berat_ayam":    "berat badan target bobot ADG pertumbuhan standar mingguan ayam broiler",
    "fcr":           "FCR feed conversion ratio konsumsi pakan efisiensi standar broiler",
    "iot_panel":     "suhu kandang kelembapan amonia CO2 ventilasi lingkungan pemeliharaan broiler",
    "eef":           "EEF indeks performa IP efisiensi produksi siklus broiler",
    "hitung_ayam":   "mortalitas deplesi kematian afkir populasi standar toleransi broiler",
    "hitung_karung": "pakan karung konsumsi harian feed intake standar broiler",
}

# ─── Endpoints ───────────────────────────────────────────────────────────────
@app.get("/health", response_model=HealthResponse)
def health_check():
    return HealthResponse(
        status="ok",
        total_chunks=collection.count(),
        embedding_model=EMBEDDING_MODEL_NAME,
    )

@app.post("/query", response_model=QueryResponse)
def query_cp707(req: QueryRequest):
    """
    Cari chunk CP707 yang relevan berdasarkan query.
    Jika topic disediakan, tambahkan hint keyword agar hasil lebih relevan.
    """
    enhanced_query = req.query
    if req.topic and req.topic in TOPIC_QUERY_HINTS:
        enhanced_query = f"{req.query} {TOPIC_QUERY_HINTS[req.topic]}"

    try:
        query_embedding = embedding_model.encode([enhanced_query])[0].tolist()

        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=min(req.n_results, collection.count()),
        )

        chunks = results["documents"][0] if results["documents"] else []
        metadatas = results["metadatas"][0] if results["metadatas"] else []

        sources = [
            f"{m.get('source', 'unknown')} (chunk {m.get('chunk_index', '?')})"
            for m in metadatas
        ]

        return QueryResponse(
            success=True,
            chunks=chunks,
            sources=sources,
            total_found=len(chunks),
            query_used=enhanced_query,
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"RAG query error: {str(e)}")


if __name__ == "__main__":
    print(f"[RAG Service] Starting on http://0.0.0.0:{RAG_PORT}")
    uvicorn.run(app, host="0.0.0.0", port=RAG_PORT, log_level="warning")