import io
import tempfile
from fastapi import FastAPI, HTTPException, File, UploadFile
from speechbrain.inference import SpeakerRecognition
from fastapi.openapi.docs import get_swagger_ui_html

# 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))

# Custom route to expose Swagger UI
@app.get("/docs")
def custom_swagger_ui_html():
    return get_swagger_ui_html(openapi_url=app.openapi_url, title="API Docs")

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