import os
import torch
import torchaudio
import pandas as pd
from tqdm import tqdm
from src.chatterbox_.tts import ChatterboxTTS, punc_norm
from src.chatterbox_.models.s3tokenizer import S3_SR
from src.utils import setup_logger



logger = setup_logger(__name__)

def preprocess_dataset_ljspeech(config, tts_engine: ChatterboxTTS):
    
    data = pd.read_csv(config.csv_path, sep="|", header=None, quoting=3)
    
    os.makedirs(config.preprocessed_dir, exist_ok=True)


    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    tts_engine.ve.to(device)
    tts_engine.s3gen.to(device)
    
    logger.info(f"Processing dataset... Total: {len(data)}")
    
    success_count = 0
    
    SPEECH_STOP_ID = getattr(tts_engine.t3.hp, 'stop_speech_token', 6562)
    for idx, row in tqdm(data.iterrows(), total=len(data)):
        
        try:
            
            filename = str(row[0])
            if not filename.endswith(".wav"): 
                filename += ".wav"
            
            wav_path = os.path.join(config.wav_dir, filename)
            
            if not os.path.exists(wav_path): 
                continue


            wav, sr = torchaudio.load(wav_path)
            
            if wav.shape[0] > 1: 
                wav = wav.mean(dim=0, keepdim=True)
                
            if sr != S3_SR:
                resampler = torchaudio.transforms.Resample(sr, S3_SR)
                wav = resampler(wav)
            
            wav = wav.to(device)


            with torch.no_grad():

                wav_np = wav.cpu().squeeze().numpy()
                
                spk_emb_np = tts_engine.ve.embeds_from_wavs([wav_np], sample_rate=S3_SR)
                speaker_emb = torch.from_numpy(spk_emb_np[0]).cpu()


                s_tokens, _ = tts_engine.s3gen.tokenizer(wav.unsqueeze(0))
                raw_speech_tokens = s_tokens.squeeze().cpu()
                stop_speech_tensor = torch.tensor([SPEECH_STOP_ID], dtype=raw_speech_tokens.dtype)
                speech_tokens = torch.cat([raw_speech_tokens, stop_speech_tensor], dim=0)


                prompt_samples = int(config.prompt_duration * S3_SR)
                
                if wav.shape[1] < prompt_samples:
                    prompt_wav = torch.nn.functional.pad(wav, (0, prompt_samples - wav.shape[1]))
                    
                else:
                    prompt_wav = wav[:, :prompt_samples]
                
                p_tokens, _ = tts_engine.s3gen.tokenizer(prompt_wav.unsqueeze(0))
                prompt_tokens = p_tokens.squeeze().cpu()


            raw_text = str(row[2]) if len(row) > 2 else str(row[1])
            
            clean_text = punc_norm(raw_text)

            # Tokenizer
            if config.is_turbo:
                token_output = tts_engine.tokenizer(clean_text, return_tensors="pt")
                raw_text_tokens = token_output.input_ids[0].cpu()
                
                if tts_engine.tokenizer.eos_token_id is not None:
                    text_eos = torch.tensor([tts_engine.tokenizer.eos_token_id], dtype=raw_text_tokens.dtype)
                    text_tokens = torch.cat([raw_text_tokens, text_eos], dim=0)
                else:
                    text_tokens = raw_text_tokens
            
            else:
                text_tokens = tts_engine.tokenizer.text_to_tokens(clean_text).squeeze(0).cpu()


            save_path = os.path.join(config.preprocessed_dir, filename.replace(".wav", ".pt"))
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            
            torch.save({
                "speech_tokens": speech_tokens,
                "speaker_emb": speaker_emb,
                "prompt_tokens": prompt_tokens,
                "text_tokens": text_tokens
            }, save_path)
            
            success_count += 1
        
        except Exception as e:
            logger.error(f"Error ({filename}): {e}")
            continue
        
    logger.info(f"Preprocessing completed! Success: {success_count}/{len(data)}")