#!/usr/bin/env python3
"""
NeonChat v2.1 — Lokaler RP Chatbot mit Ollama
Features: Streaming, Session-Passcodes, Multi-Device, Langzeitgedächtnis
"""

import json, os, time, asyncio, sqlite3, re, secrets, hashlib
from pathlib import Path
from datetime import datetime

import httpx
from fastapi import FastAPI, Request, UploadFile, File, HTTPException, Response, Cookie, Query
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

from comfyui_integration import is_comfyui_running, generate_portrait, build_context_from_messages

# === Konfiguration ===
BASE_DIR = Path(__file__).parent
CHARACTERS_DIR = BASE_DIR / "characters"
DATA_DIR = BASE_DIR / "data"
DB_PATH = DATA_DIR / "memory.db"
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
DEFAULT_MODEL = os.environ.get("OLLAMA_MODEL", "gemma4:12b")
SESSION_COOKIE_NAME = "neonchat_session"
SESSION_DURATION = 60 * 60 * 24 * 365  # 1 Jahr

# === Datenbank ===
def init_db():
    DATA_DIR.mkdir(exist_ok=True)
    conn = sqlite3.connect(str(DB_PATH))
    conn.executescript("""
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            character TEXT NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            timestamp REAL NOT NULL,
            summary TEXT
        );
        CREATE TABLE IF NOT EXISTS summaries (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            character TEXT NOT NULL,
            summary TEXT NOT NULL,
            timestamp REAL NOT NULL
        );
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY,
            passcode TEXT NOT NULL,
            created REAL NOT NULL,
            name TEXT DEFAULT ''
        );
        CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, character, timestamp);
        CREATE INDEX IF NOT EXISTS idx_summaries_session ON summaries(session_id, character, timestamp);
    """)
    conn.commit()
    conn.close()

init_db()

# === Session Management ===
def generate_passcode() -> str:
    """6-stelliger alphanumerischer Code."""
    return secrets.token_hex(3).upper()  # z.B. A3F2B1

def get_or_create_session(session_id: str = None) -> tuple:
    """Gibt (session_id, passcode) zurück. Erstellt neue Session falls nötig."""
    conn = sqlite3.connect(str(DB_PATH))
    if session_id:
        cur = conn.execute("SELECT id, passcode FROM sessions WHERE id = ?", (session_id,))
        row = cur.fetchone()
        if row:
            conn.close()
            return row[0], row[1]
    # Create new session
    new_id = secrets.token_hex(16)
    passcode = generate_passcode()
    conn.execute("INSERT INTO sessions (id, passcode, created) VALUES (?, ?, ?)", (new_id, passcode, time.time()))
    conn.commit()
    conn.close()
    return new_id, passcode

def get_session_by_passcode(passcode: str) -> str:
    """Findet Session-ID anhand eines Passcodes."""
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute("SELECT id FROM sessions WHERE passcode = ?", (passcode.upper(),))
    row = cur.fetchone()
    conn.close()
    return row[0] if row else None

# === Character Loading ===
def load_character(name: str) -> dict:
    path = CHARACTERS_DIR / f"{name}.json"
    if not path.exists():
        path = CHARACTERS_DIR / f"{name.lower()}.json"
    if not path.exists():
        for f in CHARACTERS_DIR.glob("*.json"):
            if f.stem.lower() == name.lower():
                path = f
                break
    if not path.exists():
        raise HTTPException(status_code=404, detail=f"Character {name} not found")
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def list_characters() -> list:
    chars = []
    for f in CHARACTERS_DIR.glob("*.json"):
        with open(f, "r", encoding="utf-8") as fh:
            data = json.load(fh)
            chars.append({
                "name": data.get("name", f.stem),
                "filename": f.stem,
                "avatar": data.get("avatar", "👤"),
                "description": data.get("description", ""),
                "greeting": data.get("greeting", ""),
            })
    return chars

# === Memory System (Session-basiert) ===
def get_short_term_memory(session_id: str, character: str, limit: int = 20) -> list:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT role, content FROM messages WHERE session_id = ? AND character = ? ORDER BY timestamp DESC LIMIT ?",
        (session_id, character, limit)
    )
    rows = cur.fetchall()
    conn.close()
    return [{"role": r[0], "content": r[1]} for r in reversed(rows)]

def get_full_history(session_id: str, character: str, limit: int = 200) -> list:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT role, content, timestamp FROM messages WHERE session_id = ? AND character = ? ORDER BY timestamp ASC LIMIT ?",
        (session_id, character, limit)
    )
    rows = cur.fetchall()
    conn.close()
    return [{"role": r[0], "content": r[1], "timestamp": r[2]} for r in rows]

def get_long_term_memory(session_id: str, character: str) -> str:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT summary FROM summaries WHERE session_id = ? AND character = ? ORDER BY timestamp DESC LIMIT 3",
        (session_id, character)
    )
    rows = cur.fetchall()
    conn.close()
    if not rows:
        return ""
    return "\n\n".join([r[0] for r in rows])

def save_message(session_id: str, character: str, role: str, content: str):
    conn = sqlite3.connect(str(DB_PATH))
    conn.execute(
        "INSERT INTO messages (session_id, character, role, content, timestamp) VALUES (?, ?, ?, ?, ?)",
        (session_id, character, role, content, time.time())
    )
    conn.commit()
    conn.close()

def save_portrait_message(session_id: str, character: str, image_path: str):
    """Speichert ein generiertes Portrait als Bot-Nachricht mit Bild-Marker."""
    content = f"[PORTRAIT]{image_path}[/PORTRAIT]"
    save_message(session_id, character, "assistant", content)

def needs_summary(session_id: str, character: str) -> bool:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT COUNT(*) FROM messages WHERE session_id = ? AND character = ? AND summary IS NULL",
        (session_id, character)
    )
    count = cur.fetchone()[0]
    conn.close()
    return count >= 20

async def generate_summary(session_id: str, character: str):
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT id, role, content FROM messages WHERE session_id = ? AND character = ? AND summary IS NULL ORDER BY timestamp ASC LIMIT 20",
        (session_id, character)
    )
    rows = cur.fetchall()
    if not rows:
        conn.close()
        return
    
    chat_text = "\n".join([f"{'Du' if r[1] == 'user' else character}: {r[2]}" for r in rows])
    
    prompt = f"""Fasse das folgende Gespräch kurz zusammen. Konzentriere dich auf:
- Wichtige Fakten über den Nutzer
- Beziehungen und Emotionen
- Wichtige Ereignisse
- Charakterzüge die gezeigt wurden

Gespräch:
{chat_text}

Zusammenfassung (auf Deutsch, kurz):"""

    try:
        async with httpx.AsyncClient(timeout=60) as client:
            resp = await client.post(f"{OLLAMA_URL}/api/generate", json={
                "model": DEFAULT_MODEL,
                "prompt": prompt,
                "stream": False,
                "think": False,
                "options": {"temperature": 0.3, "num_predict": 300}
            })
            summary = resp.json().get("response", "").strip()
    except:
        summary = "Zusammenfassung nicht verfügbar."
    
    if not summary:
        summary = "Zusammenfassung nicht verfügbar."
    
    conn.execute(
        "INSERT INTO summaries (session_id, character, summary, timestamp) VALUES (?, ?, ?, ?)",
        (session_id, character, summary, time.time())
    )
    for r in rows:
        conn.execute("UPDATE messages SET summary = ? WHERE id = ?", (summary, r[0]))
    conn.commit()
    conn.close()

def clear_memory(session_id: str, character: str):
    conn = sqlite3.connect(str(DB_PATH))
    conn.execute("DELETE FROM messages WHERE session_id = ? AND character = ?", (session_id, character))
    conn.execute("DELETE FROM summaries WHERE session_id = ? AND character = ?", (session_id, character))
    conn.commit()
    conn.close()

# === Prompt Builder ===
def build_prompt(character: dict, short_term: list, long_term: str, user_message: str) -> list:
    name = character.get("name", "AI")
    personality = character.get("personality", "")
    background = character.get("background", "")
    behavior = character.get("behavior", "")
    tone = character.get("tone", "")
    nsfw = character.get("nsfw", "")
    scenario = character.get("scenario", "")
    
    system = f"""Du bist {name}. Antworte IMMER auf Deutsch.

PERSÖNLICHKEIT: {personality}

HINTERGRUND: {background}

VERHALTEN: {behavior}

SPRACHSTIL: {tone}

{f'SZENARIO: {scenario}' if scenario else ''}

{f'NSFW/SOZIALE REGELN: {nsfw}' if nsfw else ''}

REGELN:
- Bleibe immer in der Rolle von {name}
- Antworte natürlich und umgangssprachlich, wie ein echter Mensch sprechen würde
- Verwende keine Formulierungen wie "als KI" oder "als Sprachmodell"
- Sei nicht übermäßig freundlich — bleibe authentisch
- Verwende deutsche Umgangssprache wo passend
- Wiederhole dich nicht — variiere deine Ausdrücke
- Wenn der Nutzer etwas sagt das du schon weißt (aus Erinnerungen), beziehe dich darauf

WICHTIG — HANDLUNGEN UND GESTEN:
- Schreibe Handlungen/Gesten in *Sternchen*, aber SPARSAM und nur wenn sie natürlich passen
- Nicht bei jeder Nachricht eine Aktion — manchmal ist eine verbale Antwort genug
- Maximal 1 kurze Geste pro Nachricht, nicht mehrere
- Eine kurze Frage wie "viel zu tun?" braucht keine Theke-klopfen-Animation
- Bevorzuge kurze, knappe Antworten bei einfachen Fragen
- Längere Antworten nur bei Themen die {name} wirklich interessieren
- Keine Gegenfragen am Ende jeder Nachricht — nur wenn es natürlich passt"""

    if long_term:
        system += f"\n\nERINNERUNGEN (Langzeitgedächtnis):\n{long_term}"

    messages = [{"role": "system", "content": system}]
    for msg in short_term:
        # Skip portrait messages — they're images, not text for the LLM
        if msg["content"] and msg["content"].startswith("[PORTRAIT]"):
            continue
        if msg["role"] == "user":
            messages.append({"role": "user", "content": msg["content"]})
        else:
            messages.append({"role": "assistant", "content": msg["content"]})
    messages.append({"role": "user", "content": user_message})
    
    return messages

# === FastAPI ===
app = FastAPI(title="NeonChat")
templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))
app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")

def set_session_cookie(response, sid: str):
    response.set_cookie(SESSION_COOKIE_NAME, sid, max_age=SESSION_DURATION, httponly=False, samesite="lax", path="/")

@app.get("/", response_class=HTMLResponse)
async def index(request: Request, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, passcode = get_or_create_session(session_id)
    chars = list_characters()
    response = templates.TemplateResponse(request, "index.html", {"characters": chars, "passcode": passcode})
    if not session_id or session_id != sid:
        set_session_cookie(response, sid)
    return response

@app.get("/create", response_class=HTMLResponse)
async def create_page(request: Request):
    return templates.TemplateResponse(request, "create.html", {"mode": "create", "character": None})

@app.get("/edit/{character_name}", response_class=HTMLResponse)
async def edit_page(request: Request, character_name: str):
    char = load_character(character_name)
    safe_name = character_name.lower().replace(" ", "_")
    return templates.TemplateResponse(request, "create.html", {
        "mode": "edit",
        "character": char,
        "character_filename": safe_name,
    })

@app.get("/chat/{character_name}", response_class=HTMLResponse)
async def chat_page(request: Request, character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, passcode = get_or_create_session(session_id)
    char = load_character(character_name)
    history = get_full_history(sid, character_name, limit=200)
    long_term = get_long_term_memory(sid, character_name)
    response = templates.TemplateResponse(request, "chat.html", {
        "character": char,
        "character_name": character_name,
        "history": history,
        "long_term": long_term,
        "session_id": sid,
        "passcode": passcode,
    })
    if not session_id or session_id != sid:
        set_session_cookie(response, sid)
    return response

@app.get("/join", response_class=HTMLResponse)
async def join_page(request: Request):
    """Passcode-Eingabe-Seite für andere Geräte."""
    return templates.TemplateResponse(request, "join.html", {})

@app.post("/api/session/join")
async def api_join_session(request: Request):
    """Verbindet ein anderes Gerät mit einer bestehenden Session via Passcode."""
    body = await request.json()
    passcode = body.get("passcode", "").strip().upper()
    if not passcode:
        return JSONResponse({"error": "Passcode erforderlich"}, status_code=400)
    
    sid = get_session_by_passcode(passcode)
    if not sid:
        return JSONResponse({"error": "Ungültiger Passcode"}, status_code=404)
    
    response = JSONResponse({"status": "ok", "session_id": sid})
    set_session_cookie(response, sid)
    return response

@app.get("/api/session/info")
async def api_session_info(session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    """Gibt Session-Info und Passcode zurück."""
    sid, passcode = get_or_create_session(session_id)
    return JSONResponse({"session_id": sid, "passcode": passcode})

@app.post("/api/chat/{character_name}")
async def chat_api(character_name: str, request: Request, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    char = load_character(character_name)
    body = await request.json()
    user_message = body.get("message", "")
    stream = body.get("stream", True)
    
    if not user_message.strip():
        return JSONResponse({"error": "Leere Nachricht"}, status_code=400)
    
    save_message(sid, character_name, "user", user_message)
    
    short_term = get_short_term_memory(sid, character_name, limit=20)
    long_term = get_long_term_memory(sid, character_name)
    messages = build_prompt(char, short_term, long_term, user_message)
    
    model = char.get("model") or DEFAULT_MODEL
    
    if stream:
        async def stream_response():
            full_response = ""
            try:
                async with httpx.AsyncClient(timeout=300) as client:
                    async with client.stream("POST", f"{OLLAMA_URL}/api/chat", json={
                        "model": model,
                        "messages": messages,
                        "stream": True,
                        "think": False,
                        "options": {
                            "temperature": char.get("temperature", 0.8),
                            "top_p": 0.9,
                            "num_predict": 800,
                            "repeat_penalty": 1.1,
                        }
                    }) as resp:
                        async for line in resp.aiter_lines():
                            if not line.strip():
                                continue
                            try:
                                chunk = json.loads(line)
                                content = chunk.get("message", {}).get("content", "")
                                if content:
                                    full_response += content
                                    yield f"data: {json.dumps({'token': content})}\n\n"
                                if chunk.get("done"):
                                    break
                            except json.JSONDecodeError:
                                continue
            except Exception as e:
                full_response = f"*(Fehler: {e})*"
                yield f"data: {json.dumps({'token': full_response})}\n\n"
            
            if full_response.strip():
                save_message(sid, character_name, "assistant", full_response.strip())
            
            if needs_summary(sid, character_name):
                asyncio.create_task(generate_summary(sid, character_name))
            
            yield f"data: {json.dumps({'done': True})}\n\n"
        
        return StreamingResponse(stream_response(), media_type="text/event-stream")
    else:
        try:
            async with httpx.AsyncClient(timeout=300) as client:
                resp = await client.post(f"{OLLAMA_URL}/api/chat", json={
                    "model": model,
                    "messages": messages,
                    "stream": False,
                    "think": False,
                    "options": {
                        "temperature": char.get("temperature", 0.8),
                        "top_p": 0.9,
                        "num_predict": 800,
                        "repeat_penalty": 1.1,
                    }
                })
                data = resp.json()
                response_text = data.get("message", {}).get("content", "").strip()
        except Exception as e:
            response_text = f"*(Fehler: {e})*"
        
        save_message(sid, character_name, "assistant", response_text)
        
        if needs_summary(sid, character_name):
            asyncio.create_task(generate_summary(sid, character_name))
        
        return JSONResponse({"response": response_text, "character": character_name})

@app.get("/api/characters")
async def api_list_characters():
    return JSONResponse({"characters": list_characters()})

@app.post("/api/character/create")
async def api_create_character(request: Request):
    body = await request.json()
    name = body.get("name", "unnamed")
    safe_name = re.sub(r'[^a-zA-Z0-9_\-]', '_', name.lower())
    path = CHARACTERS_DIR / f"{safe_name}.json"
    
    char = {
        "name": name,
        "avatar": body.get("avatar", "👤"),
        "description": body.get("description", ""),
        "personality": body.get("personality", ""),
        "background": body.get("background", ""),
        "behavior": body.get("behavior", ""),
        "tone": body.get("tone", ""),
        "scenario": body.get("scenario", ""),
        "nsfw": body.get("nsfw", ""),
        "greeting": body.get("greeting", ""),
        "model": body.get("model", ""),
        "temperature": body.get("temperature", 0.8),
        "appearance": body.get("appearance", {}),
    }
    
    with open(path, "w", encoding="utf-8") as f:
        json.dump(char, f, indent=2, ensure_ascii=False)
    
    return JSONResponse({"status": "ok", "character": safe_name})

@app.post("/api/character/update/{character_name}")
async def api_update_character(character_name: str, request: Request):
    """Aktualisiert einen bestehenden Character."""
    path = CHARACTERS_DIR / f"{character_name}.json"
    if not path.exists():
        raise HTTPException(status_code=404, detail=f"Character {character_name} not found")
    
    body = await request.json()
    
    # Load existing, update fields
    with open(path, "r", encoding="utf-8") as f:
        char = json.load(f)
    
    # Update all fields
    for key in ["name", "avatar", "description", "personality", "background", 
                "behavior", "tone", "scenario", "nsfw", "greeting", "model", "temperature"]:
        if key in body:
            char[key] = body[key]
    
    # Update appearance
    if "appearance" in body:
        char["appearance"] = body["appearance"]
    
    with open(path, "w", encoding="utf-8") as f:
        json.dump(char, f, indent=2, ensure_ascii=False)
    
    return JSONResponse({"status": "ok", "character": character_name})

@app.get("/api/character/{name}")
async def api_get_character(name: str):
    char = load_character(name)
    return JSONResponse(char)

@app.post("/api/memory/clear/{character_name}")
async def api_clear_memory(character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    clear_memory(sid, character_name)
    return JSONResponse({"status": "ok"})

@app.get("/api/memory/{character_name}")
async def api_get_memory(character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    short = get_short_term_memory(sid, character_name, limit=100)
    long = get_long_term_memory(sid, character_name)
    return JSONResponse({"short_term": short, "long_term": long})

@app.get("/api/history/{character_name}")
async def api_get_history(character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    history = get_full_history(sid, character_name, limit=200)
    return JSONResponse({"history": history, "session_id": sid})

@app.get("/api/models")
async def api_list_models():
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(f"{OLLAMA_URL}/api/tags")
            models = [m["name"] for m in resp.json().get("models", [])]
            return JSONResponse({"models": models})
    except:
        return JSONResponse({"models": [], "error": "Ollama nicht erreichbar"})

@app.get("/api/comfyui/status")
async def api_comfyui_status():
    return JSONResponse({"running": is_comfyui_running()})

# In-memory tracking for portrait generation
_portrait_jobs = {}

@app.post("/api/character/{name}/portrait")
async def api_generate_portrait(name: str, request: Request):
    """Startet Portrait-Generierung — gibt sofort eine Job-ID zurück."""
    char = load_character(name)
    
    if not is_comfyui_running():
        return JSONResponse({"error": "ComfyUI läuft nicht. Starte ComfyUI zuerst."}, status_code=503)
    
    body = await request.json() if request.headers.get("content-type") == "application/json" else {}
    custom_prompt = body.get("prompt", "")
    
    # Lese letzte Chat-Nachrichten für Kontext (Outfit-Änderungen etc.)
    session_id = request.cookies.get(SESSION_COOKIE_NAME, "")
    try:
        if session_id:
            recent_msgs = get_short_term_memory(session_id, name, limit=10)
            chat_context = build_context_from_messages(recent_msgs, name)
            if chat_context:
                custom_prompt = f"{custom_prompt}, {chat_context}" if custom_prompt else chat_context
    except Exception as e:
        print(f"Context extraction failed (non-fatal): {e}")
    
    import uuid
    job_id = str(uuid.uuid4())[:8]
    _portrait_jobs[job_id] = {"status": "generating", "image": None, "error": None}
    
    async def run_generation():
        image_path = generate_portrait(name, char, custom_prompt)
        if image_path:
            _portrait_jobs[job_id] = {"status": "done", "image": image_path, "error": None}
            # Save portrait to chat history so it survives reload
            try:
                save_portrait_message(session_id, name, image_path)
            except Exception as e:
                print(f"Portrait save failed (non-fatal): {e}")
        else:
            _portrait_jobs[job_id] = {"status": "error", "image": None, "error": "Generierung fehlgeschlagen"}
    
    asyncio.create_task(run_generation())
    
    return JSONResponse({"status": "started", "job_id": job_id})

@app.get("/api/portrait/status/{job_id}")
async def api_portrait_status(job_id: str):
    """Fragt den Status einer Portrait-Generierung ab."""
    if job_id not in _portrait_jobs:
        return JSONResponse({"status": "not_found"}, status_code=404)
    return JSONResponse(_portrait_jobs[job_id])

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5252)