#!/usr/bin/env python3
"""
ingest.py — RAG Ingestion Pipeline (v2)

Ablauf pro Datei:
    1.  Extraktion     → extractor.py  (Text, OCR, DOCX, PDF, Bilder)
    2.  Normalisierung → normalizer.py (KI: Rohtext → Template-Struktur)
    3.  Chunking       → token-basiert mit Overlap
    4.  Embedding      → multilingual-e5-small (lokal)
    5.  Speichern      → pgvector via SSH-Tunnel (Hetzner / anythingllm)

CLI-Befehle:
    python ingest.py file  <pfad>          Einzelne Datei
    python ingest.py dir   <pfad>          Ganzes Verzeichnis
    python ingest.py watch <pfad>          Ordner live beobachten
    python ingest.py list                  DB-Inhalt anzeigen
    python ingest.py delete <titel>        Dokument aus DB löschen

Flags:
    --force             Bestehende Chunks überschreiben
    --no-normalize      KI-Normalisierung überspringen
    --dry-run           Extrahieren + normalisieren, aber nicht in DB speichern
    --quality-min 0.4   Mindest-Qualitäts-Score (Standard: 0.0 = alles speichern)
"""

import os
import sys
import uuid
import json
import argparse
from datetime import datetime
from pathlib import Path

import tiktoken
import psycopg2
from psycopg2.extras import execute_values
from sshtunnel import SSHTunnelForwarder
from dotenv import load_dotenv

from extractor import extract, is_supported, SUPPORTED_EXTENSIONS
from normalizer import normalize

load_dotenv(".env.ingest")

# ---------------------------------------------------------------------------
# Konfiguration (aus .env.ingest)
# ---------------------------------------------------------------------------
SSH_HOST           = os.getenv("SSH_HOST")
SSH_USER           = os.getenv("SSH_USER", "root")
SSH_KEY_PATH       = os.path.expanduser(os.getenv("SSH_KEY_PATH", "~/.ssh/id_ed25519"))
SSH_KEY_PASSPHRASE = os.getenv("SSH_KEY_PASSPHRASE")
DB_HOST            = os.getenv("DB_HOST", "127.0.0.1")
DB_PORT            = int(os.getenv("DB_PORT", 5432))
DB_NAME            = os.getenv("DB_NAME", "anythingllm")
DB_USER            = os.getenv("DB_USER", "anythingllm")
DB_PASSWORD        = os.getenv("DB_PASSWORD")
NAMESPACE          = os.getenv("NAMESPACE", "mein-workspace")
CHUNK_SIZE         = int(os.getenv("CHUNK_SIZE", 500))
CHUNK_OVERLAP      = int(os.getenv("CHUNK_OVERLAP", 50))
EMBED_MODEL        = "intfloat/multilingual-e5-small"

_embed_model = None  # lazy load


# ---------------------------------------------------------------------------
# Chunking
# ---------------------------------------------------------------------------

def chunk_text(text: str) -> list[str]:
    """Teilt Text in überlappende Token-Chunks auf."""
    enc    = tiktoken.get_encoding("cl100k_base")
    tokens = enc.encode(text)
    chunks = []
    start  = 0
    while start < len(tokens):
        end = min(start + CHUNK_SIZE, len(tokens))
        chunks.append(enc.decode(tokens[start:end]))
        start += CHUNK_SIZE - CHUNK_OVERLAP
    return chunks


# ---------------------------------------------------------------------------
# Embedding
# ---------------------------------------------------------------------------

def embed(texts: list[str]) -> list[list[float]]:
    """Erstellt Embeddings via multilingual-e5-small (lokal, einmaliger Download ~120 MB)."""
    global _embed_model
    if _embed_model is None:
        from sentence_transformers import SentenceTransformer
        print("  ⏳ Lade Embedding-Modell (einmalig)...")
        _embed_model = SentenceTransformer(EMBED_MODEL)
    prefixed = [f"passage: {t}" for t in texts]
    return _embed_model.encode(prefixed).tolist()


# ---------------------------------------------------------------------------
# Datenbank
# ---------------------------------------------------------------------------

def _db_connect(tunnel_port: int):
    return psycopg2.connect(
        host="127.0.0.1",
        port=tunnel_port,
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD,
    )


def _tunnel():
    return SSHTunnelForwarder(
        (SSH_HOST, 22),
        ssh_username=SSH_USER,
        ssh_pkey=SSH_KEY_PATH,
        ssh_private_key_password=SSH_KEY_PASSPHRASE,
        remote_bind_address=("127.0.0.1", DB_PORT),
    )


def _count_existing(cur, title: str) -> int:
    cur.execute(
        "SELECT COUNT(*) FROM anythingllm_vectors "
        "WHERE metadata->>'title' = %s AND namespace = %s",
        (title, NAMESPACE),
    )
    return cur.fetchone()[0]


def _delete_existing(cur, title: str):
    cur.execute(
        "DELETE FROM anythingllm_vectors "
        "WHERE metadata->>'title' = %s AND namespace = %s",
        (title, NAMESPACE),
    )


def _insert_chunks(cur, chunks, embeddings, source_path: Path, extra_meta: dict):
    now     = datetime.now().isoformat()
    records = []
    for chunk, embedding in zip(chunks, embeddings):
        metadata = {
            "id":          str(uuid.uuid4()),
            "url":         f"file://{source_path.resolve()}",
            "text":        chunk,
            "title":       source_path.name,
            "docSource":   "rag-ingestion-v2",
            "published":   now,
            "wordCount":   len(chunk.split()),
            "chunkSource": str(source_path.resolve()),
            **extra_meta,
        }
        records.append((
            str(uuid.uuid4()),
            json.dumps(metadata),
            NAMESPACE,
            embedding,
        ))
    execute_values(
        cur,
        "INSERT INTO anythingllm_vectors (id, metadata, namespace, embedding) VALUES %s",
        records,
        template="(%s::uuid, %s::jsonb, %s, %s::vector)",
    )


# ---------------------------------------------------------------------------
# Kern-Pipeline
# ---------------------------------------------------------------------------

def ingest_file(
    file_path: str,
    force: bool = False,
    normalize_doc: bool = True,
    dry_run: bool = False,
    quality_min: float = 0.0,
) -> bool:
    path = Path(file_path).resolve()

    if not path.exists():
        print(f"❌ Datei nicht gefunden: {path}")
        return False

    if not is_supported(path):
        print(f"⏭️  Übersprungen (Format nicht unterstützt): {path.suffix}")
        return False

    print(f"\n📄 {path.name}")

    # ── 1. Extraktion ────────────────────────────────────────────────────────
    try:
        raw_text = extract(path)
        print(f"  ✅ Text extrahiert ({len(raw_text):,} Zeichen)")
    except Exception as e:
        print(f"  ❌ Extraktion fehlgeschlagen: {e}")
        return False

    if not raw_text.strip():
        print("  ⚠️  Kein Text extrahiert — Datei übersprungen")
        return False

    # ── 2. KI-Normalisierung ─────────────────────────────────────────────────
    try:
        norm = normalize(
            raw_text,
            filename=path.name,
            skip_normalization=not normalize_doc,
        )
        text_to_index = norm["normalized_text"]
    except Exception as e:
        print(f"  ❌ Normalisierung fehlgeschlagen: {e}")
        print("  ↩️  Verwende Rohtext als Fallback")
        norm = {
            "normalized_text": raw_text,
            "doc_type":        "unbekannt",
            "template_file":   "",
            "quality_score":   0.5,
            "missing_fields":  0,
            "was_normalized":  False,
        }
        text_to_index = raw_text

    # Qualitäts-Filter
    if norm["quality_score"] < quality_min:
        print(
            f"  ⏭️  Qualitäts-Score {norm['quality_score']:.0%} "
            f"< Mindest-Score {quality_min:.0%} — übersprungen"
        )
        return False

    extra_meta = {
        "doc_type":       norm["doc_type"],
        "quality_score":  norm["quality_score"],
        "missing_fields": norm["missing_fields"],
        "was_normalized": norm["was_normalized"],
        "template_file":  norm["template_file"],
        "source_format":  path.suffix.lower(),
    }

    # ── 3. Chunking ──────────────────────────────────────────────────────────
    chunks = chunk_text(text_to_index)
    print(f"  ✅ {len(chunks)} Chunks erstellt")

    # ── 4. Embedding ─────────────────────────────────────────────────────────
    try:
        embeddings = embed(chunks)
        print(f"  ✅ {len(embeddings)} Embeddings erstellt")
    except Exception as e:
        print(f"  ❌ Embedding fehlgeschlagen: {e}")
        return False

    # ── Dry-Run Ende ─────────────────────────────────────────────────────────
    if dry_run:
        print(f"  🔍 Dry-Run — Typ: {norm['doc_type']} | "
              f"Score: {norm['quality_score']:.0%} | "
              f"Fehlend: {norm['missing_fields']}")
        if norm["was_normalized"]:
            print("\n" + "─" * 60)
            print(text_to_index[:600] + ("..." if len(text_to_index) > 600 else ""))
            print("─" * 60)
        return True

    # ── 5. Datenbank ─────────────────────────────────────────────────────────
    print("  ⏳ Verbinde mit Datenbank (SSH-Tunnel)...")
    try:
        with _tunnel() as tunnel:
            conn = _db_connect(tunnel.local_bind_port)
            cur  = conn.cursor()

            existing = _count_existing(cur, path.name)
            if existing > 0:
                if force:
                    _delete_existing(cur, path.name)
                    print(f"  🗑️  {existing} bestehende Chunks gelöscht")
                else:
                    print(
                        f"  ⚠️  '{path.name}' bereits in DB ({existing} Chunks). "
                        "Nutze --force zum Überschreiben."
                    )
                    cur.close(); conn.close()
                    return False

            _insert_chunks(cur, chunks, embeddings, path, extra_meta)
            conn.commit()
            cur.close(); conn.close()
            print(f"  ✅ {len(chunks)} Chunks gespeichert "
                  f"(Typ: {norm['doc_type']}, Score: {norm['quality_score']:.0%})")
    except Exception as e:
        print(f"  ❌ Datenbankfehler: {e}")
        return False

    return True


def ingest_directory(
    dir_path: str,
    force: bool = False,
    normalize_doc: bool = True,
    dry_run: bool = False,
    quality_min: float = 0.0,
):
    path = Path(dir_path).resolve()
    if not path.is_dir():
        print(f"❌ Verzeichnis nicht gefunden: {path}")
        return

    files = [f for f in sorted(path.iterdir()) if f.is_file() and is_supported(f)]
    if not files:
        print(f"❌ Keine unterstützten Dateien in: {path}")
        print(f"   Unterstützte Formate: {', '.join(sorted(SUPPORTED_EXTENSIONS))}")
        return

    print(f"\n📁 {len(files)} Datei(en) in: {path}")
    success = 0
    for f in files:
        if ingest_file(str(f), force=force, normalize_doc=normalize_doc,
                       dry_run=dry_run, quality_min=quality_min):
            success += 1

    print(f"\n{'─' * 50}")
    print(f"✅ Fertig: {success}/{len(files)} Dateien verarbeitet")


def watch_directory(
    dir_path: str,
    force: bool = False,
    normalize_doc: bool = True,
    quality_min: float = 0.0,
):
    try:
        from watchdog.observers import Observer
        from watchdog.events import FileSystemEventHandler
    except ImportError:
        print("❌ watchdog nicht installiert: pip install watchdog")
        sys.exit(1)

    class _Handler(FileSystemEventHandler):
        def on_created(self, event):
            if event.is_directory:
                return
            p = Path(event.src_path)
            if is_supported(p):
                print(f"\n🆕 Neue Datei erkannt: {p.name}")
                ingest_file(str(p), force=force, normalize_doc=normalize_doc,
                            quality_min=quality_min)

    observer = Observer()
    observer.schedule(_Handler(), dir_path, recursive=False)
    observer.start()
    print(f"👁️  Beobachte: {dir_path}")
    print("   Drücke Ctrl+C zum Beenden\n")
    try:
        import time
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        observer.stop()
    observer.join()


def list_documents():
    print(f"\n📋 Dokumente in Namespace '{NAMESPACE}':\n")
    try:
        with _tunnel() as tunnel:
            conn = _db_connect(tunnel.local_bind_port)
            cur  = conn.cursor()
            cur.execute("""
                SELECT
                    metadata->>'title'         AS titel,
                    metadata->>'doc_type'      AS typ,
                    metadata->>'quality_score' AS score,
                    metadata->>'source_format' AS format,
                    COUNT(*)                   AS chunks,
                    MAX(metadata->>'published') AS datum
                FROM anythingllm_vectors
                WHERE namespace = %s
                GROUP BY titel, typ, score, format
                ORDER BY datum DESC
            """, (NAMESPACE,))
            rows = cur.fetchall()
            cur.close(); conn.close()

        if not rows:
            print("  Keine Dokumente gefunden.")
            return

        header = f"  {'Titel':<45} {'Typ':<20} {'Score':>6}  {'Fmt':>5}  {'Chunks':>6}  Datum"
        print(header)
        print("  " + "─" * (len(header) - 2))
        for titel, typ, score, fmt, chunks, datum in rows:
            score_str = f"{float(score):.0%}" if score else "–"
            print(
                f"  {(titel or '?'):<45} {(typ or '?'):<20} "
                f"{score_str:>6}  {(fmt or '?'):>5}  {chunks:>6}  "
                f"{(datum or '?')[:19]}"
            )
        print()
    except Exception as e:
        print(f"❌ Datenbankfehler: {e}")


def delete_document(title: str):
    print(f"\n🗑️  Lösche '{title}' aus Namespace '{NAMESPACE}'...")
    try:
        with _tunnel() as tunnel:
            conn     = _db_connect(tunnel.local_bind_port)
            cur      = conn.cursor()
            existing = _count_existing(cur, title)
            if existing == 0:
                print(f"  ⚠️  Kein Dokument mit Titel '{title}' gefunden")
                cur.close(); conn.close()
                return
            _delete_existing(cur, title)
            conn.commit()
            cur.close(); conn.close()
            print(f"  ✅ {existing} Chunks gelöscht")
    except Exception as e:
        print(f"  ❌ Datenbankfehler: {e}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="RAG Ingestion v2 — Dokumente in pgvector importieren",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Beispiele:
  python ingest.py file  ~/Downloads/fremdes_cv.pdf
  python ingest.py file  ~/Desktop/scan.png --dry-run
  python ingest.py dir   ~/Dokumente/bewerbung/ --force
  python ingest.py watch ~/Desktop/scan-eingang/
  python ingest.py list
  python ingest.py delete "fremdes_cv.pdf"
        """
    )
    sub = p.add_subparsers(dest="command", required=True)

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--force",        action="store_true", help="Bestehende Chunks überschreiben")
    common.add_argument("--no-normalize", action="store_true", help="KI-Normalisierung überspringen")
    common.add_argument("--quality-min",  type=float, default=0.0, metavar="0.0-1.0",
                        help="Mindest-Qualitäts-Score (Standard: 0.0)")
    common.add_argument("--dry-run",      action="store_true",
                        help="Nur extrahieren + normalisieren, nicht speichern")

    pf = sub.add_parser("file",  parents=[common], help="Einzelne Datei importieren")
    pf.add_argument("path", help="Pfad zur Datei")

    pd = sub.add_parser("dir",   parents=[common], help="Verzeichnis importieren")
    pd.add_argument("path", help="Pfad zum Verzeichnis")

    pw = sub.add_parser("watch", parents=[common], help="Ordner live beobachten")
    pw.add_argument("path", help="Pfad zum Verzeichnis")

    sub.add_parser("list", help="Alle Dokumente in der DB anzeigen")

    pdel = sub.add_parser("delete", help="Dokument aus DB löschen")
    pdel.add_argument("title", help="Titel (Dateiname) des Dokuments")

    return p


def main():
    args         = _build_parser().parse_args()
    do_normalize = not getattr(args, "no_normalize", False)
    force        = getattr(args, "force", False)
    dry_run      = getattr(args, "dry_run", False)
    quality_min  = getattr(args, "quality_min", 0.0)

    if args.command == "file":
        ingest_file(args.path, force=force, normalize_doc=do_normalize,
                    dry_run=dry_run, quality_min=quality_min)
    elif args.command == "dir":
        ingest_directory(args.path, force=force, normalize_doc=do_normalize,
                         dry_run=dry_run, quality_min=quality_min)
    elif args.command == "watch":
        watch_directory(args.path, force=force, normalize_doc=do_normalize,
                        quality_min=quality_min)
    elif args.command == "list":
        list_documents()
    elif args.command == "delete":
        delete_document(args.title)


if __name__ == "__main__":
    main()