#!/usr/bin/env python3
"""
Build RAG from BKD Sosialisasi documents
"""
import os
import re
import glob

# PDF extraction
import fitz  # pymupdf

# PPTX extraction
from markitdown import MarkItDown

# ChromaDB
import chromadb
from chromadb.config import Settings

# Embeddings
from sentence_transformers import SentenceTransformer

# Progress
from tqdm import tqdm

# Paths
BKD_DIR = "/home/ubuntu/kantor/BKD_Sosialisasi/Materi Sosialisasi 2024"
CHROMA_DIR = "/home/ubuntu/.hermes/rag_bkd_chroma"
os.makedirs(CHROMA_DIR, exist_ok=True)

# --- EXTRACTORS ---

def extract_pdf(pdf_path):
    """Extract text from PDF using pymupdf"""
    try:
        doc = fitz.open(pdf_path)
        texts = []
        for page in doc:
            t = page.get_text("text").strip()
            if t:
                texts.append(t)
        doc.close()
        return "\n".join(texts)
    except Exception as e:
        return f"[ERROR PDF {pdf_path}: {e}]"

def extract_pptx(pptx_path):
    """Extract text from PPTX using markitdown"""
    try:
        md = MarkItDown()
        result = md.convert(pptx_path)
        return result.text_content.strip()
    except Exception as e:
        return f"[ERROR PPTX {pptx_path}: {e}]"

def chunk_text(text, chunk_size=800, overlap=100):
    """Simple chunking with overlap"""
    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]
        # Try to 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 + 200:
                end = start + break_pos + 1
                chunk = text[start:end]
        chunks.append(chunk.strip())
        start = end - overlap
    return [c for c in chunks if c]

# --- COLLECT FILES ---

files = glob.glob(os.path.join(BKD_DIR, "*.pptx")) + \
        glob.glob(os.path.join(BKD_DIR, "*.pdf")) + \
        glob.glob(os.path.join(BKD_DIR, "*.PPTX")) + \
        glob.glob(os.path.join(BKD_DIR, "*.PDF"))

print(f"Found {len(files)} files")

# --- EXTRACT ALL TEXT ---

docs = []  # list of (filename, chunk_text, chunk_index)

for fpath in tqdm(files, desc="Extracting"):
    fname = os.path.basename(fpath)
    print(f"\n  Processing: {fname}")

    if fname.lower().endswith('.pdf'):
        text = extract_pdf(fpath)
    else:
        text = extract_pptx(fpath)

    if not text or len(text) < 50:
        print(f"    Skipping (no text): {fname}")
        continue

    # Quick filter — remove noise lines
    lines = text.split('\n')
    clean_lines = []
    for line in lines:
        # Skip very short lines that are likely formatting artifacts
        if len(line.strip()) < 5 and not line.strip().endswith('.'):
            continue
        clean_lines.append(line)
    text = '\n'.join(clean_lines)

    chunks = chunk_text(text)
    print(f"    -> {len(chunks)} chunks")

    for i, chunk in enumerate(chunks):
        docs.append({
            "id": f"{fname}_{i}",
            "text": chunk,
            "filename": fname,
            "chunk_idx": i,
            "total_chunks": len(chunks)
        })

print(f"\nTotal chunks: {len(docs)}")

# --- BUILD CHROMA DB ---

print("\nLoading embedding model...")
model = SentenceTransformer("all-MiniLM-L6-v2")  # fast, good quality

print("Creating Chroma collection...")
chroma_client = chromadb.PersistentClient(path=CHROMA_DIR)
collection = chroma_client.get_or_create_collection(
    name="bkd_sosialisasi",
    metadata={"description": "Dokumen sosialisasi BKD Kalimantan Selatan"}
)

# Add in batches
batch_size = 100
texts = [d["text"] for d in docs]
ids = [d["id"] for d in docs]
metadatas = [{"filename": d["filename"], "chunk": d["chunk_idx"]} for d in docs]

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

# Store in Chroma
collection.add(
    ids=ids,
    documents=texts,
    embeddings=embeddings.tolist(),
    metadatas=metadatas
)

print(f"\n✅ RAG ready! {collection.count()} chunks stored in {CHROMA_DIR}")
print(f"Collection: 'bkd_sosialisasi'")