#!/usr/bin/env python3
"""
Rebuild BKD RAG with BETTER extraction:
- PDF: fitz (pymupdf)
- PPTX: python-pptx shapes + markitdown fallback
"""
import os
import re
import glob
import fitz  # pymupdf
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from pptx import Presentation
from markitdown import MarkItDown
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)

CHUNK_SIZE = 600
OVERLAP = 80

def chunk_text(text, chunk_size=CHUNK_SIZE, overlap=OVERLAP):
    if len(text) < chunk_size:
        return [text.strip()] if text.strip() else []
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        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

def extract_pptx_text_shapes(path):
    """Extract text from PPTX using python-pptx shapes + markitdown fallback"""
    texts = []
    
    # Method 1: python-pptx shapes
    try:
        prs = Presentation(path)
        for slide_idx, slide in enumerate(prs.slides):
            slide_texts = []
            for shape in slide.shapes:
                if shape.has_text_frame:
                    for para in shape.text_frame.paragraphs:
                        t = para.text.strip()
                        if t and len(t) > 2 and not t.startswith("Picture"):
                            slide_texts.append(t)
            if slide_texts:
                combined = ' '.join(slide_texts)
                if len(combined) > 20:
                    texts.append(combined)
    except Exception as e:
        print(f"    pptx shapes error: {e}")
    
    # Method 2: markitdown fallback
    try:
        md = MarkItDown()
        result = md.convert(path)
        raw = result.text_content
        # Split by slide markers and extract meaningful content
        slides_raw = re.split(r'<!-- Slide number:', raw)
        for slide_chunk in slides_raw:
            cleaned = re.sub(r'!\[\]\(Picture\d+\.jpg\)', '', slide_chunk)
            cleaned = re.sub(r'!\[\]\(.*?\)', '', cleaned)
            cleaned = re.sub(r'\[.*?\]\(.*?\)', '', cleaned)
            cleaned = re.sub(r'\n+', ' ', cleaned)
            cleaned = cleaned.strip()
            if len(cleaned) > 50:
                texts.append(cleaned)
    except Exception as e:
        print(f"    markitdown error: {e}")
    
    return texts

def extract_pdf_text(path):
    texts = []
    doc = fitz.open(path)
    for page in doc:
        text = page.get_text("text").strip()
        if text and len(text) > 30:
            # Normalize whitespace
            text = re.sub(r'\s+', ' ', text)
            texts.append(text)
    doc.close()
    return texts

# --- COLLECT ALL DOCUMENT TEXT ---
print("Collecting documents...")
doc_data = []  # (filename, chunks_list)

files = sorted(glob.glob(os.path.join(BKD_DIR, "*")))
for filepath in tqdm(files, desc="Processing files"):
    fname = os.path.basename(filepath)
    if not (fname.endswith('.pdf') or fname.endswith('.pptx')):
        continue
    
    print(f"\n  Processing: {fname}")
    all_texts = []
    
    if filepath.endswith('.pdf'):
        texts = extract_pdf_text(filepath)
        all_texts.extend(texts)
    elif filepath.endswith('.pptx'):
        texts = extract_pptx_text_shapes(filepath)
        all_texts.extend(texts)
    
    print(f"    Extracted {len(all_texts)} text blocks")
    
    # Chunk each text block
    all_chunks = []
    for text in all_texts:
        chunks = chunk_text(text)
        all_chunks.extend(chunks)
    
    print(f"    → {len(all_chunks)} chunks")
    if all_chunks:
        doc_data.append((fname, all_chunks))

print(f"\nTotal documents: {len(doc_data)}")

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

print("Creating ChromaDB...")
chroma_client = chromadb.PersistentClient(path=CHROMA_DIR)
try:
    chroma_client.delete_collection("bkd_sosialisasi")
except:
    pass

collection = chroma_client.get_or_create_collection(
    name="bkd_sosialisasi",
    metadata={"description": "BKD Sosialisasi Dokumen Kepegawaian 2024"}
)

# Build batch
ids = []
documents = []
embeddings_list = []
metadatas = []

chunk_counter = 0
for fname, chunks in doc_data:
    for chunk in chunks:
        chunk_counter += 1
        chunk_id = f"chunk_{chunk_counter}"
        ids.append(chunk_id)
        documents.append(chunk)
        metadatas.append({"filename": fname})
        
        if chunk_counter % 500 == 0:
            print(f"  Embedding batch {chunk_counter}...")

print(f"Total chunks to embed: {len(documents)}")

# Embed in batches
batch_size = 64
all_embeddings = []
for i in tqdm(range(0, len(documents), batch_size), desc="Embedding"):
    batch = documents[i:i+batch_size]
    emb = model.encode(batch, show_progress_bar=False, batch_size=batch_size)
    all_embeddings.extend(emb.tolist())

collection.add(
    ids=ids,
    documents=documents,
    embeddings=all_embeddings,
    metadatas=metadatas
)

print(f"\n✅ RAG BKD v2 ready!")
print(f"   Collection: 'bkd_sosialisasi'")
print(f"   Total chunks: {collection.count()}")
print(f"   Chroma path: {CHROMA_DIR}")