#!/usr/bin/env python3

import os
import sys
import re
import io
import time
from datetime import datetime, timedelta
import hashlib
import sqlite3
import logging
import tempfile
import threading
from typing import Dict, List
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests
from pypdf import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# ============================================================
# CONFIG
# ============================================================
PAPERLESS_URL   = os.environ["PAPERLESS_URL"].rstrip("/")
PAPERLESS_TOKEN = os.environ["PAPERLESS_TOKEN"]

OPENWEBUI_URL   = os.environ["OPENWEBUI_URL"].rstrip("/")
OPENWEBUI_TOKEN = os.environ["OPENWEBUI_TOKEN"]

KNOWLEDGE_NAME  = os.environ["KNOWLEDGE_NAME"]
SYNC_INTERVAL   = int(os.getenv("SYNC_INTERVAL", "3600"))
MAX_WORKERS     = int(os.getenv("MAX_WORKERS", "4"))

DB_PATH         = "/data/paperless_sync.db"
CHUNK_SIZE      = int(os.getenv("CHUNK_SIZE", "1200"))
CHUNK_OVERLAP   = int(os.getenv("CHUNK_OVERLAP", "200"))
MAX_TEXT_LENGTH = int(os.getenv("MAX_TEXT_LENGTH", "2000000"))

KNOWLEDGE_ID    = None

# ============================================================
# LOGGING
# ============================================================
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger("paperless-sync")

# ============================================================
# HEADERS
# ============================================================
paperless_headers = {
    "Authorization": f"Token {PAPERLESS_TOKEN}",
    "Accept": "application/json",
}
openwebui_headers = {
    "Authorization": f"Bearer {OPENWEBUI_TOKEN}",
}

# ============================================================
# SPLITTER
# ============================================================
splitter = RecursiveCharacterTextSplitter(
    chunk_size=CHUNK_SIZE,
    chunk_overlap=CHUNK_OVERLAP,
)

# ============================================================
# SQLITE
# ============================================================
def init_db():
    conn = sqlite3.connect(DB_PATH)
    cur  = conn.cursor()
    cur.execute("""
        CREATE TABLE IF NOT EXISTS synced_documents (
            paperless_id INTEGER PRIMARY KEY,
            modified     TEXT NOT NULL,
            content_hash TEXT,
            synced_at    TEXT DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.commit()
    conn.close()

def get_synced_document(doc_id: int):
    conn = sqlite3.connect(DB_PATH)
    cur  = conn.cursor()
    cur.execute("""
        SELECT modified, content_hash
        FROM synced_documents
        WHERE paperless_id = ?
    """, (doc_id,))
    row = cur.fetchone()
    conn.close()
    return row

def get_last_sync_time() -> str | None:
    """Gibt den Zeitpunkt des letzten erfolgreichen Syncs zurück."""
    conn = sqlite3.connect(DB_PATH)
    row  = conn.execute(
        "SELECT MAX(synced_at) FROM synced_documents"
    ).fetchone()
    conn.close()
    return row[0] if row and row[0] else None

def upsert_synced_document(doc_id: int, modified: str, content_hash: str):
    conn = sqlite3.connect(DB_PATH)
    cur  = conn.cursor()
    cur.execute("""
        INSERT INTO synced_documents (paperless_id, modified, content_hash)
        VALUES (?, ?, ?)
        ON CONFLICT(paperless_id) DO UPDATE SET
            modified     = excluded.modified,
            content_hash = excluded.content_hash,
            synced_at    = CURRENT_TIMESTAMP
    """, (doc_id, modified, content_hash))
    conn.commit()
    conn.close()

# ── Thread-Lock für DB ────────────────────────────────────────
db_lock = threading.Lock()

# ============================================================
# PAPERLESS HELPERS
# ============================================================
def get_all_pages(endpoint: str):
    results = []
    page    = 1
    while True:
        url = f"{PAPERLESS_URL}{endpoint}?page={page}"
        r   = requests.get(url, headers=paperless_headers, timeout=120)
        r.raise_for_status()
        data         = r.json()
        page_results = data.get("results", [])
        if not page_results:
            break
        results.extend(page_results)
        if not data.get("next"):
            break
        page += 1
    return results

def get_all_paperless_ids() -> set:
    """
    Holt nur die IDs aller Paperless-Dokumente – sehr schnell.
    Wird für die Lösch-Erkennung verwendet.
    """
    ids  = set()
    page = 1
    while True:
        url = f"{PAPERLESS_URL}/api/documents/?page_size=100&page={page}&fields=id"
        r   = requests.get(url, headers=paperless_headers, timeout=120)
        r.raise_for_status()
        data    = r.json()
        results = data.get("results", [])
        if not results:
            break
        ids.update(doc["id"] for doc in results)
        if not data.get("next"):
            break
        page += 1
    log.info("Found %s total document IDs in Paperless", len(ids))
    return ids

def get_all_paperless_documents() -> list:
    """
    Holt Dokumente die gesynct werden müssen:
    1. Dokumente die seit dem letzten Sync geändert wurden
    2. Dokumente die noch nie gesynct wurden (nicht in SQLite)
    Beim ersten Sync (keine DB-Einträge) werden alle Dokumente geholt.
    """
    documents = []
    page      = 1
    last_sync = get_last_sync_time()

    if last_sync:
        # ── Strategie: Alle Dokumente holen, aber effizient filtern ──
        # Wir holen ALLE IDs + modified-Timestamps aus Paperless
        # und vergleichen mit unserer lokalen DB.
        # So finden wir sowohl geänderte als auch fehlende Dokumente.

        # Erst: Welche IDs haben wir schon gesynct?
        conn = sqlite3.connect(DB_PATH)
        synced_ids = {
            row[0]
            for row in conn.execute(
                "SELECT paperless_id FROM synced_documents"
            ).fetchall()
        }
        conn.close()

        # Dann: Geänderte Dokumente seit letztem Sync holen
        dt         = datetime.fromisoformat(last_sync) - timedelta(minutes=10)
        since      = dt.strftime("%Y-%m-%dT%H:%M:%S")
        filter_url = f"{PAPERLESS_URL}/api/documents/?page_size=100&modified__gt={since}"
        log.info("Fetching documents modified after %s", since)

        while True:
            url = f"{filter_url}&page={page}"
            r   = requests.get(url, headers=paperless_headers, timeout=120)
            r.raise_for_status()
            data    = r.json()
            results = data.get("results", [])
            if not results:
                break
            documents.extend(results)
            log.info("Paperless page %s: %s modified docs", page, len(results))
            if not data.get("next"):
                break
            page += 1

        modified_ids = {doc["id"] for doc in documents}

        # Jetzt: Alle Paperless-IDs holen und fehlende identifizieren
        all_paperless_ids = get_all_paperless_ids()
        missing_ids = all_paperless_ids - synced_ids - modified_ids

        if missing_ids:
            log.info(
                "Found %s documents never synced, fetching them...",
                len(missing_ids),
            )
            # Fehlende Dokumente einzeln oder in Batches nachladen
            for batch_start in range(0, len(missing_ids), 100):
                batch = list(missing_ids)[batch_start:batch_start + 100]
                for doc_id in batch:
                    try:
                        url = f"{PAPERLESS_URL}/api/documents/{doc_id}/"
                        r   = requests.get(
                            url, headers=paperless_headers, timeout=120
                        )
                        r.raise_for_status()
                        documents.append(r.json())
                    except Exception as e:
                        log.warning(
                            "Failed to fetch missing doc %s: %s", doc_id, e
                        )

        log.info(
            "Total documents to process: %s (modified: %s, missing: %s)",
            len(documents), len(modified_ids), len(missing_ids),
        )
    else:
        filter_url = f"{PAPERLESS_URL}/api/documents/?page_size=100"
        log.info("No previous sync found, fetching all documents")

        while True:
            url = f"{filter_url}&page={page}"
            r   = requests.get(url, headers=paperless_headers, timeout=120)
            r.raise_for_status()
            data    = r.json()
            results = data.get("results", [])
            if not results:
                break
            documents.extend(results)
            log.info("Paperless page %s: %s docs", page, len(results))
            if not data.get("next"):
                break
            page += 1

    return documents

# ============================================================
# METADATA LOOKUPS
# ============================================================
log.info("Loading Paperless metadata...")
TAGS = {
    x["id"]: x["name"]
    for x in get_all_pages("/api/tags/")
}
DOCUMENT_TYPES = {
    x["id"]: x["name"]
    for x in get_all_pages("/api/document_types/")
}
CORRESPONDENTS = {
    x["id"]: x["name"]
    for x in get_all_pages("/api/correspondents/")
}
log.info(
    "Loaded metadata | tags=%s | doc_types=%s | correspondents=%s",
    len(TAGS), len(DOCUMENT_TYPES), len(CORRESPONDENTS),
)

# ============================================================
# CLEANUP
# ============================================================
def clean_ocr_text(text: str):
    if not text:
        return ""
    text = text.encode("utf-8", errors="ignore").decode("utf-8", errors="ignore")
    text = text.replace("\x00", " ")
    text = "".join(
        ch for ch in text
        if ch == "\n" or ch == "\t" or ord(ch) >= 32
    )
    text = re.sub(r"\s+", " ", text)
    if len(text) > MAX_TEXT_LENGTH:
        log.warning("Truncating text (%s chars)", len(text))
        text = text[:MAX_TEXT_LENGTH]
    return text.strip()

# ============================================================
# DOWNLOAD + PDF EXTRACTION
# ============================================================
def extract_text_from_pdf_bytes(pdf_bytes: bytes) -> str:
    try:
        reader = PdfReader(io.BytesIO(pdf_bytes))
        texts  = []
        for page in reader.pages:
            try:
                text = page.extract_text()
                if text:
                    texts.append(text)
            except Exception as e:
                log.warning("PDF page extraction failed: %s", e)
        return "\n".join(texts)
    except Exception as e:
        log.warning("PDF extraction failed: %s", e)
        return ""

def download_document_text(doc_id: int) -> str:
    url = f"{PAPERLESS_URL}/api/documents/{doc_id}/download/"
    try:
        r = requests.get(
            url,
            headers={"Authorization": f"Token {PAPERLESS_TOKEN}"},
            timeout=300,
            allow_redirects=False,
        )

        if r.status_code in (301, 302, 303, 307, 308):
            log.warning(
                "Redirected document download doc=%s location=%s",
                doc_id, r.headers.get("Location"),
            )
            return ""

        if not r.ok:
            log.warning("Download failed doc=%s status=%s", doc_id, r.status_code)
            return ""

        content_type = r.headers.get("Content-Type", "").lower()

        if "pdf" in content_type:
            return extract_text_from_pdf_bytes(r.content).strip()

        # PLAIN TEXT
        text    = r.text.encode("utf-8", errors="ignore").decode("utf-8", errors="ignore")
        lowered = text.lower()

        if (
            "<html" in lowered
            or "<!doctype html" in lowered
            or "paperless-ngx sign in" in lowered
        ):
            log.warning("HTML returned for doc=%s", doc_id)
            return ""

        return text.strip()

    except Exception as e:
        log.warning("Document download failed doc=%s err=%s", doc_id, e)
        return ""

# ============================================================
# BUILD CHUNKS
# ============================================================
def build_chunk_documents(doc: Dict, text: str) -> List[str]:
    doc_id        = doc["id"]
    title         = doc.get("title") or f"document_{doc_id}"
    created       = doc.get("created")
    modified      = doc.get("modified")
    tag_names     = [TAGS[t] for t in doc.get("tags", []) if t in TAGS]
    document_type = DOCUMENT_TYPES.get(doc.get("document_type"))
    correspondent = CORRESPONDENTS.get(doc.get("correspondent"))
    chunks        = splitter.split_text(text)
    documents     = []

    for i, chunk in enumerate(chunks):
        content = f"""DOCUMENT_TITLE: {title}
PAPERLESS_ID: {doc_id}
CHUNK: {i + 1}/{len(chunks)}
CREATED: {created}
MODIFIED: {modified}
DOCUMENT_TYPE: {document_type or "Unknown"}
CORRESPONDENT: {correspondent or "Unknown"}
TAGS: {", ".join(tag_names) if tag_names else "None"}
ORIGINAL_FILENAME: {doc.get("original_file_name")}

CONTENT:
{chunk}"""
        documents.append(content.strip())

    return documents

# ============================================================
# OPENWEBUI
# ============================================================
def resolve_or_create_knowledge():
    for url in [
        f"{OPENWEBUI_URL}/api/v1/knowledge/",
        f"{OPENWEBUI_URL}/api/v1/knowledge",
    ]:
        try:
            r = requests.get(url, headers=openwebui_headers, timeout=120)
            if not r.ok:
                continue
            try:
                data = r.json()
            except Exception:
                continue

            if isinstance(data, list):
                kb_items = data
            elif isinstance(data, dict):
                kb_items = data.get("data") or data.get("items") or []
            else:
                kb_items = []

            for kb in kb_items:
                if kb.get("name") == KNOWLEDGE_NAME:
                    return kb.get("id")

        except Exception:
            log.exception("KB search failed")

    # CREATE KB
    r = requests.post(
        f"{OPENWEBUI_URL}/api/v1/knowledge/create",
        headers={**openwebui_headers, "Content-Type": "application/json"},
        json={"name": KNOWLEDGE_NAME, "description": "Paperless sync"},
        timeout=120,
    )
    r.raise_for_status()
    data  = r.json()
    kb_id = data.get("id")
    if not kb_id:
        raise Exception(f"No KB ID: {data}")
    return kb_id

def upload_to_openwebui(filename: str, content: str):
    global KNOWLEDGE_ID

    with tempfile.NamedTemporaryFile(
        suffix=".txt",
        delete=False,
        mode="w",
        encoding="utf-8",
        errors="ignore",
    ) as tmp:
        tmp.write(content)
        tmp_path = tmp.name

    try:
        # ====================================================
        # UPLOAD MIT RETRY
        # ====================================================
        file_id = None
        for attempt in range(5):
            try:
                with open(tmp_path, "rb") as f:
                    r = requests.post(
                        f"{OPENWEBUI_URL}/api/v1/files/",
                        headers=openwebui_headers,
                        files={"file": (filename, f, "text/plain")},
                        timeout=300,
                    )
                if not r.ok:
                    log.error(
                        "UPLOAD FAILED status=%s body=%s",
                        r.status_code,
                        r.text[:5000],
                    )
                    raise Exception("Upload failed")
                data    = r.json()
                file_id = data.get("id")
                if not file_id:
                    raise Exception(f"No file ID: {data}")
                log.info("Uploaded file -> %s", file_id)
                break  # Erfolg → raus aus Retry-Schleife

            except requests.exceptions.ConnectionError as e:
                if attempt < 4:
                    wait = 10 * (2 ** attempt)  # 10, 20, 40, 80s
                    log.warning(
                        "Connection failed (attempt %s/5), waiting %ss: %s",
                        attempt + 1, wait, e,
                    )
                    time.sleep(wait)
                else:
                    log.error("All 5 upload attempts failed.")
                    raise

        # ====================================================
        # WAIT FOR EMBEDDING
        # ====================================================
        start = time.time()
        wait  = 0.3
        while True:
            if time.time() - start > 300:
                raise Exception("Processing timeout")
            status_r = requests.get(
                f"{OPENWEBUI_URL}/api/v1/files/{file_id}/process/status",
                headers=openwebui_headers,
                timeout=120,
            )
            status_r.raise_for_status()
            status_data = status_r.json()
            status      = status_data.get("status")
            log.info("Processing status: %s", status)
            if status == "completed":
                break
            if status == "failed":
                raise Exception(f"Processing failed: {status_data}")
            time.sleep(wait)
            wait = min(wait * 1.5, 15)  # exponentieller Backoff, max 15s

        # ====================================================
        # ADD TO KB
        # ====================================================
        kb_r = requests.post(
            f"{OPENWEBUI_URL}/api/v1/knowledge/{KNOWLEDGE_ID}/file/add",
            headers={**openwebui_headers, "Content-Type": "application/json"},
            json={"file_id": file_id},
            timeout=300,
        )

        # 400 = Datei bereits in KB → kein echter Fehler
        if kb_r.status_code == 400:
            log.warning(
                "File already in KB (400), treating as success: %s",
                file_id,
            )
        elif not kb_r.ok:
            kb_r.raise_for_status()

        log.info("Added file to KB")
        return file_id

    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass

# ============================================================
# SYNC – parallelisiert
# ============================================================
def process_single_doc(doc, i: int, total: int):
    """Verarbeitet ein einzelnes Dokument – thread-safe."""
    doc_id   = doc["id"]
    title    = doc.get("title") or f"document_{doc_id}"
    modified = doc.get("modified") or ""

    with db_lock:
        existing = get_synced_document(doc_id)

    text = download_document_text(doc_id)
    if not text.strip():
        log.warning("Skipping empty doc=%s", doc_id)
        return "skipped"

    text         = clean_ocr_text(text)
    chunk_docs   = build_chunk_documents(doc, text)
    content_hash = hashlib.sha256(
        ("".join(chunk_docs)).encode("utf-8")
    ).hexdigest()

    if existing:
        existing_modified, existing_hash = existing
        if existing_modified == modified and existing_hash == content_hash:
            log.debug("Document %s: %s exists and has not changed, hence skip.", doc_id, title)
            return "skipped"
        log.debug("Document %s: %s exists but has changed, hence update.", doc_id, title)
    else:
        log.debug("Document %s: %s does not exist yet.", doc_id, title)

    for idx, chunk in enumerate(chunk_docs):
        filename = f"paperless_{doc_id}_chunk_{idx + 1}.txt"
        upload_to_openwebui(filename, chunk)

    with db_lock:
        upsert_synced_document(doc_id, modified, content_hash)

    time.sleep(float(os.getenv("DOC_PAUSE", "2.0")))

    if existing:
        log.info("Updated: %s", title)
        return "updated"
    else:
        log.info("Uploaded: %s [%d/%d]", title, i, total)
        return "uploaded"

def sync_documents():
    log.info("Starting sync")

    # Nur geänderte Dokumente für Upload
    docs  = get_all_paperless_documents()
    total = len(docs)
    log.info("Found %s changed/new documents", total)

    # Alle IDs für Lösch-Erkennung (nur IDs, sehr schnell)
    all_paperless_ids = get_all_paperless_ids()

    uploaded = 0
    updated  = 0
    skipped  = 0
    failed   = 0
    deleted  = 0

    # ── Lösch-Erkennung ───────────────────────────────────────
    conn   = sqlite3.connect(DB_PATH)
    db_ids = {
        row[0]
        for row in conn.execute(
            "SELECT paperless_id FROM synced_documents"
        ).fetchall()
    }
    conn.close()

    deleted_ids = db_ids - all_paperless_ids
    if deleted_ids:
        log.info(
            "Found %s deleted documents, removing from OpenWebUI...",
            len(deleted_ids),
        )
        for doc_id in deleted_ids:
            try:
                r = requests.get(
                    f"{OPENWEBUI_URL}/api/v1/files/",
                    headers=openwebui_headers,
                    timeout=60,
                )
                r.raise_for_status()
                data      = r.json()
                all_files = (
                    data
                    if isinstance(data, list)
                    else data.get("data", [])
                )

                prefix = f"paperless_{doc_id}_chunk_"
                for f in all_files:
                    filename = (
                        f.get("meta", {}).get("name", "")
                        or f.get("filename", "")
                    )
                    if filename.startswith(prefix):
                        requests.delete(
                            f"{OPENWEBUI_URL}/api/v1/files/{f['id']}",
                            headers=openwebui_headers,
                            timeout=30,
                        )
                        log.info(
                            "Deleted file %s (%s) from OpenWebUI",
                            f["id"], filename,
                        )

                conn = sqlite3.connect(DB_PATH)
                conn.execute(
                    "DELETE FROM synced_documents WHERE paperless_id = ?",
                    (doc_id,),
                )
                conn.commit()
                conn.close()
                log.info("Removed doc %s from sync state", doc_id)
                deleted += 1

            except Exception:
                log.exception("Failed to delete doc %s from OpenWebUI", doc_id)

    # ── Normale Sync-Schleife (parallelisiert) ────────────────
    if total > 0:
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = {}
            for i, doc in enumerate(docs):
                futures[executor.submit(process_single_doc, doc, i + 1, total)] = doc
                time.sleep(0.5)  # Versatz damit Worker nicht alle gleichzeitig starten
            for future in as_completed(futures):
                try:
                    result = future.result()
                    if result == "uploaded": uploaded += 1
                    elif result == "updated": updated  += 1
                    elif result == "skipped": skipped  += 1
                    elif result == "failed":  failed   += 1
                except Exception:
                    failed += 1
                    log.exception("Unexpected error in future")

    log.info(
        "Sync complete | uploaded=%s updated=%s skipped=%s failed=%s deleted=%s",
        uploaded, updated, skipped, failed, deleted,
    )

# ============================================================
# WAIT FOR OPENWEBUI
# ============================================================
def wait_for_openwebui():
    log.info("Waiting for OpenWebUI...")
    while True:
        try:
            if requests.get(f"{OPENWEBUI_URL}/health", timeout=5).status_code == 200:
                log.info("OpenWebUI is ready!")
                return
        except Exception:
            pass
        log.info("Not ready yet, retrying in 10s...")
        time.sleep(10)

# ============================================================
# MAIN
# ============================================================
def main():
    global KNOWLEDGE_ID
    init_db()
    wait_for_openwebui()
    KNOWLEDGE_ID = resolve_or_create_knowledge()
    log.info("Using KB ID: %s", KNOWLEDGE_ID)
    sync_documents()

def calc_wait_time(target_time_str):
    n        = datetime.now()
    hrs, mins = map(int, target_time_str.split(":"))
    target   = n.replace(hour=hrs, minute=mins, second=0, microsecond=0)
    if target <= n:
        target += timedelta(days=1)
    return (target - n).total_seconds()

if __name__ == "__main__":
    RUN_AT = os.getenv("SYNC_TIME", "04:00")
    print(f"Container started, task will run daily at {RUN_AT}", flush=True)
    while True:
        try:
            wait_seconds = calc_wait_time(RUN_AT)
            print(f"Next run in {wait_seconds / 3600:.2f} hours.", flush=True)
            time.sleep(wait_seconds)
            main()
        except KeyboardInterrupt:
            print("Container stopped manually", flush=True)
            sys.exit(0)
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr, flush=True)
            time.sleep(60)