import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from model import generate_code

app = FastAPI()

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 256

@app.get("/")
def home():
    return {"message": "Code Generation API is running!"}

@app.post("/generate")
def generate(request: GenerateRequest):
    if not request.prompt:
        raise HTTPException(status_code=400, detail="Prompt cannot be empty.")

    try:
        code = generate_code(request.prompt, request.max_tokens)
        return {"generated_code": code}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")

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