import os, json, time, threading
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from huggingface_hub import hf_hub_download

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://aibruh-crane-builder-v1.static.hf.space",
                   "https://aibruh-crane-builder-v1.hf.space",
                   "https://huggingface.co", "*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

llm = None
status = {"ready": False, "msg": "Downloading model..."}

def load():
    global llm, status
    try:
        path = hf_hub_download(
            repo_id="mradermacher/Qwen2.5-Coder-1.5B-Instruct-abliterated-i1-GGUF",
            filename="Qwen2.5-Coder-1.5B-Instruct-abliterated.i1-Q4_K_M.gguf",
            cache_dir="/data/models"
        )
        from llama_cpp import Llama
        llm = Llama(model_path=path, n_ctx=8192, n_threads=4, verbose=False)
        status = {"ready": True, "msg": "ready"}
    except Exception as e:
        status = {"ready": False, "msg": str(e)}

threading.Thread(target=load, daemon=True).start()

@app.get("/")
def root():
    return status

@app.get("/v1/models")
def models():
    return {"object": "list", "data": [{"id": "Qwen2.5-Coder-1.5B-abliterated", "object": "model"}]}

@app.post("/v1/chat/completions")
async def chat(req: Request):
    if not llm:
        return JSONResponse({"error": status["msg"]}, status_code=503)
    body = await req.json()
    msgs = body.get("messages", [])
    stream = body.get("stream", False)
    prompt = "\n".join(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>" for m in msgs) + "\n<|im_start|>assistant\n"
    max_tokens = body.get("max_tokens", 2048)

    if stream:
        def gen():
            for chunk in llm(prompt, max_tokens=max_tokens, stream=True, stop=["<|im_end|>"]):
                token = chunk["choices"][0]["text"]
                data = json.dumps({"choices": [{"delta": {"content": token}, "finish_reason": None}]})
                yield f"data: {data}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(gen(), media_type="text/event-stream")
    else:
        out = llm(prompt, max_tokens=max_tokens, stop=["<|im_end|>"])
        text = out["choices"][0]["text"]
        return {"choices": [{"message": {"role": "assistant", "content": text}}]}