from fastapi import FastAPI, UploadFile
from fastapi.middleware.cors import CORSMiddleware

import utils.handle as handle

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.post("/training_data_from_utags_json")
async def training_data_from_utags_json(file: UploadFile, savePath: str):
    return handle.trainingDataFromUTagsJSON(file, savePath)


@app.post("/training_data_from_prompts_for_bert")
async def training_data_from_prompts_for_bert(file: UploadFile, savePath: str):
    return handle.trainingDataFromPromptsForBERT(file, savePath)


@app.post("/augment_data_using_vector_space_algorithm")
async def augment_data_using_vector_space_algorithm(file: UploadFile, savePath: str):
    return handle.augmentDataUsingVectorSpaceAlgorithm(file, savePath)


@app.post("/get_symptoms_causes_and_disease_name_from_json")
async def get_symptoms_causes_and_disease_name_from_json(file: UploadFile, savePath: str):
    return handle.getSymptomsCausesAndDiseaseNameFromJSON(file, savePath)


@app.post("/train_model_on_sagemaker")
async def train_model_on_sagemaker(trainDataPath: str, testDataPath: str, file: UploadFile | None = None):
    return handle.trainModelOnSageMaker(trainDataPath, testDataPath, file)