#!/usr/bin/env python3
"""
NeonChat — Lokaler RP Chatbot mit Ollama
Character-basiert, mit Kurz- und Langzeitgedächtnis
"""

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

import httpx
from fastapi import FastAPI, Request, UploadFile, File, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
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", "llama3.1:8b")

# === 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,
            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,
            character TEXT NOT NULL,
            summary TEXT NOT NULL,
            timestamp REAL NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_messages_character ON messages(character, timestamp);
        CREATE INDEX IF NOT EXISTS idx_summaries_character ON summaries(character, timestamp);
    """)
    conn.commit()
    conn.close()

init_db()

# === Character Loading ===
def load_character(name: str) -> dict:
    path = CHARACTERS_DIR / f"{name}.json"
    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),
                "avatar": data.get("avatar", "👤"),
                "description": data.get("description", ""),
                "greeting": data.get("greeting", ""),
            })
    return chars

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

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

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

def needs_summary(character: str) -> bool:
    """Prüft ob eine neue Zusammenfassung nötig ist (alle 20 Nachrichten)."""
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT COUNT(*) FROM messages WHERE character = ? AND summary IS NULL",
        (character,)
    )
    count = cur.fetchone()[0]
    conn.close()
    return count >= 20

async def generate_summary(character: str):
    """Erstellt eine Zusammenfassung der letzten 20 Nachrichten."""
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT id, role, content FROM messages WHERE character = ? AND summary IS NULL ORDER BY timestamp ASC LIMIT 20",
        (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() as client:
            resp = await client.post(f"{OLLAMA_URL}/api/generate", json={
                "model": DEFAULT_MODEL,
                "prompt": prompt,
                "stream": False,
                "options": {"temperature": 0.3, "num_predict": 300}
            })
            summary = resp.json().get("response", "").strip()
    except:
        summary = "Zusammenfassung nicht verfügbar."
    
    conn.execute(
        "INSERT INTO summaries (character, summary, timestamp) VALUES (?, ?, ?)",
        (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(character: str):
    conn = sqlite3.connect(str(DB_PATH))
    conn.execute("DELETE FROM messages WHERE character = ?", (character,))
    conn.execute("DELETE FROM summaries WHERE character = ?", (character,))
    conn.commit()
    conn.close()

# === Prompt Builder ===
def build_prompt(character: dict, short_term: list, long_term: str, user_message: str) -> str:
    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"""

    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):
    chars = list_characters()
    return templates.TemplateResponse(request, "index.html", {"characters": chars})

@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):
    char = load_character(character_name)
    history = get_short_term_memory(character_name, limit=50)
    long_term = get_long_term_memory(character_name)
    return templates.TemplateResponse(request, "chat.html", {
        "character": char,
        "history": history,
        "long_term": long_term,
    })

@app.post("/api/chat/{character_name}")
async def chat_api(character_name: str, request: Request):
    char = load_character(character_name)
    body = await request.json()
    user_message = body.get("message", "")
    
    if not user_message.strip():
        return JSONResponse({"error": "Leere Nachricht"}, status_code=400)
    
    save_message(character_name, "user", user_message)
    
    short_term = get_short_term_memory(character_name, limit=20)
    long_term = get_long_term_memory(character_name)
    messages = build_prompt(char, short_term, long_term, user_message)
    
    model = char.get("model") or DEFAULT_MODEL
    
    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,
                "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()
            # Some models (gemma4) put output in "thinking" field
            if not response_text:
                thinking = data.get("message", {}).get("thinking", "")
                if thinking:
                    response_text = thinking.strip()
    except Exception as e:
        response_text = f"*(Fehler: {e})*"
    
    save_message(character_name, "assistant", response_text)
    
    if needs_summary(character_name):
        asyncio.create_task(generate_summary(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", DEFAULT_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):
    clear_memory(character_name)
    return JSONResponse({"status": "ok"})

@app.get("/api/memory/{character_name}")
async def api_get_memory(character_name: str):
    short = get_short_term_memory(character_name, limit=100)
    long = get_long_term_memory(character_name)
    return JSONResponse({"short_term": short, "long_term": long})

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