from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from dotenv import load_dotenv

from promptius_gui_schema import PromptiusGuiSchema
import uvicorn

load_dotenv()

app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

llm = ChatOpenAI(model_name="gpt-4.1-mini", temperature=0)
#llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-preview-05-20", temperature=0)
llm_with_struct = llm.with_structured_output(PromptiusGuiSchema)

class GenerateUIRequest(BaseModel):
    prompt: str

@app.get("/health")
def health_check():
    return {"status": "ok"}

@app.post("/generate_ui")
def generate_ui(request: GenerateUIRequest):
    """
    Generates a UI schema based on the user's prompt.
    """
    print("Received prompt:", request.prompt)
    answer: PromptiusGuiSchema = llm_with_struct.invoke([SystemMessage(content="You are a UI generator, you are required to generate UI, even if user is not providing sufficient data you are supposed to generate mock values. Keep the styling compact, use grid when required. You need to ensure that the UI looks good, think like a graphic designer"), HumanMessage(content=request.prompt)])
    print("Generated UI Schema:", answer)
    return answer.model_dump()

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