import os
import uuid
import json
import sqlite3
import threading
import numpy as np
import uvicorn
import gradio as gr
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse
from pydantic import BaseModel
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans

PORT = int(os.environ.get("PORT", 7860))
STATIC_DIR = "./static"
DB_PATH = "./neuraxon.db"

class Centroid3D(BaseModel):
    id: int
    position: list

class PuzzleSubmission(BaseModel):
    puzzle_id: str
    player_centroids: list[Centroid3D]
    player_name: str = "Satoshi_Quant"

PUZZLE_CACHE = {}
MP_ROOMS: dict[str, dict[str, dict]] = {}

def calculate_snr_db(points, centroids, labels):
    signal_power = float(np.mean(points ** 2))
    if signal_power == 0:
        return -999.0
    reconstructed = centroids[labels]
    noise_power = float(np.mean((points - reconstructed) ** 2))
    if noise_power <= 1e-12:
        return 999.0
    return float(10 * np.log10(signal_power / noise_power))

def nearest_cluster_labels(points: np.ndarray, centroids: np.ndarray) -> np.ndarray:
    """Assign each point to its nearest centroid (matches client Slam / drag math)."""
    dists = np.linalg.norm(points[:, None, :] - centroids[None, :, :], axis=2)
    return np.argmin(dists, axis=1)

def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    conn = get_db()
    conn.executescript("""
        CREATE TABLE IF NOT EXISTS recipes (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            puzzle_id TEXT NOT NULL,
            tensor_source TEXT NOT NULL,
            player_name TEXT,
            baseline_snr_db REAL NOT NULL,
            player_snr_db REAL NOT NULL,
            snr_delta_db REAL NOT NULL,
            trust_level TEXT DEFAULT 'THEORETICAL',
            centroid_positions TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        CREATE INDEX IF NOT EXISTS idx_player ON recipes(player_name);
        CREATE TABLE IF NOT EXISTS players (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT UNIQUE NOT NULL,
            synaptic_weight REAL DEFAULT 0.0,
            total_snr_improved REAL DEFAULT 0.0,
            recipes_submitted INTEGER DEFAULT 0,
            highest_snr_delta REAL DEFAULT 0.0
        );
    """)
    conn.commit()
    conn.close()

DOMAIN_CONFIG = {
    "qwen": {"num_vectors": 1000, "dimensions": 256, "num_clusters": 6,
             "baseline_snr": 3.2, "target_snr": 11.5},
    "mri": {"num_vectors": 800, "dimensions": 128, "num_clusters": 8,
            "baseline_snr": 4.1, "target_snr": 13.8},
    "genomic": {"num_vectors": 1200, "dimensions": 64, "num_clusters": 4,
                "baseline_snr": 5.6, "target_snr": 14.2},
    "neural": {"num_vectors": 600, "dimensions": 512, "num_clusters": 10,
               "baseline_snr": 2.1, "target_snr": 10.9}
}

def generate_puzzle(domain: str) -> dict:
    cfg = DOMAIN_CONFIG.get(domain, DOMAIN_CONFIG["qwen"])
    n, dims, k = cfg["num_vectors"], cfg["dimensions"], cfg["num_clusters"]
    raw = np.zeros((n, dims), dtype=np.float32)
    centers = np.random.randn(k, dims).astype(np.float32) * 5.0
    for i in range(n):
        raw[i] = centers[i % k] + np.random.randn(dims).astype(np.float32) * 1.5
    pca = PCA(n_components=3)
    pts_3d = pca.fit_transform(raw).astype(np.float32)
    km = KMeans(n_clusters=k, n_init=10, random_state=42)
    labels = km.fit_predict(pts_3d)
    base_cents = km.cluster_centers_.astype(np.float32)
    base_snr = calculate_snr_db(pts_3d, base_cents, labels)
    # Target is always baseline + domain improvement delta (not a fixed 11.5 dB)
    improve_db = float(cfg["target_snr"] - cfg["baseline_snr"])
    target_snr = round(base_snr + improve_db, 2)
    pid = str(uuid.uuid4())
    PUZZLE_CACHE[pid] = {
        "points_3d": pts_3d, "labels": labels,
        "baseline_centroids": base_cents, "baseline_snr": base_snr,
        "target_snr": target_snr, "improve_db": improve_db,
        "domain": domain, "config": cfg
    }
    return {
        "puzzle_id": pid,
        "explained_variance_ratio": round(float(np.sum(pca.explained_variance_ratio_)), 4),
        "baseline_snr_db": round(base_snr, 2),
        "target_snr_db": target_snr,
        "improve_db": round(improve_db, 2),
        "points": [{"pos": pts_3d[i].tolist(), "cluster": int(labels[i])} for i in range(n)],
        "centroids": [{"id": i, "position": base_cents[i].tolist()} for i in range(k)],
        "data_domain": domain
    }

api = FastAPI(title="VectorVoid", version="1.0.0")
api.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

def mp_players_list(room: dict[str, dict]) -> list:
    return [
        {"id": pid, "name": p["name"], "x": p["x"], "y": p["y"], "z": p["z"], "rotY": p["rotY"]}
        for pid, p in room.items()
    ]

async def mp_broadcast(room_id: str, payload: dict, except_id: str | None = None):
    room = MP_ROOMS.get(room_id, {})
    msg = json.dumps(payload)
    dead = []
    for pid, p in room.items():
        if pid == except_id:
            continue
        try:
            await p["ws"].send_text(msg)
        except Exception:
            dead.append(pid)
    for pid in dead:
        room.pop(pid, None)

@api.websocket("/ws/multiplayer")
async def multiplayer_ws(websocket: WebSocket):
    await websocket.accept()
    room_id = None
    player_id = None
    try:
        while True:
            data = json.loads(await websocket.receive_text())
            if data.get("type") == "join":
                room_id = str(data.get("room", "")).upper()[:8]
                if not room_id:
                    await websocket.send_text(json.dumps({"type": "error", "message": "Room required"}))
                    continue
                room = MP_ROOMS.setdefault(room_id, {})
                if len(room) >= 2:
                    await websocket.send_text(json.dumps({"type": "full"}))
                    await websocket.close()
                    return
                player_id = f"p_{uuid.uuid4().hex[:8]}"
                room[player_id] = {
                    "name": data.get("name") or "Player",
                    "x": 0.0, "y": 2.4, "z": 18.0, "rotY": 0.0,
                    "ws": websocket,
                }
                await websocket.send_text(json.dumps({
                    "type": "joined",
                    "room": room_id,
                    "playerId": player_id,
                    "players": mp_players_list(room),
                }))
                await mp_broadcast(room_id, {"type": "players", "players": mp_players_list(room)}, player_id)
            elif data.get("type") == "update" and room_id and player_id:
                room = MP_ROOMS.get(room_id, {})
                p = room.get(player_id)
                if not p:
                    continue
                p["x"] = float(data.get("x", p["x"]))
                p["y"] = float(data.get("y", p["y"]))
                p["z"] = float(data.get("z", p["z"]))
                p["rotY"] = float(data.get("rotY", p["rotY"]))
                await mp_broadcast(room_id, {"type": "players", "players": mp_players_list(room)}, player_id)
    except WebSocketDisconnect:
        pass
    finally:
        if room_id and player_id:
            room = MP_ROOMS.get(room_id, {})
            room.pop(player_id, None)
            if room:
                await mp_broadcast(room_id, {"type": "players", "players": mp_players_list(room)})
            else:
                MP_ROOMS.pop(room_id, None)

@api.get("/api/health")
def health():
    return {"status": "VectorVoid Online", "version": "1.0.0", "gamepad": "PS4 supported"}

@api.get("/api/puzzle/get")
def get_puzzle(domain: str = "qwen"):
    if domain not in DOMAIN_CONFIG:
        raise HTTPException(400, f"Unknown domain. Use: {list(DOMAIN_CONFIG.keys())}")
    return generate_puzzle(domain)

@api.post("/api/puzzle/submit")
def submit_puzzle(sub: PuzzleSubmission):
    pid = sub.puzzle_id
    if pid not in PUZZLE_CACHE:
        raise HTTPException(404, "Puzzle expired. Generate a new one.")
    cache = PUZZLE_CACHE[pid]
    pts, base_snr = cache["points_3d"], cache["baseline_snr"]
    domain = cache["domain"]
    pd = {c.id: np.array(c.position, dtype=np.float32) for c in sub.player_centroids}
    k = cache["config"]["num_clusters"]
    pc = np.zeros((k, 3), dtype=np.float32)
    for i in range(k):
        pc[i] = pd.get(i, cache["baseline_centroids"][i])
    labels = nearest_cluster_labels(pts, pc)
    ps = calculate_snr_db(pts, pc, labels)
    delta = ps - base_snr
    improved = delta > 0.001
    if improved:
        conn = get_db()
        try:
            conn.execute("""INSERT INTO recipes (puzzle_id,tensor_source,player_name,
                baseline_snr_db,player_snr_db,snr_delta_db,trust_level,centroid_positions)
                VALUES (?,?,?,?,?,?,?,?)""",
                (pid, domain, sub.player_name, round(base_snr, 4), round(ps, 4),
                 round(delta, 4), 'THEORETICAL',
                 json.dumps([c.position for c in sub.player_centroids])))
            conn.execute("""INSERT INTO players (name,synaptic_weight,total_snr_improved,recipes_submitted,highest_snr_delta)
                VALUES (?,0.1,?,1,?)
                ON CONFLICT(name) DO UPDATE SET
                    synaptic_weight=synaptic_weight+0.1,
                    total_snr_improved=total_snr_improved+?,
                    recipes_submitted=recipes_submitted+1,
                    highest_snr_delta=MAX(highest_snr_delta,?)""",
                (sub.player_name, delta, delta, delta, delta))
            conn.commit()
        finally:
            conn.close()
    return {"status": "success", "is_improvement": improved,
            "baseline_snr_db": round(base_snr, 2),
            "player_snr_db": round(ps, 2),
            "snr_delta_db": round(delta, 4)}

@api.get("/api/leaderboard")
def leaderboard(domain: str = "all", limit: int = 20):
    conn = get_db()
    try:
        if domain == "all":
            rows = conn.execute("""SELECT player_name,MAX(player_snr_db)as best,COUNT(*)as n
                FROM recipes GROUP BY player_name ORDER BY best DESC LIMIT ?""", (limit,)).fetchall()
        else:
            rows = conn.execute("""SELECT player_name,MAX(player_snr_db)as best,COUNT(*)as n
                FROM recipes WHERE tensor_source=? GROUP BY player_name ORDER BY best DESC LIMIT ?""",
                (domain, limit)).fetchall()
        return {"leaderboard": [{"player": r["player_name"], "best_snr": round(r["best"], 2), "games": r["n"]} for r in rows]}
    finally:
        conn.close()

@api.get("/api/domains")
def domains():
    return {"domains": [{"id": k, "baseline": v["baseline_snr"], "target": v["target_snr"]} for k, v in DOMAIN_CONFIG.items()]}

@api.get("/api/impact")
def impact():
    conn = get_db()
    try:
        total = conn.execute("SELECT COUNT(*)as c FROM recipes").fetchone()["c"]
        total_snr = conn.execute("SELECT COALESCE(SUM(snr_delta_db),0)as s FROM recipes").fetchone()["s"]
        return {"metrics": {
            "mri_slices_compressed": 14262 + total * 3,
            "clinics_equipped": 146,
            "genomes_mapped_bp": 360531 + total * 50,
            "comms_lines_meters": 984131 + total * 100,
            "bandwidth_saved_gb": round(1870.1 + total_snr * 0.5, 1),
            "active_grid": "QWEN TENSOR",
            "total_submissions": total,
            "total_snr_improved_db": round(total_snr, 2)
        }}
    finally:
        conn.close()

@api.get("/", include_in_schema=False)
def root():
    idx = os.path.join(STATIC_DIR, "index.html")
    if os.path.exists(idx):
        return FileResponse(idx)
    return HTMLResponse("""
    <!DOCTYPE html>
    <html style="margin:0;padding:0;height:100vh;background:#050508;">
    <head><meta charset="UTF-8"><title>VectorVoid Grid Wars</title></head>
    <body style="margin:0;display:flex;align-items:center;justify-content:center;height:100vh;font-family:monospace;">
    <div style="text-align:center;">
    <h1 style="color:#FF6B35;font-size:4em;margin:0;text-shadow:0 0 30px #FF6B35;">VECTORVOID</h1>
    <h2 style="color:#00B4D8;font-size:1.5em;letter-spacing:0.3em;">THE GRID WAR</h2>
    <p style="color:#888;">Every Shot Advances Science</p>
    <p style="color:#FF6B35;margin-top:20px;">Backend Online. Upload static/ for full game.</p>
    </div></body></html>""")

@api.get("/{path:path}", include_in_schema=False)
def spa(path: str):
    idx = os.path.join(STATIC_DIR, "index.html")
    fp = os.path.join(STATIC_DIR, path)
    if os.path.exists(fp) and os.path.isfile(fp):
        return FileResponse(fp)
    # Missing game assets must 404 — returning index.html breaks GLTF/GLB loaders.
    if path.startswith("assets/"):
        raise HTTPException(404, f"Asset not found: {path}")
    if os.path.exists(idx):
        return FileResponse(idx)
    return HTMLResponse("<h1>404</h1>")

init_db()

if os.environ.get("SPACE_ID"):
    threading.Thread(
        target=lambda: uvicorn.run(api, host="0.0.0.0", port=PORT, log_level="warning"),
        daemon=True,
    ).start()

with gr.Blocks(title="VectorVoid Grid Wars", css="footer {visibility: hidden}") as demo:
    gr.HTML(
        '<iframe src="/" style="width:100vw;height:100vh;border:none;margin:0;padding:0;" '
        'allow="gamepad *"></iframe>'
    )

demo.queue()
demo.launch(server_name="0.0.0.0", server_port=PORT)