"""
PERTEMUAN 10: RAG Part 1 - Embedding & Vector Database
Complete Lab Runner: Lab 1 (Chunking), Lab 2 (Indexing & Search), Lab 3 (Hit@k Evaluation)
"""

from pathlib import Path
import random
from pypdf import PdfReader
import chromadb
from chromadb.utils import embedding_functions

# ---------------------------------------------------------
# Utility: PDF Reader & Fixed-Size Splitter
# ---------------------------------------------------------
def baca_pdf(path: Path) -> str:
    """Extracts text from all pages of a PDF."""
    reader = PdfReader(path)
    return "\n".join((p.extract_text() or "") for p in reader.pages)

def potong(teks: str, ukuran: int = 800, tumpang: int = 100) -> list[str]:
    """Fixed-size chunking with overlap."""
    chunks, start = [], 0
    while start < len(teks):
        chunk = teks[start:start + ukuran].strip()
        if chunk:
            chunks.append(chunk)
        start += ukuran - tumpang
    return chunks

# ---------------------------------------------------------
# Lab 1: Chunking Experiment
# ---------------------------------------------------------
def run_lab1(pdf_path: Path):
    print("\n" + "=" * 60)
    print("LAB 1: CHUNKING EXPERIMENT")
    print("=" * 60)

    text = baca_pdf(pdf_path)
    print(f"\n[1] Reading PDF: {pdf_path.name} (Total length: {len(text)} chars)")
    print("-" * 50)
    print("First 500 characters preview:")
    print(text[:500])
    print("-" * 50)

    configs = [(500, 50), (1000, 100), (1500, 150)]
    print("\n[2] Chunk Counts Across Configurations (overlap ~10%):")
    for size, overlap in configs:
        chunks = potong(text, ukuran=size, tumpang=overlap)
        print(f"  - Chunk Size {size} (overlap {overlap}) -> {len(chunks)} chunks")

    print("\n[3] Inspection for Truncated Sentences (3 sample chunks each):")
    random.seed(42)
    for size, overlap in configs:
        chunks = potong(text, ukuran=size, tumpang=overlap)
        sample_count = min(3, len(chunks))
        samples = random.sample(chunks, sample_count)
        print(f"\n--- Chunk Size {size} Samples ({sample_count} samples) ---")
        for i, s in enumerate(samples, 1):
            starts_mid = not s[0].isupper()
            ends_mid = s[-1] not in ".!?"
            print(f"Sample {i} (len {len(s)}):")
            print(f"  Preview: {s[:80]}... [ends with: '{s[-40:]}']")
            print(f"  Cutoff check: Starts mid-sentence={starts_mid}, Ends mid-sentence={ends_mid}")

# ---------------------------------------------------------
# Lab 2: Indexing & Querying
# ---------------------------------------------------------
def setup_collection(client: chromadb.PersistentClient, name: str, emb_fn, chunk_size: int, chunk_overlap: int, docs_dir: Path):
    # Reset collection if exists for clean repeatable runs
    try:
        client.delete_collection(name=name)
    except Exception:
        pass

    col = client.create_collection(
        name=name,
        embedding_function=emb_fn,
        metadata={"hnsw:space": "cosine"}
    )

    for pdf in docs_dir.glob("*.pdf"):
        raw_text = baca_pdf(pdf)
        chunks = potong(raw_text, ukuran=chunk_size, tumpang=chunk_overlap)
        if not chunks:
            continue
        col.upsert(
            ids=[f"{pdf.stem}-sz{chunk_size}-{i}" for i in range(len(chunks))],
            documents=chunks,
            metadatas=[{"sumber": pdf.name, "chunk": i, "size": chunk_size} for i in range(len(chunks))]
        )
    return col

def run_lab2(col):
    print("\n" + "=" * 60)
    print("LAB 2: BUILD INDEX & SEARCH")
    print("=" * 60)

    # 5 standard queries
    queries = [
        "What is the minimum password length required?",
        "What are the core working hours for remote staff?",
        "What is the company policy on harassment and discrimination?",
        "Where must electronic company data be stored?",
        "What are the consequences of violating the Code of Conduct?"
    ]

    print("\n[1] Testing 5 Standard Queries (top-3 results):")
    for q in queries:
        res = col.query(query_texts=[q], n_results=3)
        top_doc = res["documents"][0][0]
        top_meta = res["metadatas"][0][0]
        top_dist = res["distances"][0][0]
        print(f"\nQuery: '{q}'")
        print(f"  -> Top match: [{top_meta['sumber']} | chunk {top_meta['chunk']}] (distance: {top_dist:.4f})")
        print(f"  -> Snippet: {top_doc[:120].replace(chr(10), ' ')}...")

    # 2 paraphrased queries (different wording, semantic match)
    paraphrased = [
        ("Can I log in using a free cafe internet connection?", "Remote work / VPN on public Wi-Fi"),
        ("What is the penalty if I take a second job or outside gig?", "Conflict of interest / outside employment")
    ]

    print("\n[2] Testing 2 Paraphrased Queries (Semantic vs Keyword):")
    for q, intent in paraphrased:
        res = col.query(query_texts=[q], n_results=1)
        top_doc = res["documents"][0][0]
        top_meta = res["metadatas"][0][0]
        top_dist = res["distances"][0][0]
        print(f"\nParaphrased: '{q}' (Target: {intent})")
        print(f"  -> Retrieved: [{top_meta['sumber']} | chunk {top_meta['chunk']}] (distance: {top_dist:.4f})")
        print(f"  -> Snippet: {top_doc[:140].replace(chr(10), ' ')}...")

# ---------------------------------------------------------
# Lab 3: Evaluation (Hit@k)
# ---------------------------------------------------------
# 15 targeted (query, required_keyword) pairs across all 3 PDFs
TEST_SET = [
    ("What is the minimum password length required?", "12 characters"),
    ("How often must employees update their passwords?", "90 days"),
    ("Is multi-factor authentication mandatory?", "MFA"),
    ("Where should suspected security incidents be reported?", "security@nexusglobal.dummy"),
    ("What is the shortcut to lock the screen on Windows?", "Win + L"),
    ("What are the mandatory core hours for remote employees?", "10:00 AM to 3:00 PM"),
    ("How quickly must equipment hardware issues be reported?", "24 hours"),
    ("What is required before connecting to public Wi-Fi?", "VPN"),
    ("Who decides if an employee is eligible for remote work?", "department head"),
    ("What should be done with sensitive paper documents?", "shredded"),
    ("What outside business activities must be disclosed to HR?", "outside employment"),
    ("What is the zero-tolerance policy regarding discrimination?", "zero-tolerance"),
    ("What are examples of intellectual property assets?", "client lists"),
    ("What disciplinary actions can follow policy violations?", "termination"),
    ("Where must electronic corporate data be stored?", "encrypted company servers")
]

def hit_at_k(col, test_set, k: int) -> float:
    """Computes Hit@k: proportion of queries where answer keyword appears in top-k chunks."""
    hits = 0
    for query, keyword in test_set:
        docs = col.query(query_texts=[query], n_results=k)["documents"][0]
        if any(keyword.lower() in d.lower() for d in docs):
            hits += 1
    return hits / len(test_set)

def run_lab3(col_500, col_1000):
    print("\n" + "=" * 60)
    print("LAB 3: MEASURE HIT@K EVALUATION")
    print("=" * 60)

    print(f"Evaluating {len(TEST_SET)} test queries across two chunk size collections:\n")
    print(f"{'Collection':<25} | {'Hit@1':<8} | {'Hit@3':<8} | {'Hit@5':<8}")
    print("-" * 55)

    for name, col in [("Collection (size=500)", col_500), ("Collection (size=1000)", col_1000)]:
        h1 = hit_at_k(col, TEST_SET, k=1)
        h3 = hit_at_k(col, TEST_SET, k=3)
        h5 = hit_at_k(col, TEST_SET, k=5)
        print(f"{name:<25} | {h1:<8.2f} | {h3:<8.2f} | {h5:<8.2f}")

    print("\nDetailed Miss Analysis for size=500 at k=1:")
    for query, keyword in TEST_SET:
        docs = col_500.query(query_texts=[query], n_results=1)["documents"][0]
        if not any(keyword.lower() in d.lower() for d in docs):
            print(f"  - Missed: '{query}' (Target keyword: '{keyword}')")
            print(f"    Top retrieved snippet: {docs[0][:90].replace(chr(10), ' ')}...")

# ---------------------------------------------------------
# Main Execution
# ---------------------------------------------------------
def main():
    docs_dir = Path("dokumen")
    pdf_files = list(docs_dir.glob("*.pdf"))
    if not pdf_files:
        print("Error: No PDF files found in 'dokumen/' folder.")
        return

    # Initialize Embedding Function (runs locally)
    print("Loading multilingual embedding model (paraphrase-multilingual-MiniLM-L12-v2)...")
    emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="paraphrase-multilingual-MiniLM-L12-v2"
    )

    client = chromadb.PersistentClient(path="./chroma_db")

    # Lab 1
    run_lab1(pdf_files[0])

    # Index two collections for comparison (size 500 & size 1000)
    print("\nIndexing documents into ChromaDB collections...")
    col_500 = setup_collection(client, "pedoman_500", emb_fn, chunk_size=500, chunk_overlap=50, docs_dir=docs_dir)
    col_1000 = setup_collection(client, "pedoman_1000", emb_fn, chunk_size=1000, chunk_overlap=100, docs_dir=docs_dir)
    print("Indexed collections 'pedoman_500' and 'pedoman_1000' successfully.")

    # Lab 2 (using col_500)
    run_lab2(col_500)

    # Lab 3
    run_lab3(col_500, col_1000)

if __name__ == "__main__":
    main()