from fastapi import FastAPI, Request, Header
from fastapi.responses import JSONResponse
import uvicorn, os, time, json, torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ==== CONFIG ====
MODEL_REPO = os.getenv("MODEL_REPO", "retaj/my-qwen14b-finetune")  # your PRIVATE model
HF_TOKEN   = os.getenv("HF_TOKEN")  # REQUIRED (read token) because repo is private
SERVER_API_KEY = os.getenv("SERVER_API_KEY")  # optional: require Bearer auth from clients
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 256))
TEMPERATURE    = float(os.getenv("TEMPERATURE", 0.7))
TOP_P          = float(os.getenv("TOP_P", 0.95))
PORT           = int(os.getenv("PORT", "7860"))  # HF Spaces uses this

device = torch.device("cpu")  # free tier is CPU-only

# ==== LOAD MODEL ====
def load_model():
    tok = AutoTokenizer.from_pretrained(
        MODEL_REPO, token=HF_TOKEN, use_fast=True, trust_remote_code=True
    )
    mdl = AutoModelForCausalLM.from_pretrained(
        MODEL_REPO,
        torch_dtype=torch.float32,
        device_map=None,
        token=HF_TOKEN,
        trust_remote_code=True,
    ).to(device)
    mdl.eval()
    return tok, mdl

tokenizer, model = load_model()

# ==== Qwen chat template ====
IM_START = "<|im_start|>"
IM_END   = "<|im_end|>"

def build_prompt(messages):
    parts = []
    for m in messages:
        role = m.get("role", "user")
        content = m.get("content", "")
        parts.append(f"{IM_START}{role}\n{content}{IM_END}\n")
    parts.append(f"{IM_START}assistant\n")
    return "".join(parts)

def check_auth(header: str | None):
    if SERVER_API_KEY is None:
        return True
    return header and header.startswith("Bearer ") and header.split(" ",1)[1] == SERVER_API_KEY

# ==== FastAPI app ====
app = FastAPI()

@app.get("/")
def root():
    return {"status": "ok", "model": MODEL_REPO}

@app.post("/v1/chat/completions")
async def chat_completions(request: Request, authorization: str | None = Header(default=None)):
    if not check_auth(authorization):
        return JSONResponse({"error": {"message": "Unauthorized"}}, status_code=401)

    body = await request.json()
    messages = body.get("messages", [])
    temperature = float(body.get("temperature", TEMPERATURE))
    top_p = float(body.get("top_p", TOP_P))
    max_tokens = int(body.get("max_tokens", MAX_NEW_TOKENS))

    if not messages:
        return JSONResponse({"error": {"message": "messages required"}}, status_code=400)

    prompt = build_prompt(messages)
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id,
        )

    decoded = tokenizer.decode(output_ids[0], skip_special_tokens=False)
    marker = f"{IM_START}assistant\n"
    resp = decoded.split(marker)[-1].split(IM_END)[0].strip()

    now = int(time.time())
    return {
        "id": f"chatcmpl-{now}",
        "object": "chat.completion",
        "created": now,
        "model": MODEL_REPO,
        "choices": [
            {"index": 0, "message": {"role": "assistant", "content": resp}, "finish_reason": "stop"}
        ],
        "usage": None,
    }

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=PORT)