import os
import random
import torch
from torch.utils.data import Dataset
from torch.nn.utils.rnn import pad_sequence

from src.utils import setup_logger


logger = setup_logger(__name__)



class ChatterboxDataset(Dataset):
    
    def __init__(self, config):
        self.cfg = config
        self.preprocessed_dir = config.preprocessed_dir
        
        if not os.path.exists(self.preprocessed_dir):
            raise FileNotFoundError(f"Preprocessing folder not found: {self.preprocessed_dir}.")
            
        self.files = [f for f in os.listdir(self.preprocessed_dir) if f.endswith(".pt")]
        
        if len(self.files) == 0:
            raise RuntimeError(f"There are no .pt files in the folder: {self.preprocessed_dir}")
            
        logger.info(f"Dataset loaded. Total sample: {len(self.files)}")

        self.sot_token = config.start_text_token 
        self.eot_token = config.stop_text_token


    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        
        try:
            
            filename = self.files[idx]
            
            pt_path = os.path.join(self.preprocessed_dir, filename)
            
            data = torch.load(pt_path)
            
            
            text_tokens = data["text_tokens"]
            if text_tokens.size(0) > self.cfg.max_text_len - 2:
                text_tokens = text_tokens[:self.cfg.max_text_len - 2]
                
            sot = torch.tensor([self.sot_token], dtype=torch.long)
            eot = torch.tensor([self.eot_token], dtype=torch.long)
            text_tokens = torch.cat([sot, text_tokens, eot])

            speech_tokens = data["speech_tokens"]
            if speech_tokens.size(0) > self.cfg.max_speech_len:
                speech_tokens = speech_tokens[:self.cfg.max_speech_len]

            speaker_emb = data["speaker_emb"]
            prompt_tokens = data["prompt_tokens"]

            if random.random() < 0.20:
                speaker_emb = torch.zeros_like(speaker_emb)
                prompt_tokens = torch.zeros(1, dtype=torch.long)


            return {
                "text_tokens": text_tokens,
                "speech_tokens": speech_tokens,
                "speaker_emb": speaker_emb,
                "prompt_tokens": prompt_tokens
            }


        except Exception as e:
            logger.error(f"Error loading {filename}: {e}")
            return None


def data_collator_standart(batch):

    batch = [item for item in batch if item is not None]
    if not batch:
        return {}

    # Padding
    text_tokens = pad_sequence([x["text_tokens"] for x in batch], batch_first=True, padding_value=0)
    speech_tokens = pad_sequence([x["speech_tokens"] for x in batch], batch_first=True, padding_value=0)
    prompt_tokens = pad_sequence([x["prompt_tokens"] for x in batch], batch_first=True, padding_value=0)
    # print(text_tokens)
    speaker_embs = torch.stack([x["speaker_emb"] for x in batch])

    # Lengths
    text_lens = torch.tensor([len(x["text_tokens"]) for x in batch], dtype=torch.long)
    speech_lens = torch.tensor([len(x["speech_tokens"]) for x in batch], dtype=torch.long)


    return {
        "text_tokens": text_tokens,
        "text_token_lens": text_lens,
        "speech_tokens": speech_tokens,
        "speech_token_lens": speech_lens,
        "speaker_emb": speaker_embs,
        "prompt_tokens": prompt_tokens
    }
    
    


def data_collator_turbo(batch):

    batch = [item for item in batch if item is not None]
    if not batch: 
        return {}

    # 1. Text Tokens Padding
    text_tokens = pad_sequence([x["text_tokens"] for x in batch], batch_first=True, padding_value=0)
    text_lens = torch.tensor([len(x["text_tokens"]) for x in batch], dtype=torch.long)

    # 2. Speech Tokens Padding
    speech_tokens = pad_sequence([x["speech_tokens"] for x in batch], batch_first=True, padding_value=0)
    speech_lens = torch.tensor([len(x["speech_tokens"]) for x in batch], dtype=torch.long)

    # 3. Prompt Tokens Padding
    prompt_tokens = pad_sequence([x["prompt_tokens"] for x in batch], batch_first=True, padding_value=0)
    prompt_lens = torch.tensor([x["prompt_tokens"].shape[0] for x in batch], dtype=torch.long)

    # 4. Speaker Embedding
    speaker_embs = torch.stack([x["speaker_emb"] for x in batch])

    return {
        "text_tokens": text_tokens,
        "text_token_lens": text_lens,
        "speech_tokens": speech_tokens,
        "speech_token_lens": speech_lens,
        "speaker_emb": speaker_embs,
        "prompt_tokens": prompt_tokens,
        "prompt_lens": prompt_lens
    }