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

THRESHOLD = 0.75
device = "cpu"

class EndpointHandler:
    def __init__(self, model_dir):
        self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_dir)
        self.model = Wav2Vec2ForSequenceClassification.from_pretrained(
            model_dir,
            low_cpu_mem_usage=True
        ).to(device)

        self.model.eval()

    def load_mp3_from_base64(self, 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(self, 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 = self.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 = self.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

    def __call__(self, data):
        # Accept both formats
        if "inputs" in data:
            req = data["inputs"]
        else:
            req = data

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

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

        if ai_prob >= THRESHOLD:
            return {
                "status": "success",
                "language": language,
                "classification": "AI_GENERATED",
                "confidenceScore": round(ai_prob, 4),
                "explanation": "Unnatural pitch consistency and synthetic speech patterns detected"
            }
        else:
            return {
                "status": "success",
                "language": language,
                "classification": "HUMAN",
                "confidenceScore": round(1 - ai_prob, 4),
                "explanation": "Natural pitch variation and human speech characteristics detected"
            }