#!/usr/bin/env python3
import os, sys, time, threading
from pathlib import Path
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import uvicorn

MODEL_PATH = Path(os.getenv("HOME", "/home/user")) / "models" / os.getenv("MODEL_FILE", "qwen2.5-7b-instruct-q3_k_m.gguf")
MODEL_NAME = os.getenv("SERVED_MODEL_NAME", "qwen")
API_KEY = os.getenv("API_KEY", "")

app = FastAPI(title="Qwen 2.5 API")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
                   allow_methods=["*"], allow_headers=["*"])

llm = None
llm_lock = threading.Lock()

def get_llm():
    global llm
    if llm is not None:
        return llm
    with llm_lock:
        if llm is not None:
            return llm
        print(f"Loading model from {MODEL_PATH}...", flush=True)
        t0 = time.time()
        from llama_cpp import Llama
        llm = Llama(
            model_path=str(MODEL_PATH),
            n_gpu_layers=0,
            n_ctx=int(os.getenv("N_CTX", "8192")),
            n_threads=int(os.getenv("N_THREADS", "2")),
            n_batch=int(os.getenv("N_BATCH", "256")),
            n_ubatch=int(os.getenv("N_UBATCH", "128")),
            verbose=False,
        )
        print(f"Model loaded in {time.time()-t0:.1f}s", flush=True)
        return llm

class Message(BaseModel):
    role: str; content: str

class ChatRequest(BaseModel):
    model: str = MODEL_NAME
    messages: List[Message]
    max_tokens: Optional[int] = 512
    temperature: Optional[float] = 0.7

async def check_auth(request: Request):
    if API_KEY:
        auth = request.headers.get("Authorization", "")
        if auth != f"Bearer {API_KEY}":
            raise HTTPException(status_code=401, detail="Invalid API key")

@app.get("/")
@app.get("/v1/models")
async def models():
    return {"data": [{"id": MODEL_NAME, "object": "model", "owned_by": "local"}]}

@app.get("/health")
async def health():
    return {"status": "ok"}

@app.post("/v1/chat/completions")
async def chat(req: ChatRequest, request: Request):
    await check_auth(request)
    instance = get_llm()
    messages = [{"role": m.role, "content": m.content} for m in req.messages]
    r = instance.create_chat_completion(
        messages=messages, max_tokens=req.max_tokens,
        temperature=req.temperature,
    )
    return {
        "id": "chatcmpl-1", "object": "chat.completion",
        "model": req.model,
        "choices": [{"index": 0, "message": r["choices"][0]["message"],
                    "finish_reason": r["choices"][0]["finish_reason"]}]
    }

if __name__ == "__main__":
    print("Starting uvicorn (lazy model load)...", flush=True)
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")