from fastapi import FastAPI, HTTPException, Depends
from pydantic import BaseModel
from app.llm_service import get_llm_response
import time
import psutil
import os
from fastapi.middleware.cors import CORSMiddleware
from typing import Dict, Any


app = FastAPI()


# CORS Middleware (Restrict Origins in Production!)
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],  # Allows all HTTP methods (GET, POST, PUT, DELETE, etc.)
    allow_headers=["*"],  # Allows all headers
)

class PromptRequest(BaseModel):
    prompt: str
    llm_name: str

class LLMResponse(BaseModel):
    response: str
    response_time: float
    cost: float
    memory_usage: float

@app.post("/generate", response_model=LLMResponse)
async def generate_text(request: PromptRequest) -> LLMResponse:
    """Generates text using the specified LLM."""
    start_time = time.time()
    process = psutil.Process(os.getpid())
    initial_memory = process.memory_info().rss

    try:
        response_data: Dict[str, Any] = await get_llm_response(request.llm_name, request.prompt)  #get_llm_response now returns a dictionary.
    except ValueError as ve:
        raise HTTPException(status_code=400, detail=str(ve)) #Value Error will give 400
    except Exception as e:
        print(f"Unexpected error: {e}")
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}") #All other errors will give 500.

    end_time = time.time()
    response_time = end_time - start_time
    final_memory = process.memory_info().rss
    memory_usage = (final_memory - initial_memory) / (1024 * 1024)  # Memory usage in MB

    #Cost calculation is now taken from get_llm_response()
    cost = response_data["cost"]

    return LLMResponse(response=response_data["response"], response_time=response_time, cost=cost, memory_usage = memory_usage)