import joblib
from transformers import AutoFeatureExtractor, Wav2Vec2Model
import torch
import librosa
from scipy.special import expit
import json
import os

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

FEATURE_EXTRACTOR = None
classifier = None
scaler = None
thresh = None

class CustomWav2Vec2Model(Wav2Vec2Model):
    def __init__(self, config):
        super().__init__(config)
        self.encoder.layers = self.encoder.layers[:9]


class HuggingFaceFeatureExtractor:
    def __init__(self, model, feature_extractor_name):
        self.device = device
        self.feature_extractor = AutoFeatureExtractor.from_pretrained(feature_extractor_name)
        self.model = model
        self.model.eval()
        self.model.to(self.device)

    def __call__(self, audio, sr):
        inputs = self.feature_extractor(
            audio,
            sampling_rate=sr,
            return_tensors="pt",
            padding=True,
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        with torch.no_grad():
            outputs = self.model(**inputs, output_hidden_states=True)
        return outputs.hidden_states[9] 


def load_models(folder_name='model'):
    global FEATURE_EXTRACTOR, classifier, scaler, thresh
    try:
        truncated_model = CustomWav2Vec2Model.from_pretrained(os.path.join(folder_name,"wav2vec2-xls-r-2b_truncated"))
        FEATURE_EXTRACTOR = HuggingFaceFeatureExtractor(truncated_model, "facebook/wav2vec2-xls-r-2b")
        classifier, scaler, thresh = joblib.load(os.path.join(folder_name,'logreg_margin_pruning_ALL_with_scaler_threshold.joblib'))

    except Exception as e:
        raise RuntimeError(f"failed to load the models: {e}")


def segment_audio(audio, sr, segment_duration):
    segment_samples = int(segment_duration * sr)
    total_samples = len(audio)
    segments = [audio[i:i + segment_samples] for i in range(0, total_samples, segment_samples)]
    segments_check = []
    for seg in segments:
        # if the segment is shorter than 0.7s, skip it to avoid complications inside wav2vec2
        if len(seg) > 0.7 * sr:
           segments_check.append(seg)
    return segments_check

def process_audio(input_data, segment_duration=60):
    audio, sr = librosa.load(input_data, sr=16000)
    if len(audio.shape) > 1:
        audio = audio[0]
    segments = segment_audio(audio, sr, segment_duration)
    segment_predictions = []
    all_probs = []
    confidence_scores_fake_sum = 0
    fake_segments = 0
    confidence_scores_real_sum = 0
    real_segments = 0
    eer_threshold = thresh - 5e-3 # small margin error due to feature extractor space differences
    for idx, segment in enumerate(segments):
        features = FEATURE_EXTRACTOR(segment, sr)
        features_avg = torch.mean(features, dim=1).cpu().numpy()
        features_avg = features_avg.reshape(1, -1)
        decision_score = classifier.decision_function(features_avg)
        decision_score_scaled = scaler.transform(decision_score.reshape(-1, 1)).flatten()
        decision_value = decision_score_scaled[0]
        all_probs.append(expit(decision_score_scaled).item())
        pred = 1 if decision_value >= eer_threshold else 0
        if pred == 0:
            confidence_percentage = 1 - expit(decision_score).item()
            confidence_scores_fake_sum +=confidence_percentage
            fake_segments +=1
        else:
            confidence_percentage = expit(decision_score).item()
            confidence_scores_real_sum +=confidence_percentage
            real_segments +=1
        segment_predictions.append(pred)
    avg_whole_audio_prob = sum(all_probs) / len(all_probs)
    score_audio = 1 - avg_whole_audio_prob ## the model was trained with asvspoof labels
    output_dict = {
    "label": "real" if sum(segment_predictions) > (len(segment_predictions) / 2) else "fake",
    "confidence score": f'{confidence_scores_real_sum/real_segments:.2f}' if sum(segment_predictions) > (len(segment_predictions) / 2) else f'{confidence_scores_fake_sum/fake_segments:.2f}',
    "score_audio": score_audio,
    "segments_processed": len(all_probs)
        }
    json_output = json.dumps(output_dict, indent=4)
    return json_output