import os
from fastapi import FastAPI
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

app = FastAPI(title="Qwen Coder Engine")

# 1. Define Model Settings
MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"

print("Downloading model weights... (This takes a few minutes on first boot)")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)

print("Loading model into memory...")
# We use 4096 context to leave room for code chunks, and 2 threads for HF Free Tier
llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_threads=2,
    n_batch=512,
    verbose=False
)

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 1024
    temperature: float = 0.2 # Low temperature for accurate coding

@app.get("/")
def health_check():
    return {"status": "Heavy Coder is Online and Ready"}

@app.post("/generate")
async def generate(request: GenerateRequest):
    # Format the prompt using Qwen's specific ChatML syntax
    formatted_prompt = f"<|im_start|>user\n{request.prompt}<|im_end|>\n<|im_start|>assistant\n"
    
    def token_generator():
        stream = llm(
            formatted_prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            stream=True
        )
        for output in stream:
            token = output["choices"][0]["text"]
            if token:
                yield {"data": token}

    return EventSourceResponse(token_generator())