#!/usr/bin/env python3
"""Extract large NIST PDFs locally, then upload as .txt to RAG service.

Workaround for embedding-service container crashing on large PDFs (>5 MB).
Runs pdfplumber + normalization locally, uploads extracted text as .txt.

Usage (on Mac Mini):
    python3 control-pipeline/scripts/extract_and_upload_nist.py
"""

import json
import os
import re
import sys
import tempfile
import unicodedata

import httpx
import pdfplumber

RAG_URL = "https://localhost:8097"
QDRANT_URL = "http://localhost:6333"

DOCS = [
    {
        "object_name": "compliance/bund/compliance/2026/NIST_SP_800_53r5.pdf",
        "collection": "bp_compliance_datenschutz",
        "filename": "NIST_SP_800_53r5.txt",
        "extra_metadata": {
            "regulation_id": "nist_sp800_53r5",
            "source_id": "nist",
            "doc_type": "controls_catalog",
            "guideline_name": "NIST SP 800-53 Rev. 5 Security and Privacy Controls",
            "license": "public_domain_us_gov",
            "attribution": "NIST",
            "source": "nist.gov",
        },
    },
    {
        "object_name": "compliance/bund/compliance/2026/nist_sp_800_82r3.pdf",
        "collection": "bp_compliance_ce",
        "filename": "nist_sp_800_82r3.txt",
        "extra_metadata": {
            "regulation_id": "nist_sp_800_82r3",
            "regulation_name_de": "NIST SP 800-82 Rev. 3 — Guide to OT Security",
            "regulation_name_en": "NIST SP 800-82 Rev. 3 — Guide to OT Security",
            "regulation_short": "NIST SP 800-82",
            "category": "ot_security",
            "license": "public_domain_us",
            "source": "nist.gov",
        },
    },
    {
        "object_name": "compliance/bund/compliance/2026/nist_sp_800_160v1r1.pdf",
        "collection": "bp_compliance_ce",
        "filename": "nist_sp_800_160v1r1.txt",
        "extra_metadata": {
            "regulation_id": "nist_sp_800_160v1r1",
            "regulation_name_de": "NIST SP 800-160 Vol. 1 Rev. 1",
            "regulation_name_en": "NIST SP 800-160 Vol. 1 Rev. 1",
            "regulation_short": "NIST SP 800-160",
            "category": "security_engineering",
            "license": "public_domain_us",
            "source": "nist.gov",
        },
    },
    {
        "object_name": "compliance/bund/compliance/2026/NIST_SP_800_207.pdf",
        "collection": "bp_compliance_datenschutz",
        "filename": "NIST_SP_800_207.txt",
        "extra_metadata": {
            "regulation_id": "nist_sp800_207",
            "source_id": "nist",
            "doc_type": "architecture",
            "guideline_name": "NIST SP 800-207 Zero Trust Architecture",
            "license": "public_domain_us_gov",
            "attribution": "NIST",
            "source": "nist.gov",
        },
    },
]


def normalize_pdf_text(text: str) -> str:
    """Fix broken spacing from multi-column PDF extraction."""
    text = unicodedata.normalize('NFKC', text)
    text = text.replace('\u00ad', '').replace('\u200b', '')
    prev = None
    while prev != text:
        prev = text
        text = re.sub(r'(\d+)\s+\.\s+(\d+)', r'\1.\2', text)
    text = re.sub(r'\b([A-Z]{2,4})\s+-\s+(\d+)\b', r'\1-\2', text)
    text = re.sub(
        r'\b([A-Z]{2})\s*\.\s*([A-Z]{2})\s*-\s*(\d{2})\b', r'\1.\2-\3', text
    )
    text = re.sub(r'\(\s+(\d+)\s+\)', r'(\1)', text)
    text = re.sub(r'[^\S\n]{2,}', ' ', text)
    return text


def extract_pdf_locally(pdf_bytes: bytes) -> str:
    """Extract text from PDF using pdfplumber with normalization."""
    import io
    text_parts = []
    with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf:
        print(f"    Pages: {len(pdf.pages)}")
        for i, page in enumerate(pdf.pages):
            text = page.extract_text(x_tolerance=3, y_tolerance=4)
            if text:
                text_parts.append(text)
            if (i + 1) % 50 == 0:
                print(f"    Extracted {i + 1}/{len(pdf.pages)} pages...")
    raw = "\n\n".join(text_parts)
    return normalize_pdf_text(raw)


def download_from_minio(object_name: str) -> bytes:
    """Download file from MinIO via RAG service."""
    with httpx.Client(timeout=60.0, verify=False) as c:
        resp = c.get(f"{RAG_URL}/api/v1/documents/download/{object_name}")
        resp.raise_for_status()
        url = resp.json()["url"]
    with httpx.Client(timeout=300.0, verify=False) as c:
        resp = c.get(url)
        resp.raise_for_status()
        return resp.content


def upload_text(
    text: str, filename: str, collection: str, extra_metadata: dict,
) -> dict:
    """Upload extracted text to RAG service as .txt."""
    form_data = {
        "collection": collection,
        "data_type": "compliance",
        "bundesland": "bund",
        "use_case": "compliance",
        "year": "2026",
        "chunk_strategy": "recursive",
        "chunk_size": "1500",
        "chunk_overlap": "100",
        "metadata_json": json.dumps(extra_metadata, ensure_ascii=False),
    }
    text_bytes = text.encode("utf-8")
    with httpx.Client(timeout=1800.0, verify=False) as c:
        resp = c.post(
            f"{RAG_URL}/api/v1/documents/upload",
            files={"file": (filename, text_bytes, "text/plain")},
            data=form_data,
        )
        resp.raise_for_status()
        return resp.json()


def count_chunks(collection: str, regulation_id: str) -> int:
    """Count chunks for a regulation in Qdrant."""
    with httpx.Client(timeout=30.0) as c:
        resp = c.post(
            f"{QDRANT_URL}/collections/{collection}/points/count",
            json={
                "filter": {
                    "must": [{
                        "key": "regulation_id",
                        "match": {"value": regulation_id},
                    }]
                },
                "exact": True,
            },
        )
        resp.raise_for_status()
        return resp.json()["result"]["count"]


def check_section_rate(collection: str, regulation_id: str) -> tuple:
    """Returns (total_chunks, chunks_with_section)."""
    total = 0
    with_sec = 0
    offset = None
    with httpx.Client(timeout=60.0) as c:
        while True:
            body = {
                "filter": {
                    "must": [{
                        "key": "regulation_id",
                        "match": {"value": regulation_id},
                    }]
                },
                "limit": 100,
                "with_payload": ["section"],
            }
            if offset is not None:
                body["offset"] = offset
            resp = c.post(
                f"{QDRANT_URL}/collections/{collection}/points/scroll",
                json=body,
            )
            resp.raise_for_status()
            data = resp.json()["result"]
            for pt in data["points"]:
                total += 1
                s = pt.get("payload", {}).get("section", "")
                if s and s.strip():
                    with_sec += 1
            offset = data.get("next_page_offset")
            if offset is None:
                break
    return total, with_sec


def main():
    print("=" * 60)
    print("NIST PDF Local Extraction + Upload")
    print("=" * 60)

    results = []
    for i, doc in enumerate(DOCS, 1):
        reg_id = doc["extra_metadata"]["regulation_id"]
        print(f"\n[{i}/{len(DOCS)}] {doc['filename']} → {doc['collection']}")

        # 1. Check current state
        existing = count_chunks(doc["collection"], reg_id)
        print(f"  Existing chunks: {existing}")

        # 2. Download PDF from MinIO
        print(f"  Downloading from MinIO...")
        pdf_bytes = download_from_minio(doc["object_name"])
        print(f"  Downloaded {len(pdf_bytes) / 1024 / 1024:.1f} MB")

        # 3. Extract text locally with pdfplumber
        print(f"  Extracting text locally...")
        text = extract_pdf_locally(pdf_bytes)
        print(f"  Extracted {len(text):,} chars, {text.count(chr(10)):,} lines")

        # 4. Save extracted text temporarily (for debugging)
        tmp_path = f"/tmp/nist_{reg_id}.txt"
        with open(tmp_path, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"  Saved to {tmp_path}")

        # 5. Upload as .txt
        print(f"  Uploading as .txt to RAG service...")
        result = upload_text(text, doc["filename"], doc["collection"],
                             doc["extra_metadata"])
        new_chunks = result.get("chunks_count", 0)
        new_doc_id = result.get("document_id", "")
        print(f"  Uploaded: {new_chunks} chunks (doc_id={new_doc_id})")

        # 6. Check section rate
        if new_chunks > 0:
            total, with_sec = check_section_rate(doc["collection"], reg_id)
            pct = (with_sec / total * 100) if total > 0 else 0
            print(f"  Section rate: {with_sec}/{total} ({pct:.0f}%)")
        else:
            pct = 0
            print("  WARNING: 0 chunks created!")

        results.append({
            "file": doc["filename"],
            "old": existing,
            "new": new_chunks,
            "section_rate": round(pct, 1),
        })

    # Summary
    print("\n" + "=" * 60)
    print("RESULTS")
    print("=" * 60)
    for r in results:
        print(f"  {r['file']:<40} old={r['old']} new={r['new']} sect={r['section_rate']}%")

    total_new = sum(r["new"] for r in results)
    print(f"\nTotal new chunks: {total_new}")

    if any(r["new"] == 0 for r in results):
        print("\nWARNING: Some documents produced 0 chunks!")
        sys.exit(1)


if __name__ == "__main__":
    main()
