import pandas as pd
import os
import librosa
import soundfile as sf
from tqdm import tqdm

# --- Configurações ---
TARGET_SR = 16000

# Caminhos de entrada
BASE_DATA_PATH = '../data/common_voice/cv-corpus-22.0-2025-06-20/pt/'
METADATA_FILE = os.path.join(BASE_DATA_PATH, 'validated.tsv')
AUDIO_CLIPS_PATH = os.path.join(BASE_DATA_PATH, 'clips')

# Caminho de saída
OUTPUT_PATH = '../dataset_preparado/'
OUTPUT_BR_PATH = os.path.join(OUTPUT_PATH, 'pt_br')
OUTPUT_PT_PATH = os.path.join(OUTPUT_PATH, 'pt_pt')

# --- Mapeamento de Labels ---
VARIANT_MAP = {
    'Portuguese (Brasil)': 'pt_br',
    'Portuguese (Portugal)': 'pt_pt'
}
PATH_MAP = {
    'pt_br': OUTPUT_BR_PATH,
    'pt_pt': OUTPUT_PT_PATH
}

def preprocess_common_voice():
    print("Iniciando o pré-processamento do Common Voice...")

    os.makedirs(OUTPUT_BR_PATH, exist_ok=True)
    os.makedirs(OUTPUT_PT_PATH, exist_ok=True)

    try:
        df = pd.read_csv(METADATA_FILE, sep='\t')
    except FileNotFoundError:
        print(f"Erro: Arquivo de metadados não encontrado em {METADATA_FILE}")
        return

    df = df[['path', 'variant']].dropna()
    df = df[df['variant'].isin(VARIANT_MAP.keys())]
    df['label'] = df['variant'].map(VARIANT_MAP)

    print(f"Amostras encontradas por sotaque (antes do balanceamento):\n{df['label'].value_counts()}")

    # --- Lógica de Amostragem Aleatória e Balanceamento ---
    class_counts = df['label'].value_counts().to_dict()
    if not class_counts or len(class_counts) < 2:
        print("Erro: Não foram encontradas amostras suficientes de ambas as classes para balancear.")
        return
        
    min_samples = min(class_counts.values())
    print(f"Classe minoritária tem {min_samples} amostras. Usando este valor para o balanceamento.")

    # Embaralhar o dataframe para garantir amostragem aleatória
    df = df.sample(frac=1).reset_index(drop=True)
    
    # O limite agora é dinâmico
    max_samples_per_class = min_samples 
    counters = {'pt_br': 0, 'pt_pt': 0}

    for _, row in tqdm(df.iterrows(), total=df.shape[0], desc="Processando áudios"):
        label = row['label']
        
        if counters[label] >= max_samples_per_class:
            continue

        source_audio_path = os.path.join(AUDIO_CLIPS_PATH, row['path'])
        filename = os.path.splitext(row['path'])[0] + '.wav'
        dest_path = os.path.join(PATH_MAP[label], f"cv_{filename}")

        try:
            audio, sr = librosa.load(source_audio_path, sr=TARGET_SR, mono=True)
            sf.write(dest_path, audio, TARGET_SR)
            counters[label] += 1
        except Exception as e:
            print(f"Aviso: Não foi possível processar {source_audio_path}. Erro: {e}")

    print("\nPré-processamento do Common Voice concluído!")
    print(f"Amostras salvas (balanceado): {counters}")

if __name__ == '__main__':
    preprocess_common_voice()