import os
import torchaudio
import gradio as gr
from fastapi import FastAPI, HTTPException, File, UploadFile
from speechbrain.inference import SpeakerRecognition
from fastapi.responses import JSONResponse

# Initialize the speaker verification model
speaker_verification = SpeakerRecognition.from_hparams(
    source="speechbrain/spkrec-ecapa-voxceleb",
    savedir="tmp_model"
)

# Temporary folder to save uploaded files
UPLOAD_FOLDER = "uploaded_audio"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Function to calculate similarity score
def get_similarity(audio_path1: str, audio_path2: str):
    try:
        # Load audio files
        signal1, _ = torchaudio.load(audio_path1)
        signal2, _ = torchaudio.load(audio_path2)

        # Get similarity score and prediction
        score, prediction = speaker_verification.verify_batch(signal1, signal2)
        return float(score), "Yes" if prediction else "No"
    except Exception as e:
        return str(e), None
    finally:
        # Clean up temporary files
        if os.path.exists(audio_path1):
            os.remove(audio_path1)
        if os.path.exists(audio_path2):
            os.remove(audio_path2)

# API function to compare voices
def compare_voices(file1, file2):
    # Save uploaded files temporarily
    file1_path = os.path.join(UPLOAD_FOLDER, file1.name)
    file2_path = os.path.join(UPLOAD_FOLDER, file2.name)

    with open(file1_path, "wb") as f1:
        f1.write(file1.read())
    with open(file2_path, "wb") as f2:
        f2.write(file2.read())

    # Get similarity score
    score, is_same_user = get_similarity(file1_path, file2_path)

    if is_same_user is None:
        return "Error: " + score  # This will return the error message

    return {"Similarity Score": f"{score:.4f}", "Same User Prediction": is_same_user}

# FastAPI app
app = FastAPI()

@app.post("/compare_voices/")
async def compare_voices_api(file1: UploadFile = File(...), file2: UploadFile = File(...)):
    """
    Compare two audio files and return the similarity score and prediction.
    """
    # Save uploaded files temporarily
    file1_path = os.path.join(UPLOAD_FOLDER, file1.filename)
    file2_path = os.path.join(UPLOAD_FOLDER, file2.filename)

    with open(file1_path, "wb") as f1:
        f1.write(await file1.read())
    with open(file2_path, "wb") as f2:
        f2.write(await file2.read())

    # Get similarity score
    score, is_same_user = get_similarity(file1_path, file2_path)

    if is_same_user is None:
        raise HTTPException(status_code=500, detail="Error in processing files: " + score)

    return JSONResponse(content={"Similarity Score": f"{score:.4f}", "Same User Prediction": is_same_user})

# Gradio interface function
def gradio_interface():
    return gr.Interface(
        fn=compare_voices,
        inputs=[
            gr.Audio(type="file", label="First Audio File"),  # Updated to use `type="file"`
            gr.Audio(type="file", label="Second Audio File")  # Updated to use `type="file"`
        ],
        outputs="json",  # Output results as JSON
        live=False  # No live interface, just the API
    )

# Launch Gradio as a web interface
@app.on_event("startup")
async def startup():
    gr.Interface(fn=compare_voices, inputs=[
        gr.Audio(type="file", label="First Audio File"),  # Updated to use `type="file"`
        gr.Audio(type="file", label="Second Audio File")  # Updated to use `type="file"`
    ], outputs="json", live=False).launch(share=True, inline=True)

# Running the FastAPI app with Gradio
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000)