import base64
import io
import numpy as np
import librosa
import torch
from fastapi import FastAPI
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification

THRESHOLD = 0.75
device = "cpu"

app = FastAPI()

print("Loading model...")

feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("/repository")
model = Wav2Vec2ForSequenceClassification.from_pretrained("/repository").to(device)
model.eval()

print("Model loaded.")

def load_mp3_from_base64(b64):
    audio_bytes = base64.b64decode(b64)
    with io.BytesIO(audio_bytes) as f:
        y, sr = librosa.load(f, sr=16000)
    return y

def predict_chunked(y):
    chunk_len = 16000
    probs = []

    for start in range(0, len(y), chunk_len):
        chunk = y[start:start + chunk_len]
        if len(chunk) < 4000:
            continue

        inputs = feature_extractor(
            chunk,
            sampling_rate=16000,
            return_tensors="pt"
        )

        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            logits = model(**inputs).logits
            p = torch.softmax(logits, dim=1)[0][1].item()

        probs.append(p)

    return float(np.mean(probs)) if probs else 0.0


@app.post("/")
async def predict(data: dict):

    language = data.get("language")
    audio_base64 = data.get("audioBase64")

    y = load_mp3_from_base64(audio_base64)
    ai_prob = predict_chunked(y)

    if ai_prob >= THRESHOLD:
        return {
            "status": "success",
            "language": language,
            "classification": "AI_GENERATED",
            "confidenceScore": round(ai_prob, 4),
            "explanation": "Synthetic patterns detected"
        }
    else:
        return {
            "status": "success",
            "language": language,
            "classification": "HUMAN",
            "confidenceScore": round(1 - ai_prob, 4),
            "explanation": "Natural speech patterns detected"
        }