from dataclasses import dataclass, field
from typing import List

@dataclass
class TrainConfig:
    # --- Paths ---
    # Directory where setup.py downloaded the files
    model_dir: str = "./pretrained_models"
    
    # Path to your metadata CSV (Format: ID|RawText|NormText)
    csv_path: str = "./MyTTSDataset/metadata.csv"
    #metadata_path: str = "./metadata.json"
    
    # Directory containing WAV files
    wav_dir: str = "./MyTTSDataset/wavs"
    #wav_dir: str = "./FileBasedDataset"
    
    preprocessed_dir = "./MyTTSDataset/preprocess"
    #preprocessed_dir = "./FileBasedDataset/preprocess"
    
    # Output directory for the finetuned model
    output_dir: str = "./chatterbox_output"
    
    is_inference = False
    inference_prompt_path: str = "./speaker_reference/reference.wav"
    # Reusing the Hindi "[hi]" tag for Nepali (see prepare_dataset.py's docstring for why)
    # -- must match what preprocess_ljspeech.py/inference.py feed to the tokenizer.
    inference_test_text: str = "[hi]त्यो त राम्रो भयो। अनि, मसँग कुरा गर्नुभएकोमा धन्यवाद।"


    ljspeech = True # Set True if the dataset format is ljspeech, and False if it's file-based.
    json_format = False # Set True if the dataset format is json, and False if it's file-based or ljspeech.
    preprocess = True # If you've already done preprocessing once, set it to false.
    
    is_turbo: bool = False # Set True if you're training Turbo, False if you're training Normal.
    is_lora: bool = True   # True: Efficient LoRA training (Recommended for < 10h data)
                           # False: Full Fine-Tune (High VRAM, for massive datasets)

    lora_r: int = 128
    lora_alpha: int = 256
    turbo_lora_target_modules: List[str] = field(default_factory=lambda: ["c_attn", "c_proj", "c_fc", "spkr_enc"])
    lora_target_modules: List[str] = field(default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "spkr_enc"])
    lora_modules_to_save: List[str] = field(default_factory=lambda: ["text_emb", "text_head"])
    

    # --- Vocabulary ---
    # The size of the NEW vocabulary (from tokenizer.json)
    # Ensure this matches the JSON file generated by your tokenizer script.
    # For Turbo mode: Use the exact number provided by setup.py (e.g., 52260)
    new_vocab_size: int = 52260 if is_turbo else 2454 

    # --- Hyperparameters ---
    batch_size: int = 32         # Adjust based on VRAM (2, 4, 8)
    grad_accum: int = 1        # Effective Batch Size = Batch * Accum
    learning_rate: float = 1e-4 if is_lora else 1e-5  # T3 is sensitive, keep low
    num_epochs: int = 10 if is_lora else 30

    save_steps: int = 500
    save_total_limit: int = 5
    dataloader_num_workers: int = 8

    # --- Constraints ---
    start_text_token = 255
    stop_text_token = 0
    max_text_len: int = 256
    max_speech_len: int = 850   # Truncates very long audio
    prompt_duration: float = 3.0 # Duration for the reference prompt (seconds)