"""
HBL Production PDF Extractor
----------------------------
PDF ingestion pipeline for RAG.
Input:
    /pdfs/*.pdf
Output:
    extracted/
        documents.json
        chunks.json
        markdown/
Pipeline:
PDF
 |
Docling layout extraction
 |
OCR fallback
 |
Document classification
 |
Metadata enrichment
 |
Semantic chunks
 |
JSON ready for embeddings
"""

import os

# Raw strings so backslashes are never silently mis-parsed
os.environ["HF_HOME"] = r"D:\hf_cache"
os.environ["HF_HUB_DISABLE_XET"] = "1"
os.environ["HF_HUB_OFFLINE"] = "1"

import logging
import json
import gc
import time
import hashlib
from pathlib import Path
from datetime import datetime

from tqdm import tqdm

from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions

import fitz  # pymupdf

from langdetect import detect

logging.basicConfig(level=logging.INFO)

# --------------------------------------------------
# CONFIG
# --------------------------------------------------

PDF_FOLDER = Path("./classified_pdfs/text_pdfs")

OUTPUT_FOLDER = Path("./extracted_textpdfs")

MARKDOWN_FOLDER = OUTPUT_FOLDER / "markdown"

DOCUMENT_JSON = OUTPUT_FOLDER / "documents.json"

CHUNKS_JSON = OUTPUT_FOLDER / "chunks.json"

# Checkpoint files let the script resume after a crash instead of
# reprocessing everything from scratch.
DOCUMENTS_CKPT = OUTPUT_FOLDER / "documents_checkpoint.jsonl"
CHUNKS_CKPT = OUTPUT_FOLDER / "chunks_checkpoint.jsonl"
FAILED_LOG = OUTPUT_FOLDER / "failed.log"

CHUNK_SIZE = 800
CHUNK_OVERLAP = 150


# --------------------------------------------------
# INIT DOCLING
# --------------------------------------------------

print("Initialising Docling models....")
pdf_options = PdfPipelineOptions()

pdf_options.artifacts_path = r"D:\hf_cache\docling_artifacts"
pdf_options.do_ocr = False
pdf_options.do_table_structure = False
pdf_options.generate_page_images = False
pdf_options.generate_picture_images = False

converter = DocumentConverter(
    format_options={
        InputFormat.PDF: PdfFormatOption(
            pipeline_options=pdf_options
        )
    }
)

print("Docling ready, starting extraction now.....")


# --------------------------------------------------
# HASH
# --------------------------------------------------

def file_hash(path):
    md5 = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            md5.update(chunk)
    return md5.hexdigest()


# --------------------------------------------------
# PDF METADATA
# --------------------------------------------------

def get_pdf_metadata(path):
    doc = fitz.open(path)

    images = 0
    pages = len(doc)

    for p in doc:
        images += len(p.get_images())

    doc.close()

    return {
        "pages": pages,
        "images": images,
        "has_images": images > 0
    }


# --------------------------------------------------
# CLASSIFIER
# --------------------------------------------------

def classify_document(text):
    t = text.lower()

    if any(x in t for x in [
        "application form",
        "account opening",
        "signature",
        "cnic",
        "customer information"
    ]):
        return "FORM"

    if any(x in t for x in [
        "circular",
        "notification",
        "effective date"
    ]):
        return "CIRCULAR"

    if any(x in t for x in [
        "policy",
        "procedure",
        "guidelines"
    ]):
        return "POLICY"

    if any(x in t for x in [
        "rate",
        "pricing",
        "charges",
        "fee"
    ]):
        return "RATE_SHEET"

    return "GENERAL"


# --------------------------------------------------
# FORM DETECTION
# --------------------------------------------------

def detect_form(pdf_meta, text):
    keywords = [
        "signature",
        "tick",
        "checkbox",
        "applicant name",
        "date of birth"
    ]

    score = 0

    for k in keywords:
        if k in text.lower():
            score += 1

    if pdf_meta["has_images"]:
        score += 1

    return score >= 3


# --------------------------------------------------
# LANGUAGE
# --------------------------------------------------

def detect_language(text):
    try:
        return detect(text[:1000])
    except Exception:
        return "unknown"


# --------------------------------------------------
# EXTRACT PDF
# --------------------------------------------------

def extract_pdf(pdf_path):
    print(f"\nProcessing {pdf_path.name}")

    print(f"\n Starting conversion for {pdf_path.name}...")
    start = time.time()
    result = converter.convert(pdf_path)
    print(f"Conversion completed in {time.time() - start:.2f} seconds.")

    start = time.time()
    markdown = result.document.export_to_markdown()
    print(f"Markdown export completed in {time.time() - start:.2f} seconds")

    pdf_meta = get_pdf_metadata(pdf_path)

    metadata = {
        "filename": pdf_path.name,
        "hash": file_hash(pdf_path),
        "size_bytes": pdf_path.stat().st_size,
        "pages": pdf_meta["pages"],
        "extracted_at": datetime.utcnow().isoformat(),
        "language": detect_language(markdown),
        "document_type": classify_document(markdown),
        "is_form": detect_form(pdf_meta, markdown)
    }

    return metadata, markdown


# --------------------------------------------------
# CHUNKER
# --------------------------------------------------

def create_chunks(text, metadata):
    words = text.split()

    chunks = []
    start = 0
    chunk_id = 0

    while start < len(words):
        end = start + CHUNK_SIZE
        chunk_words = words[start:end]
        chunk = " ".join(chunk_words)

        chunks.append({
            "chunk_id": chunk_id,
            "text": chunk,
            "metadata": metadata
        })

        chunk_id += 1
        start = end - CHUNK_OVERLAP

    return chunks


# --------------------------------------------------
# CHECKPOINT HELPERS
# --------------------------------------------------

def load_processed_hashes():
    """Read the checkpoint file to find which PDFs are already done,
    so a crash + rerun doesn't reprocess them."""
    processed = set()
    if DOCUMENTS_CKPT.exists():
        with open(DOCUMENTS_CKPT, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    doc = json.loads(line)
                    processed.add(doc["hash"])
                except Exception:
                    continue
    return processed


def append_jsonl(path, obj):
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(obj, ensure_ascii=False) + "\n")


def rebuild_final_json_from_checkpoints():
    """Convert the append-only checkpoint files into the final
    documents.json / chunks.json the rest of the pipeline expects."""
    documents = []
    if DOCUMENTS_CKPT.exists():
        with open(DOCUMENTS_CKPT, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    documents.append(json.loads(line))

    chunks = []
    if CHUNKS_CKPT.exists():
        with open(CHUNKS_CKPT, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    chunks.append(json.loads(line))

    DOCUMENT_JSON.write_text(
        json.dumps(documents, indent=2, ensure_ascii=False),
        encoding="utf-8"
    )

    CHUNKS_JSON.write_text(
        json.dumps(chunks, indent=2, ensure_ascii=False),
        encoding="utf-8"
    )

    return documents, chunks


# --------------------------------------------------
# MAIN
# --------------------------------------------------

def main():
    OUTPUT_FOLDER.mkdir(exist_ok=True)
    MARKDOWN_FOLDER.mkdir(exist_ok=True)

    pdfs = list(PDF_FOLDER.glob("*.pdf"))
    print(f"Found {len(pdfs)} PDFs")

    # Resume support: skip PDFs whose hash is already in the checkpoint.
    seen_hashes = load_processed_hashes()
    if seen_hashes:
        print(f"Resuming: {len(seen_hashes)} PDFs already processed, skipping those.")

    for pdf in tqdm(pdfs):
        try:
            # Hash first, before the expensive Docling conversion,
            # so duplicates and already-processed files cost almost nothing.
            h = file_hash(pdf)

            if h in seen_hashes:
                print("Already processed / duplicate, skipping", pdf.name)
                continue

            meta, markdown = extract_pdf(pdf)

            # extract_pdf recomputes the hash internally; keep them consistent
            meta["hash"] = h

            md_file = MARKDOWN_FOLDER / (pdf.stem + ".md")
            md_file.write_text(markdown, encoding="utf-8")
            meta["markdown_file"] = str(md_file)

            doc_entry = {**meta, "text_length": len(markdown)}

            # Write to checkpoint immediately so a crash on the NEXT
            # file doesn't lose this one's results.
            append_jsonl(DOCUMENTS_CKPT, doc_entry)

            if not meta["is_form"]:
                for chunk in create_chunks(markdown, meta):
                    append_jsonl(CHUNKS_CKPT, chunk)

            seen_hashes.add(h)

        except Exception as e:
            print("FAILED", pdf.name, e)
            with open(FAILED_LOG, "a", encoding="utf-8") as f:
                f.write(f"{datetime.utcnow().isoformat()} {pdf.name} {e}\n")
        finally:
            gc.collect()

    documents, chunks = rebuild_final_json_from_checkpoints()

    print("\nDONE")
    print("Documents:", len(documents))
    print("Chunks:", len(chunks))


if __name__ == "__main__":
    main()