#!/usr/bin/env python3
"""
Build RAG from DOHI/IOHI document (846 pages)
"""
import os
import fitz  # pymupdf
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

# Paths
DOHI_PATH = "/home/ubuntu/kantor/DOHI_IOHI/Index_Obat_Hewan_Indonesia_IOHI_2014.pdf"
CHROMA_DIR = "/home/ubuntu/.hermes/rag_dohi_chroma"
os.makedirs(CHROMA_DIR, exist_ok=True)

def chunk_text(text, chunk_size=600, overlap=80):
    """Chunk text by sentences, respecting chunk_size"""
    if len(text) < chunk_size:
        return [text] if text.strip() else []
    
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        
        # Break at sentence boundary
        if end < len(text):
            last_period = chunk.rfind('. ')
            last_newline = chunk.rfind('\n')
            break_pos = max(last_period + 1, last_newline)
            if break_pos > start + 150:
                end = start + break_pos + 1
                chunk = text[start:end]
        
        stripped = chunk.strip()
        if stripped:
            chunks.append(stripped)
        
        start = end - overlap
        if start >= len(text):
            break
    
    return chunks

# --- EXTRACT ---
print(f"Opening: {DOHI_PATH}")
doc = fitz.open(DOHI_PATH)
print(f"Total pages: {len(doc)}")

docs = []
page_data = []  # (page_num, chunk_idx, chunk_text)

for page_num in tqdm(range(len(doc)), desc="Extracting pages"):
    page = doc[page_num]
    text = page.get_text("text").strip()
    
    if not text or len(text) < 30:
        continue
    
    # Clean: normalize whitespace
    lines = text.split('\n')
    clean_lines = []
    for line in lines:
        # Skip very short artifact lines
        stripped = line.strip()
        if len(stripped) < 3 and stripped and not stripped.endswith('.'):
            continue
        clean_lines.append(stripped)
    clean_text = ' '.join(clean_lines)
    
    if len(clean_text) < 30:
        continue
    
    chunks = chunk_text(clean_text)
    for i, chunk in enumerate(chunks):
        page_data.append((page_num + 1, i, chunk))

doc.close()
print(f"Total chunks: {len(page_data)}")

# --- STORE IN CHROMA ---
print("\nLoading embedding model...")
model = SentenceTransformer("all-MiniLM-L6-v2")

print("Creating Chroma collection...")
chroma_client = chromadb.PersistentClient(path=CHROMA_DIR)

# Recreate to avoid duplicates if rerunning
try:
    chroma_client.delete_collection("dohi_iohi")
except:
    pass

collection = chroma_client.get_or_create_collection(
    name="dohi_iohi",
    metadata={"description": "Index Obat Hewan Indonesia (IOHI) Edisi IX 2014 - DOHI"}
)

ids = [f"page_{pg}_chunk_{i}" for pg, i, _ in page_data]
texts = [text for _, _, text in page_data]
metadatas = [{"filename": "Index_Obat_Hewan_Indonesia_IOHI_2014.pdf", "halaman": pg} for pg, _, _ in page_data]

print(f"Embedding {len(docs)} chunks...")
batch_size = 64
embeddings = model.encode(texts, show_progress_bar=True, batch_size=batch_size)

collection.add(
    ids=ids,
    documents=texts,
    embeddings=embeddings.tolist(),
    metadatas=metadatas
)

print(f"\n✅ RAG DOHI ready!")
print(f"   Collection: 'dohi_iohi'")
print(f"   Path: {CHROMA_DIR}")
print(f"   Chunks: {collection.count()}")