#!/usr/bin/env python3
"""
RAG Ingestion Script für Job-Matching Workflow
Liest MD/TXT/PDF Dokumente, chunked sie und speichert Embeddings
in der lokalen pgvector Instanz auf der NAS (Port 5433).

Verwendung:
    python ingest_job_matching.py file <pfad>         # Einzelne Datei
    python ingest_job_matching.py dir <pfad>          # Verzeichnis
    python ingest_job_matching.py list                # Alle Dokumente anzeigen
    python ingest_job_matching.py clear <filename>    # Dokument löschen
"""

import os
import uuid
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
import tiktoken

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

# Lade job_matching spezifische .env Datei
env_path = Path(__file__).parent / ".env.job_matching"
load_dotenv(dotenv_path=env_path)

# --- Konfiguration ---
OPENAI_API_KEY  = os.getenv("OPENAI_API_KEY")
DB_HOST         = os.getenv("JM_DB_HOST", "192.168.178.128")
DB_PORT         = int(os.getenv("JM_DB_PORT", 5433))
DB_NAME         = os.getenv("JM_DB_NAME", "job_matching")
DB_USER         = os.getenv("JM_DB_USER", "jobmatch")
DB_PASSWORD     = os.getenv("JM_DB_PASSWORD")
CHUNK_SIZE      = int(os.getenv("CHUNK_SIZE", 500))
CHUNK_OVERLAP   = int(os.getenv("CHUNK_OVERLAP", 50))
EMBED_MODEL     = "text-embedding-3-small"
EMBED_DIM       = 1536

client = OpenAI(api_key=OPENAI_API_KEY)


# --- Datenbank ---

def get_db_connection():
    """Direkte Verbindung zur pgvector Instanz auf der NAS."""
    return psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD
    )


def ensure_table(conn):
    """Stellt sicher dass die profile_documents Tabelle existiert."""
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
        cur.execute(f"""
            CREATE TABLE IF NOT EXISTS profile_documents (
                id SERIAL PRIMARY KEY,
                filename TEXT NOT NULL,
                chunk_index INTEGER NOT NULL,
                content TEXT NOT NULL,
                embedding vector({EMBED_DIM}),
                created_at TIMESTAMP DEFAULT NOW(),
                UNIQUE(filename, chunk_index)
            );
        """)
        cur.execute("""
            CREATE INDEX IF NOT EXISTS profile_documents_embedding_idx
            ON profile_documents
            USING ivfflat (embedding vector_cosine_ops)
            WITH (lists = 10);
        """)
        conn.commit()
    print("  ✅ Tabelle profile_documents bereit")


# --- 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 OpenAI text-embedding-3-small."""
    response = client.embeddings.create(
        model=EMBED_MODEL,
        input=texts
    )
    return [item.embedding for item in response.data]


# --- Ingestion ---

def check_existing(conn, filename: str) -> int:
    """Prüft ob Dokument bereits in der DB vorhanden ist."""
    with conn.cursor() as cur:
        cur.execute(
            "SELECT COUNT(*) FROM profile_documents WHERE filename = %s",
            (filename,)
        )
        return cur.fetchone()[0]


def delete_existing(conn, filename: str):
    """Löscht alle Chunks eines Dokuments."""
    with conn.cursor() as cur:
        cur.execute(
            "DELETE FROM profile_documents WHERE filename = %s",
            (filename,)
        )
        conn.commit()
    print(f"  🗑️  Bestehende Chunks für '{filename}' gelöscht.")


def insert_chunks(conn, chunks: list[str], embeddings: list[list[float]], filename: str):
    """Fügt Chunks mit Embeddings in pgvector ein."""
    records = [
        (filename, i, chunk, embedding)
        for i, (chunk, embedding) in enumerate(zip(chunks, embeddings))
    ]
    with conn.cursor() as cur:
        execute_values(
            cur,
            """
            INSERT INTO profile_documents (filename, chunk_index, content, embedding)
            VALUES %s
            ON CONFLICT (filename, chunk_index) DO UPDATE
                SET content = EXCLUDED.content,
                    embedding = EXCLUDED.embedding,
                    created_at = NOW()
            """,
            records,
            template="(%s, %s, %s, %s::vector)"
        )
        conn.commit()


def ingest_file(file_path: str, force: bool = False) -> bool:
    """Verarbeitet eine einzelne Datei."""
    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 ({EMBED_MODEL})...")
    try:
        embeddings = []
        batch_size = 100
        for i in range(0, len(chunks), batch_size):
            batch = chunks[i:i + batch_size]
            embeddings.extend(get_embeddings(batch))
        print(f"  ✅ {len(embeddings)} Embeddings erstellt")
    except Exception as e:
        print(f"  ❌ Fehler bei Embeddings: {e}")
        return False

    # Datenbank
    print(f"  ⏳ Verbinde mit pgvector auf {DB_HOST}:{DB_PORT}...")
    try:
        conn = get_db_connection()
        ensure_table(conn)

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

        insert_chunks(conn, chunks, embeddings, path.name)
        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"))
    )
    files = [f for f in files if not f.name.startswith(".")]

    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 sorted(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 profile_documents ({DB_HOST}:{DB_PORT}/{DB_NAME}):")
    try:
        conn = get_db_connection()
        with conn.cursor() as cur:
            cur.execute("""
                SELECT filename, COUNT(*) as chunks, MAX(created_at) as imported_at
                FROM profile_documents
                GROUP BY filename
                ORDER BY MAX(created_at) DESC
            """)
            rows = cur.fetchall()
        conn.close()

        if not rows:
            print("  Keine Dokumente gefunden.")
        else:
            print(f"\n  {'Dateiname':<55} {'Chunks':>6}  {'Importiert'}")
            print("  " + "-" * 80)
            for filename, chunks, imported_at in rows:
                ts = imported_at.strftime("%d.%m.%Y %H:%M") if imported_at else "-"
                print(f"  {filename:<55} {chunks:>6}  {ts}")
        print()
    except Exception as e:
        print(f"❌ Fehler: {e}")


def clear_document(filename: str):
    """Löscht ein Dokument aus der Datenbank."""
    try:
        conn = get_db_connection()
        existing = check_existing(conn, filename)
        if existing == 0:
            print(f"⚠️  '{filename}' nicht in der Datenbank gefunden.")
        else:
            delete_existing(conn, filename)
            print(f"✅ '{filename}' gelöscht ({existing} Chunks).")
        conn.close()
    except Exception as e:
        print(f"❌ Fehler: {e}")


# --- CLI ---

if __name__ == "__main__":
    import argparse

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

    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")

    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")

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

    p_clear = subparsers.add_parser("clear", help="Dokument aus DB löschen")
    p_clear.add_argument("filename", help="Dateiname (z.B. 'Lebenslauf - 2026 - Sebastian Fröhlich.md')")

    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()
    elif args.command == "clear":
        clear_document(args.filename)
    else:
        parser.print_help()