#!/usr/bin/env python3
"""
RAG Ingestion Script für apply4jobs.de
Liest Dokumente (MD, TXT, PDF), chunked sie und speichert Embeddings in pgvector.

Späterer Ausbau: Tkinter UI
"""

import os
import uuid
import json
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv
from openai import OpenAI
import psycopg2
from psycopg2.extras import execute_values
from sshtunnel import SSHTunnelForwarder
import tiktoken

# PDF Support (optional)
try:
    from pypdf import PdfReader
    PDF_SUPPORT = True
except ImportError:
    PDF_SUPPORT = False

load_dotenv(".env.anythingllm")

# --- Konfiguration ---
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
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", None)
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"

client = OpenAI(api_key=OPENAI_API_KEY)
embed_model = None  # lazy load


# --- Text-Extraktion ---

def read_file(path: Path) -> str:
    """Liest MD, TXT oder PDF und gibt den Text zurück."""
    suffix = path.suffix.lower()
    if suffix in (".md", ".txt"):
        return path.read_text(encoding="utf-8")
    elif suffix == ".pdf":
        if not PDF_SUPPORT:
            raise ImportError("pypdf nicht installiert: pip install pypdf")
        reader = PdfReader(str(path))
        return "\n\n".join(page.extract_text() or "" for page in reader.pages)
    else:
        raise ValueError(f"Nicht unterstütztes Dateiformat: {suffix}")


# --- Chunking ---

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


# --- Embeddings ---

def get_embeddings(texts: list[str]) -> list[list[float]]:
    """Erstellt Embeddings via multilingual-e5-small (lokal)."""
    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 get_db_connection(tunnel_port: int):
    return psycopg2.connect(
        host="127.0.0.1",
        port=tunnel_port,
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD
    )


def check_existing(cur, source_title: str) -> int:
    """Prüft ob Dokument bereits in der DB vorhanden ist."""
    cur.execute(
        "SELECT COUNT(*) FROM anythingllm_vectors WHERE metadata->>'title' = %s AND namespace = %s",
        (source_title, NAMESPACE)
    )
    return cur.fetchone()[0]


def delete_existing(cur, source_title: str):
    """Löscht alle Chunks eines Dokuments."""
    cur.execute(
        "DELETE FROM anythingllm_vectors WHERE metadata->>'title' = %s AND namespace = %s",
        (source_title, NAMESPACE)
    )
    print(f"  🗑️  Bestehende Chunks für '{source_title}' gelöscht.")


def insert_chunks(cur, chunks: list[str], embeddings: list[list[float]], source_path: Path):
    """Fügt Chunks mit Embeddings in pgvector ein."""
    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-script",
            "published": now,
            "wordCount": len(chunk.split()),
            "chunkSource": str(source_path.resolve()),
        }
        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)"
    )


# --- Haupt-Ingestion ---

def ingest_file(file_path: str, force: bool = False):
    """
    Verarbeitet eine einzelne Datei:
    1. Text lesen
    2. Chunken
    3. Embeddings erstellen
    4. In pgvector speichern (via SSH Tunnel)
    """
    path = Path(file_path).resolve()
    if not path.exists():
        print(f"❌ Datei nicht gefunden: {path}")
        return False

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

    # Text lesen
    try:
        text = read_file(path)
        print(f"  ✅ Text gelesen ({len(text)} Zeichen)")
    except Exception as e:
        print(f"  ❌ Fehler beim Lesen: {e}")
        return False

    # Chunken
    chunks = chunk_text(text)
    print(f"  ✅ {len(chunks)} Chunks erstellt")

    # Embeddings
    print(f"  ⏳ Erstelle Embeddings via OpenAI...")
    try:
        embeddings = get_embeddings(chunks)
        print(f"  ✅ {len(embeddings)} Embeddings erstellt")
    except Exception as e:
        print(f"  ❌ Fehler bei Embeddings: {e}")
        return False

    # SSH Tunnel + DB
    print(f"  ⏳ Verbinde mit Datenbank via SSH Tunnel...")
    try:
        with 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)
        ) as tunnel:
            conn = get_db_connection(tunnel.local_bind_port)
            cur = conn.cursor()

            # Prüfen ob Dokument bereits existiert
            existing = check_existing(cur, path.name)
            if existing > 0:
                if force:
                    delete_existing(cur, path.name)
                else:
                    print(f"  ⚠️  '{path.name}' bereits in DB ({existing} Chunks). Nutze --force zum Überschreiben.")
                    cur.close()
                    conn.close()
                    return False

            # Einfügen
            insert_chunks(cur, chunks, embeddings, path)
            conn.commit()
            cur.close()
            conn.close()
            print(f"  ✅ {len(chunks)} Chunks in pgvector gespeichert")

    except Exception as e:
        print(f"  ❌ Datenbankfehler: {e}")
        return False

    return True


def ingest_directory(dir_path: str, force: bool = False):
    """Verarbeitet alle MD, TXT und PDF Dateien in einem Verzeichnis."""
    path = Path(dir_path).resolve()
    if not path.is_dir():
        print(f"❌ Verzeichnis nicht gefunden: {path}")
        return

    files = list(path.glob("*.md")) + list(path.glob("*.txt")) + list(path.glob("*.pdf"))
    if not files:
        print(f"❌ Keine unterstützten Dateien in: {path}")
        return

    print(f"\n📁 Verarbeite {len(files)} Dateien aus: {path}")
    success = 0
    for f in files:
        if ingest_file(str(f), force=force):
            success += 1

    print(f"\n✅ Fertig: {success}/{len(files)} Dateien erfolgreich importiert.")


def list_documents():
    """Zeigt alle Dokumente in der Datenbank an."""
    print(f"\n📋 Dokumente in Namespace '{NAMESPACE}':")
    try:
        with 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)
        ) as tunnel:
            conn = get_db_connection(tunnel.local_bind_port)
            cur = conn.cursor()
            cur.execute("""
                SELECT metadata->>'title', COUNT(*), MAX(metadata->>'published')
                FROM anythingllm_vectors
                WHERE namespace = %s
                GROUP BY metadata->>'title'
                ORDER BY MAX(metadata->>'published') DESC
            """, (NAMESPACE,))
            rows = cur.fetchall()
            cur.close()
            conn.close()

        if not rows:
            print("  Keine Dokumente gefunden.")
        else:
            print(f"  {'Titel':<50} {'Chunks':>6}  {'Erstellt'}")
            print("  " + "-" * 75)
            for title, count, published in rows:
                print(f"  {(title or 'Unknown'):<50} {count:>6}  {published or '-'}")
        print()
    except Exception as e:
        print(f"❌ Fehler: {e}")


# --- CLI ---

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="RAG Ingestion Script – Dokumente in pgvector importieren"
    )
    subparsers = parser.add_subparsers(dest="command")

    # ingest file
    p_file = subparsers.add_parser("file", help="Einzelne Datei importieren")
    p_file.add_argument("path", help="Pfad zur Datei (MD, TXT, PDF)")
    p_file.add_argument("--force", action="store_true", help="Bestehende Chunks überschreiben")

    # ingest directory
    p_dir = subparsers.add_parser("dir", help="Verzeichnis importieren")
    p_dir.add_argument("path", help="Pfad zum Verzeichnis")
    p_dir.add_argument("--force", action="store_true", help="Bestehende Chunks überschreiben")

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

    args = parser.parse_args()

    if args.command == "file":
        ingest_file(args.path, force=args.force)
    elif args.command == "dir":
        ingest_directory(args.path, force=args.force)
    elif args.command == "list":
        list_documents()
    else:
        parser.print_help()