app.py

from fastapi import (FastAPI,
                     UploadFile,
                     File,
                     Header,
                     BackgroundTasks,
                     Body,
                     HTTPException)

from fastapi.staticfiles import StaticFiles
from vitpose import VitPose
from dotenv import load_dotenv
from tasks import process_video,process_salto_alto
from fastapi.responses import JSONResponse
from config import AI_API_TOKEN
import logging
import json
from src.exercises.dead_lift import analyze_dead_lift


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI()
vitpose = VitPose()
# vitpose.pipeline.warmup()

load_dotenv() 


app.mount("/static", StaticFiles())

@app.get("/")
def read_root():
    return {"message": "Hello World"}

@app.get("/test")
def test():
    return {"message": "from test"}

@app.post("/upload")
async def upload(background_tasks: BackgroundTasks,
                 file: UploadFile = File(...), 
                 token: str = Header(...),
                 user_id: str = Body(...),
                 player_id: str = Body(...)):
    
    if token != AI_API_TOKEN:
        return JSONResponse(content={"message": "Unauthorized", 
                                     "status": 401})
    
    logger.info("reading contents")
    contents = await file.read()
    
    # Save the file to the local directory
    logger.info("saving file")
    with open(file.filename, "wb") as f:
        f.write(contents)
    
    logger.info(f"file saved {file.filename}")
    
    # Create a clone of the file with content already read
    background_tasks.add_task(process_video, 
                              file.filename,
                              vitpose,
                              user_id,
                              player_id)  

    # Return the file as a response
    return JSONResponse(content={"message": "Video uploaded successfully", 
                                 "status": 200})
    
@app.post("/exercise/salto_alto")
async def upload(background_tasks: BackgroundTasks,
                 file: UploadFile = File(...), 
                 token: str = Header(...),
                 player_data: str = Body(...),
                 repetitions: int|str = Body(...),
                 exercise_id: str = Body(...),
                 ):
    
    
    player_data = json.loads(player_data)

    if token != AI_API_TOKEN:
        
        raise HTTPException(status_code=401, detail="Unauthorized")
    
    logger.info("reading contents")
    contents = await file.read()
    
    # Save the file to the local directory
    logger.info("saving file")
    with open(file.filename, "wb") as f:
        f.write(contents)
    
    logger.info(f"file saved {file.filename}")
    
    # Create a clone of the file with content already read
    background_tasks.add_task(process_salto_alto, 
                              file.filename,
                              vitpose, 
                              player_data,
                              exercise_id,
                              repetitions)  

    # Return the file as a response
    print(f"returning response")
    return JSONResponse(content={"message": "Video uploaded successfully",
                                 "status": 200})



@app.post("/exercise/dead_lift")
async def upload(background_tasks: BackgroundTasks,
                 file: UploadFile = File(...), 
                 token: str = Header(...),
                 player_data: str = Body(...),
                 repetitions: int|str = Body(...),
                 exercise_id: str = Body(...),
                 weight: int|str = Body(...)
                 ):
    
    player_data = json.loads(player_data)

    if token != AI_API_TOKEN:
        
        raise HTTPException(status_code=401, detail="Unauthorized")
    
    logger.info("reading contents")
    contents = await file.read()
    
    # Save the file to the local directory
    logger.info("saving file")
    with open(file.filename, "wb") as f:
        f.write(contents)
    
    logger.info(f"file saved {file.filename}")
    
    
    
    background_tasks.add_task(analyze_dead_lift, 
                              file.filename,
                              repetitions,
                              weight,
                              player_data["height"],
                              vitpose, 
                              player_data["id"],
                              exercise_id)  
    
    return JSONResponse(content={"message": "Video uploaded successfully",
                                 "status": 200})