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

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

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

# === 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,
            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 get_or_create_session(session_id: str = None) -> str:
    conn = sqlite3.connect(str(DB_PATH))
    if session_id:
        cur = conn.execute("SELECT id FROM sessions WHERE id = ?", (session_id,))
        if cur.fetchone():
            conn.close()
            return session_id
    # Create new session
    new_id = secrets.token_hex(16)
    conn.execute("INSERT INTO sessions (id, created) VALUES (?, ?)", (new_id, time.time()))
    conn.commit()
    conn.close()
    return new_id

def migrate_old_messages():
    """Migriert alte Nachrichten ohne session_id zu einer default Session."""
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute("SELECT COUNT(*) FROM messages WHERE session_id = '' OR session_id IS NULL")
    count = cur.fetchone()[0]
    if count > 0:
        # Create default session
        default_session = "default_" + secrets.token_hex(8)
        conn.execute("INSERT OR IGNORE INTO sessions (id, created, name) VALUES (?, ?, 'default')", (default_session, time.time()))
        conn.execute("UPDATE messages SET session_id = ? WHERE session_id = '' OR session_id IS NULL", (default_session,))
        conn.execute("UPDATE summaries SET session_id = ? WHERE session_id = '' OR session_id IS NULL", (default_session,))
        conn.commit()
    conn.close()

migrate_old_messages()

# === 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:
    """Holt vollständigen Chat-Verlauf für die UI-Anzeige."""
    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 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
- Verwende keine Formulierungen wie "als KI" oder "als Sprachmodell"
- Reagiere emotional angemessen
- 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"""

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

    messages = [{"role": "system", "content": system}]
    for msg in short_term:
        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")

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

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

@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 = get_or_create_session(session_id)
    char = load_character(character_name)
    # Fix: Load FULL history for display, not just short-term
    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,
        "history": history,
        "long_term": long_term,
        "session_id": sid,
    })
    if not session_id or session_id != sid:
        response.set_cookie(SESSION_COOKIE_NAME, sid, max_age=SESSION_DURATION, httponly=False, samesite="lax")
    return response

@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:
        # Streaming via Server-Sent Events
        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"
            
            # Save complete response
            if full_response.strip():
                save_message(sid, character_name, "assistant", full_response.strip())
            
            # Check if summary needed
            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:
        # Non-streaming fallback
        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),
    }
    
    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.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)):
    """API endpoint to get full chat history — for reconnecting devices."""
    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"})

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