import os
import json
import pickle
from pathlib import Path

from sentence_transformers import SentenceTransformer
import numpy as np
from PyPDF2 import PdfReader

DATA_DIR = Path("data")
PDF_PATH = DATA_DIR / "Permenpar 6 Tahun 2025.pdf"

OUT_TEXT = DATA_DIR / "permenpar_text.txt"
OUT_CHUNKS = DATA_DIR / "permenpar_chunks.json"
OUT_VECTORS = DATA_DIR / "permenpar_vectors.pkl"

MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"


def load_pdf(path: Path):
    print(f"[INFO] Membaca PDF: {path}")
    reader = PdfReader(str(path))
    text = ""

    for i, page in enumerate(reader.pages):
        try:
            txt = page.extract_text()
            if txt:
                text += txt + "\n\n"
        except:
            continue

    return text


def chunk_text(text, max_chars=1200):
    """
    Potong isi PDF menjadi chunk kecil
    Agar pencarian lebih akurat & embedding efisien
    """
    print("[INFO] Memotong teks menjadi chunk...")
    chunks = []
    current = ""

    for paragraph in text.split("\n"):
        if len(current) + len(paragraph) < max_chars:
            current += paragraph + "\n"
        else:
            chunks.append(current.strip())
            current = paragraph + "\n"

    if current.strip():
        chunks.append(current.strip())

    print(f"[INFO] Total chunk: {len(chunks)}")
    return chunks


def main():
    DATA_DIR.mkdir(exist_ok=True)

    # Step 1 — Load PDF
    if not PDF_PATH.exists():
        print(f"[ERROR] PDF tidak ditemukan: {PDF_PATH}")
        return

    full_text = load_pdf(PDF_PATH)

    with open(OUT_TEXT, "w", encoding="utf-8") as f:
        f.write(full_text)

    # Step 2 — Chunking
    chunks = chunk_text(full_text)

    with open(OUT_CHUNKS, "w", encoding="utf-8") as f:
        json.dump(chunks, f, indent=2, ensure_ascii=False)

    # Step 3 — Embeddings
    print("[INFO] Loading model:", MODEL_NAME)
    model = SentenceTransformer(MODEL_NAME)

    print("[INFO] Embedding chunk...")
    embeddings = model.encode(chunks, show_progress_bar=True)
    embeddings = np.array(embeddings)

    with open(OUT_VECTORS, "wb") as f:
        pickle.dump(embeddings, f)

    print("\n=== SELESAI ===")
    print(f"Teks     → {OUT_TEXT}")
    print(f"Chunks   → {OUT_CHUNKS}")
    print(f"Vectors  → {OUT_VECTORS}")
    print("\nKnowledge base siap digunakan oleh API!")


if __name__ == "__main__":
    main()
