import io
import tempfile
from fastapi import FastAPI, HTTPException, File, UploadFile
from speechbrain.inference import SpeakerRecognition
import gradio as gr

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

# FastAPI app
app = FastAPI()

# Function to calculate similarity score
def get_similarity(file1_path, file2_path):
    try:
        # Use `verify_files` to compare the audio files
        score, prediction = verification.verify_files(file1_path, file2_path)

        # Return the result as a dictionary
        return {
            "Similarity Score": f"{score:.4f}",
            "Same User Prediction": "Yes" if prediction else "No"
        }

    except Exception as e:
        return {"error": str(e)}

# API function to compare voices
@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.
    """
    try:
        # Create temporary files for the uploaded audio
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmpfile1, \
             tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmpfile2:
            
            # Write audio data to the temporary files
            tmpfile1.write(await file1.read())
            tmpfile2.write(await file2.read())

            # Get the file paths
            file1_path = tmpfile1.name
            file2_path = tmpfile2.name

        # Call the get_similarity function with file paths
        result = get_similarity(file1_path, file2_path)
        
        return result

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

# Gradio interface function
def gradio_interface():
    return gr.Interface(
        fn=compare_voices_api,  # FastAPI function is wrapped here
        inputs=[
            gr.Audio(type="filepath", label="First Audio File"),  # Audio file input
            gr.Audio(type="filepath", label="Second Audio File")  # Audio file input
        ],
        outputs="json",  # Output as JSON
        live=False  # No live interface, just the API
    )

# Launch Gradio interface
@app.on_event("startup")
async def startup():
    gr.Interface(fn=compare_voices_api, inputs=[
        gr.Audio(type="filepath", label="First Audio File"),  # Audio file input
        gr.Audio(type="filepath", label="Second Audio File")  # Audio file input
    ], 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)