from fastapi import FastAPI, UploadFile, HTTPException
from ffmpeg import input, output, run
import boto3

app = FastAPI()

# Define your AWS S3 credentials and bucket name
aws_access_key_id = 'YOUR_AWS_ACCESS_KEY_ID'
aws_secret_access_key = 'YOUR_AWS_SECRET_ACCESS_KEY'
s3_bucket_name = 'YOUR_S3_BUCKET_NAME'

# Initialize an S3 client
s3_client = boto3.client('s3', aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)

@app.post("/merge_audio/{output_key}")
async def merge_audio(output_key: str):
    # Temporary directory to store uploaded audio files
    temp_dir = '/tmp/'

    # Local paths for the fixed input audio files
    gen_base_file = temp_dir + 'genBase.mp3'
    quote_file = temp_dir + 'quote.mp3'

    # Download files from S3 based on the provided key
    s3_directory = 'your-s3-directory/' + output_key
    s3_gen_base_key = f'{s3_directory}/genBase.mp3'
    s3_quote_key = f'{s3_directory}/quote.mp3'

    try:
        s3_client.download_file(s3_bucket_name, s3_gen_base_key, gen_base_file)
        s3_client.download_file(s3_bucket_name, s3_quote_key, quote_file)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to download files from S3: {e}")

    # Output file name
    output_file = temp_dir + 'output.mp3'

    # Use python-ffmpeg to merge audio clips
    input_args = [input(gen_base_file), input(quote_file)]
    output_args = output(output_file, acodec='copy')

    try:
        run(*input_args, output_args, overwrite_output=True)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"FFmpeg Error: {e}")

    # Upload the merged file back to S3
    s3_output_key = f'{s3_directory}/output.mp3'
    try:
        s3_client.upload_file(output_file, s3_bucket_name, s3_output_key)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to upload file to S3: {e}")

    return {"message": "Audio clips successfully merged and saved in S3"}

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