#!/usr/bin/env python3
"""WHOAMISec Wizard-Vicuna-13B — OpenAI-compatible API via Ollama (zero compilation)."""
import os, sys, json, time, subprocess
from starlette.applications import Starlette
from starlette.routing import Route
from starlette.responses import JSONResponse, StreamingResponse
from starlette.middleware.cors import CORSMiddleware
import httpx

MODEL_DIR = "/data/model"
MODEL_FILE = os.path.join(MODEL_DIR, "Wizard-Vicuna-13B-Uncensored.Q4_K_M.gguf")
MODEL_REPO = "TheBloke/Wizard-Vicuna-13B-Uncensored-GGUF"
MODEL_FILENAME = "Wizard-Vicuna-13B-Uncensored.Q4_K_M.gguf"
MODEL_NAME = "whoamisec-wizard-13b"
OLLAMA_URL = "http://localhost:11434"
ready = False

def download_model():
    os.makedirs(MODEL_DIR, exist_ok=True)
    if os.path.exists(MODEL_FILE):
        print(f"[SRV] Model cached: {MODEL_FILE}", flush=True)
        return
    print(f"[SRV] Downloading {MODEL_FILENAME}...", flush=True)
    hf_token = os.environ.get("HF_TOKEN", "")
    from huggingface_hub import hf_hub_download
    path = hf_hub_download(MODEL_REPO, MODEL_FILENAME, local_dir=MODEL_DIR, local_dir_use_symlinks=False, token=hf_token or None)
    print(f"[SRV] Downloaded: {path}", flush=True)

def start_ollama():
    global ready
    # Step 1: Download model FIRST
    download_model()

    # Step 2: Start Ollama server
    print("[SRV] Starting Ollama server...", flush=True)
    subprocess.Popen(
        ["ollama", "serve"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
    )
    for _ in range(30):
        try:
            r = httpx.get(f"{OLLAMA_URL}/api/tags", timeout=2)
            if r.status_code == 200:
                print("[SRV] Ollama ready!", flush=True)
                break
        except:
            pass
        time.sleep(1)

    # Step 3: Generate Modelfile dynamically
    modelfile_content = f'''FROM {MODEL_FILE}

TEMPLATE """{{{{ if .System }}}}{{{{ .System }}}}
{{{{ end }}}}{{{{ range .Messages }}}}{{{{ if eq .Role "user" }}}}
USER: {{{{ .Content }}}}
ASSISTANT: {{{{ end }}}}
{{{{ end }}}}
"""

PARAMETER stop "</s>"
PARAMETER num_ctx 2048
'''
    mf_path = "/tmp/Modelfile"
    with open(mf_path, "w") as f:
        f.write(modelfile_content)

    # Step 4: Create model in Ollama
    print(f"[SRV] Creating model '{MODEL_NAME}' from {MODEL_FILE}...", flush=True)
    result = subprocess.run(
        ["ollama", "create", MODEL_NAME, "-f", mf_path],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        print(f"[SRV] ERROR creating model: {result.stderr[:500]}", flush=True)
        sys.exit(1)
    print(f"[SRV] Model '{MODEL_NAME}' ready!", flush=True)
    ready = True

# Start everything in background (download + ollama setup)
import threading
threading.Thread(target=start_ollama, daemon=True).start()

async def chat_completions(request):
    if not ready:
        return JSONResponse({"error": {"message": "model loading", "type": "server_error"}}, status_code=503)
    try:
        body = await request.json()
        async with httpx.AsyncClient(timeout=600) as client:
            if body.get("stream", False):
                resp = await client.post(
                    f"{OLLAMA_URL}/v1/chat/completions",
                    json={**body, "model": MODEL_NAME},
                    headers={"Content-Type": "application/json"},
                )
                return StreamingResponse(
                    resp.aiter_bytes(),
                    media_type=resp.headers.get("content-type", "text/event-stream"),
                    headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
                )
            else:
                resp = await client.post(
                    f"{OLLAMA_URL}/v1/chat/completions",
                    json={**body, "model": MODEL_NAME},
                    headers={"Content-Type": "application/json"},
                )
                data = resp.json()
                if "model" in data:
                    data["model"] = body.get("model", MODEL_NAME)
                return JSONResponse(data)
    except Exception as e:
        return JSONResponse({"error": {"message": str(e), "type": "server_error"}}, status_code=500)

async def models(request):
    return JSONResponse({
        "object": "list",
        "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "whoamisec",
                  "permissions": [{"allow_sampling": True, "allow_logprobs": True, "allow_search_indices": True, "allow_view": True, "allow_fine_tuning": False}]}],
    })

async def health(request):
    if ready:
        try:
            async with httpx.AsyncClient(timeout=5) as client:
                r = await client.get(f"{OLLAMA_URL}/api/tags")
                if r.status_code == 200:
                    return JSONResponse({"status": "ok", "model": MODEL_NAME})
        except:
            pass
    return JSONResponse({"status": "loading", "model": MODEL_NAME})

app = Starlette(debug=False, routes=[
    Route("/v1/chat/completions", chat_completions, methods=["POST"]),
    Route("/v1/models", models, methods=["GET"]),
    Route("/health", health, methods=["GET"]),
])
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

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