"""molab_setup.py — Dual LLM server (DC + MC) on molab GPU.
DC: Mistral-Small-3.2-24B-abliterated (chat humano, español, Discord)
MC: Qwen3-14B-abliterated (agentic, tool calling, Minecraft)
"""
import os, sys, re, time, threading, urllib.request
import torch, requests
from transformers import AutoTokenizer, AutoModelForCausalLM
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn

# Install missing deps
os.system("pip install -q sentencepiece tiktoken accelerate > /tmp/pip.log 2>&1")

LOG = "/tmp/setup.log"
TUNNEL_URL_FILE = "/tmp/tunnel_url.txt"
PORT = 8000
CF_PATH = "/usr/local/bin/cloudflared"

# Modelos SIN CENSURA
DC_MODEL = "huihui-ai/Huihui-Mistral-Small-3.2-24B-Instruct-2506-abliterated"
MC_MODEL = "mlabonne/Qwen3-14B-abliterated"

def log(msg):
    line = "[{}] {}".format(time.strftime("%H:%M:%S"), msg)
    try:
        with open(LOG, "a") as f: f.write(line + "\n")
    except: pass
    print(line, flush=True)

open(LOG, "w").close()
log("=== SETUP START ===")
log("DC model: {}".format(DC_MODEL))
log("MC model: {}".format(MC_MODEL))

# 1. cloudflared
if not os.path.exists(CF_PATH) or os.path.getsize(CF_PATH) < 1_000_000:
    log("Downloading cloudflared...")
    urllib.request.urlretrieve(
        "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64",
        CF_PATH)
    os.chmod(CF_PATH, 0o755)
log("cloudflared OK")

# 2. Load DC model (Mistral 24B abliterated)
log("Loading DC model (Mistral 24B abliterated)...")
dc_tokenizer = AutoTokenizer.from_pretrained(DC_MODEL, trust_remote_code=True)
try:
    dc_model = AutoModelForCausalLM.from_pretrained(
        DC_MODEL, dtype=torch.bfloat16, trust_remote_code=True).to("cuda")
except ValueError:
    # Mistral3 needs Mistral3ForConditionalGeneration
    from transformers import Mistral3ForConditionalGeneration
    dc_model = Mistral3ForConditionalGeneration.from_pretrained(
        DC_MODEL, dtype=torch.bfloat16, trust_remote_code=True).to("cuda")
log("DC model loaded! VRAM: {:.2f} GB".format(torch.cuda.memory_allocated()/1024**3))

# 3. Load MC model (Qwen3 14B abliterated)
log("Loading MC model (Qwen3 14B abliterated)...")
mc_tokenizer = AutoTokenizer.from_pretrained(MC_MODEL, trust_remote_code=True)
mc_model = AutoModelForCausalLM.from_pretrained(
    MC_MODEL, dtype=torch.bfloat16, trust_remote_code=True).to("cuda")
log("MC model loaded! VRAM: {:.2f} GB".format(torch.cuda.memory_allocated()/1024**3))

# 4. FastAPI
app = FastAPI(title="Zelin Dual LLM API")
app.add_middleware(CORSMiddleware,
    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

class ChatReq(BaseModel):
    message: str
    max_tokens: int = 200
    temperature: float = 0.7
    model: str = "dc"  # "dc" or "mc"

@app.get("/")
def root():
    return {
        "status": "ok",
        "dc_model": "Mistral-Small-3.2-24B-abliterated",
        "mc_model": "Qwen3-14B-abliterated",
        "vram_gb": round(torch.cuda.memory_allocated()/1024**3, 2),
    }

@app.get("/health")
def health():
    return {
        "status": "healthy",
        "vram_gb": round(torch.cuda.memory_allocated()/1024**3, 2),
        "vram_free_gb": round(torch.cuda.mem_get_info()[0]/1024**3, 2),
    }

def generate_response(model, tokenizer, message, max_tokens, temperature):
    # Máximo contexto posible: 16K tokens (deja margen seguro)
    inputs = tokenizer(message, return_tensors="pt", truncation=True, max_length=16384).to("cuda")
    t0 = time.time()
    with torch.no_grad():
        out = model.generate(**inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            max_length=16384 + max_tokens)  # context + generation
    elapsed = time.time() - t0
    resp = tokenizer.decode(
        out[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True)
    tokens = out.shape[1] - inputs["input_ids"].shape[1]
    return {
        "response": resp,
        "tokens": tokens,
        "elapsed_s": round(elapsed, 2),
        "tokens_per_second": round(tokens/elapsed, 1) if elapsed > 0 else 0,
    }

@app.post("/chat")
def chat(req: ChatReq):
    if req.model == "mc":
        result = generate_response(mc_model, mc_tokenizer, req.message, req.max_tokens, req.temperature)
        result["model"] = "Qwen3-14B-abliterated"
    else:
        result = generate_response(dc_model, dc_tokenizer, req.message, req.max_tokens, req.temperature)
        result["model"] = "Mistral-Small-3.2-24B-abliterated"
    return result

@app.post("/dc")
def dc_chat(req: ChatReq):
    result = generate_response(dc_model, dc_tokenizer, req.message, req.max_tokens, req.temperature)
    result["model"] = "Mistral-Small-3.2-24B-abliterated"
    return result

@app.post("/mc")
def mc_chat(req: ChatReq):
    result = generate_response(mc_model, mc_tokenizer, req.message, req.max_tokens, req.temperature)
    result["model"] = "Qwen3-14B-abliterated"
    return result

# 5. Server background
def run_server():
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="error")
threading.Thread(target=run_server, daemon=True).start()
time.sleep(3)
log("FastAPI on port {}".format(PORT))

# 6. Tunnel
log("Starting tunnel...")
os.system("{} tunnel --url http://localhost:{} --no-autoupdate > /tmp/tunnel.log 2>&1 &".format(CF_PATH, PORT))
tunnel_url = None
for _ in range(15):
    time.sleep(2)
    try:
        l = open("/tmp/tunnel.log").read()
        m = re.search(r'https://[a-z0-9-]+\.trycloudflare\.com', l)
        if m:
            tunnel_url = m.group(0)
            break
    except: pass

# 7. Keep-alive
def keep_alive():
    while True:
        try: requests.get("http://localhost:{}/health".format(PORT), timeout=5)
        except: pass
        time.sleep(60)
threading.Thread(target=keep_alive, daemon=True).start()
log("Keep-alive active")

# 8. Test both models
if tunnel_url:
    log("TUNNEL_URL: {}".format(tunnel_url))
    with open(TUNNEL_URL_FILE, "w") as f:
        f.write(tunnel_url)
    try:
        # Health check
        r = requests.get("{}/health".format(tunnel_url), timeout=15)
        log("Health: {}".format(r.json()))

        # Test DC model
        log("Testing DC model...")
        r2 = requests.post("{}/dc".format(tunnel_url),
            json={"message": "Hola! Eres Zelin? Como estas?"}, timeout=60)
        d = r2.json()
        log("DC: {} tok/s".format(d.get("tokens_per_second", "?")))
        log("DC Response: {}".format(d.get("response", "")[:200]))

        # Test MC model
        log("Testing MC model...")
        r3 = requests.post("{}/mc".format(tunnel_url),
            json={"message": "How do I craft a diamond sword in Minecraft?"}, timeout=60)
        d2 = r3.json()
        log("MC: {} tok/s".format(d2.get("tokens_per_second", "?")))
        log("MC Response: {}".format(d2.get("response", "")[:200]))
    except Exception as e:
        log("Test error: {}".format(e))
else:
    log("Tunnel not ready")

log("=== SETUP DONE ===")