import json
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
from pydantic import BaseModel
from typing import List

# Initialize FastAPI
app = FastAPI()

# Enable CORS for frontend communication
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load the model
# We use n_ctx=512 to keep memory footprint low
print("Loading model... please wait.")
MODEL_PATH = hf_hub_download(
    repo_id="Qwen/Qwen2.5-Coder-1.5B-Instruct-GGUF", 
    filename="qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"
)

# Force CPU mode (n_gpu_layers=0) for maximum stability on Free Tier
llm = Llama(
    model_path=MODEL_PATH, 
    n_ctx=512, 
    n_gpu_layers=0, 
    n_threads=1
)
print("Model loaded successfully.")

# Pydantic models for request validation
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[ChatMessage]

# Endpoints
@app.get("/")
def read_root():
    return {"message": "Server is up and waiting for chat!"}

@app.post("/v1/chat/completions")
async def chat_completion(request: ChatRequest):
    print("Received chat request...")
    try:
        # Convert request to history format
        history = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        # Run inference with stream=False to ensure stable connection
        response = llm.create_chat_completion(
            messages=history,
            temperature=0.2,
            stream=False 
        )
        
        print("Inference finished.")
        return response
        
    except Exception as e:
        print(f"ERROR: {e}")
        raise HTTPException(status_code=500, detail=str(e))