from fastapi import FastAPI
from pydantic import BaseModel
import google.generativeai as genai
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Set up Gemini API Key
API_KEY = os.getenv("GEMINI_API_KEY")
if not API_KEY:
    raise ValueError("❌ Missing GEMINI_API_KEY. Please set it in the .env file.")

genai.configure(api_key=API_KEY)

# Load Gemini Model
model = genai.GenerativeModel("gemini-2.0-flash")

# Initialize FastAPI app
app = FastAPI(title="APTS Injury Recovery AI API")

# Define input structure
class InjuryInput(BaseModel):
    injury_type: str
    age: int
    past_injuries: str
    sport: str
    mobility: int  # 1-3 scale
    pressure: int  # 1-3 scale
    weight_bearing: int  # 1-3 scale

# Function to analyze injury and suggest recovery
def generate_recovery_plan(data: InjuryInput):
    severity_level = max(data.mobility, data.pressure, data.weight_bearing)

    recovery_prompt = f"""
    Generate a **personalized recovery plan** for an athlete with the following details:
    - **Injury Type**: {data.injury_type}
    - **Age**: {data.age} years
    - **Sport**: {data.sport}
    - **Past Injuries**: {data.past_injuries}
    - **Injury Severity Level**: {severity_level} (1 = Low, 2 = Medium, 3 = High)

    Provide:
    1️⃣ **Rehabilitation Plan**
    2️⃣ **Estimated Recovery Time**
    3️⃣ **Diet & Supplements**
    4️⃣ **Precautions**
    """

    recovery_response = model.generate_content(recovery_prompt)

    return {
        "injury_type": data.injury_type,
        "age": data.age,
        "sport": data.sport,
        "severity_level": severity_level,
        "recovery_plan": recovery_response.text
    }

# API Endpoint
@app.post("/analyze-injury")
async def analyze_injury(data: InjuryInput):
    response = generate_recovery_plan(data)
    return response