#!/usr/bin/env python3
"""
Dataset Generator - Fully Async Pipeline
Q&A -> TTS (2 GPUs) -> Features (2 GPUs) -> Save
"""
import os
import sys
import re
import time
import gc
import logging
import multiprocessing as mp
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
import numpy as np
import requests
import torch

# Global logger
logger = None

def setup_logging(log_file=None):
    """Setup logging to both console and file."""
    global logger
    logger = logging.getLogger("dataset_generator")
    logger.setLevel(logging.INFO)
    logger.handlers.clear()

    # Console handler
    console = logging.StreamHandler(sys.stdout)
    console.setLevel(logging.INFO)
    console.setFormatter(logging.Formatter('%(message)s'))
    logger.addHandler(console)

    # File handler (if specified)
    if log_file:
        file_handler = logging.FileHandler(log_file, mode='a')
        file_handler.setLevel(logging.INFO)
        file_handler.setFormatter(logging.Formatter('%(asctime)s | %(message)s', datefmt='%H:%M:%S'))
        logger.addHandler(file_handler)

    return logger

def log(msg):
    """Log message to both console and file."""
    global logger
    if logger:
        logger.info(msg)
        # Force flush all handlers
        for handler in logger.handlers:
            handler.flush()
    else:
        print(msg)
        sys.stdout.flush()


def load_dotenv(env_path=None):
    """Load environment variables from .env file."""
    if env_path is None:
        # Look for .env in current dir or parent dirs
        for path in [Path(".env"), Path(__file__).parent.parent / ".env"]:
            if path.exists():
                env_path = path
                break
    if env_path and Path(env_path).exists():
        with open(env_path) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    os.environ.setdefault(key.strip(), value.strip())
        log(f"[ENV] Loaded from {env_path}")


# Load .env file
load_dotenv()

# Configuration
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
GROQ_MODEL = "openai/gpt-oss-20b"
GROQ_PARALLEL_REQUESTS = 10
QA_PER_REQUEST = 100


def get_num_gpus() -> int:
    """Auto-detect number of available GPUs."""
    if not torch.cuda.is_available():
        return 1
    return max(1, torch.cuda.device_count())


# Auto-detect GPUs at import time (can be overridden by --gpus)
NUM_GPUS = get_num_gpus()

# Memory estimates per component (GB per item in batch)
MEMORY_PER_ITEM = {
    "tts": 0.04,       # ~40MB per TTS item (question+answer pair)
    "whisper": 0.02,   # ~20MB per Whisper encoding (16kHz audio)
    "snac": 0.15,      # ~150MB per SNAC encoding (24kHz audio)
}

# Default batch sizes (will be auto-adjusted based on VRAM)
DEFAULT_BATCH_SIZES = {
    "tts": 200,        # TTS batch (processes questions + answers)
    "whisper": 8,      # Whisper parallel workers (each uses GPU)
    "snac": 50,        # SNAC encoding batch
}


def get_gpu_vram_gb() -> float:
    """Get total GPU VRAM in GB."""
    if not torch.cuda.is_available():
        return 0
    return torch.cuda.get_device_properties(0).total_memory / 1024**3


def calculate_batch_sizes(vram_gb: float = None, shared_gpu: bool = True) -> dict:
    """
    Calculate optimal batch sizes based on available VRAM.

    When shared_gpu=True (default), TTS and Features share the same GPU,
    so batch sizes must be more conservative to avoid OOM.

    VRAM tiers (with shared GPU adjustment):
    - 80GB+ (H100/H200): Large batches
    - 40-80GB (A100/H100-64GB): Medium-large batches
    - 24-40GB (RTX 4090): Medium batches
    - 16-24GB (RTX 3090/4080): Smaller batches
    - <16GB: Minimum safe values

    Returns dict with: tts, whisper (workers), snac batch sizes
    """
    if vram_gb is None:
        vram_gb = get_gpu_vram_gb()

    # Determine scale factor based on VRAM
    if vram_gb >= 80:
        tts_scale = 1.0
        whisper_workers = 8  # Parallel Whisper threads
        snac_scale = 0.6
    elif vram_gb >= 40:
        tts_scale = 0.75
        whisper_workers = 6
        snac_scale = 0.5
    elif vram_gb >= 24:
        tts_scale = 0.5
        whisper_workers = 4  # RTX 4090 - moderate parallelism
        snac_scale = 0.4
    elif vram_gb >= 16:
        tts_scale = 0.33
        whisper_workers = 2
        snac_scale = 0.25
    else:
        tts_scale = 0.2
        whisper_workers = 1
        snac_scale = 0.15

    # If not sharing GPU, can use more memory
    if not shared_gpu:
        snac_scale = min(1.0, snac_scale * 1.5)
        whisper_workers = min(8, whisper_workers + 2)

    # Calculate batch sizes
    batch_sizes = {
        "tts": max(10, int(DEFAULT_BATCH_SIZES["tts"] * tts_scale)),
        "whisper": whisper_workers,  # Thread count for parallel Whisper
        "snac": max(8, int(DEFAULT_BATCH_SIZES["snac"] * snac_scale)),
    }

    return batch_sizes


def print_batch_config(batch_sizes: dict, vram_gb: float):
    """Print batch configuration for transparency."""
    log(f"[Config] GPU VRAM: {vram_gb:.1f}GB")
    log(f"[Config] Batch sizes: TTS={batch_sizes['tts']}, Whisper={batch_sizes['whisper']}, SNAC={batch_sizes['snac']}")


def parse_qa(content):
    pairs = []
    content = content.replace("**", "")
    current_q, current_a = None, None
    for line in content.split("\n"):
        line = line.strip()
        if not line:
            continue
        qm = re.match(r"^[\d\.\)\-\*]*\s*[Qq][:\.]?\s*(.+)", line)
        am = re.match(r"^[\d\.\)\-\*]*\s*[Aa][:\.]?\s*(.+)", line)
        if qm:
            current_q = qm.group(1).strip()
        elif am:
            current_a = am.group(1).strip()
            if current_q and current_a and 2 <= len(current_q.split()) <= 25 and len(current_a) > 3:
                pairs.append({"q": current_q, "a": current_a})
            current_q, current_a = None, None
    return pairs


def make_groq_request(request_id: int) -> list:
    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    prompt = f"""Generate {QA_PER_REQUEST} unique Q&A pairs on diverse topics.
Format: Q: [question]
A: [answer]
Questions 2-25 words, answers 1-3 sentences."""

    for retry in range(3):
        try:
            response = requests.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers=headers,
                json={"model": GROQ_MODEL, "messages": [{"role": "user", "content": prompt}], "max_tokens": 8000, "temperature": 1.0},
                timeout=60
            )
            if response.status_code == 429:
                time.sleep(2 ** retry)
                continue
            response.raise_for_status()
            return parse_qa(response.json()["choices"][0]["message"]["content"])
        except:
            if retry < 2:
                time.sleep(1)
    return []


def qa_producer(target_count: int, tts_queue: mp.Queue, batch_size: int, status_queue: mp.Queue, num_workers: int = 2):
    """Produces Q&A batches for TTS - exactly target_count pairs."""
    seen = set()
    pairs = []
    pending = []
    batch_idx = 0
    t0 = time.time()

    while len(pairs) < target_count:
        with ThreadPoolExecutor(max_workers=GROQ_PARALLEL_REQUESTS) as ex:
            futures = [ex.submit(make_groq_request, i) for i in range(GROQ_PARALLEL_REQUESTS)]
            for f in as_completed(futures):
                for p in f.result():
                    if len(pairs) >= target_count:
                        break
                    norm = p["q"].lower().strip().rstrip("?")
                    if norm not in seen:
                        seen.add(norm)
                        pairs.append(p)
                        pending.append(p)
                        if len(pending) >= batch_size:
                            tts_queue.put((batch_idx, pending[:batch_size]))
                            batch_idx += 1
                            pending = pending[batch_size:]
                if len(pairs) >= target_count:
                    break

        status_queue.put(("qa", len(pairs), target_count, len(pairs)/(time.time()-t0)))

    if pending:
        tts_queue.put((batch_idx, pending))
        batch_idx += 1

    for _ in range(num_workers):
        tts_queue.put(None)

    status_queue.put(("qa_done", len(pairs), batch_idx))


def tts_worker(gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: dict, num_gpus: int = 1, gpu_offset: int = 0):
    """TTS worker - processes batches with VRAM-adjusted batch sizes on specific GPU."""
    import torch
    _orig_load = torch.load
    torch.load = lambda *a, **kw: _orig_load(*a, **{**kw, 'weights_only': False})

    # Set specific GPU for this worker (distribute across available GPUs with offset)
    actual_gpu = gpu_offset + (gpu_id % num_gpus)
    torch.cuda.set_device(actual_gpu)
    print(f"[TTS-GPU{gpu_id}] Assigned to CUDA device {actual_gpu}")

    # Patch JSON encoder to handle torch.dtype (transformers bug workaround)
    import json
    _orig_default = json.JSONEncoder.default
    def _patched_default(self, obj):
        if isinstance(obj, torch.dtype):
            return str(obj).split('.')[-1]  # torch.float32 -> "float32"
        return _orig_default(self, obj)
    json.JSONEncoder.default = _patched_default

    from soprano import SopranoTTS
    # Use lmdeploy backend for 2000x real-time speed (much faster than transformers)
    # Scale decoder_batch_size based on VRAM
    # Note: Soprano TTS only accepts "cuda", not "cuda:N"
    # torch.cuda.set_device() already selected the correct GPU above
    vram_gb = torch.cuda.get_device_properties(actual_gpu).total_memory / 1024**3
    gpu_name = torch.cuda.get_device_properties(actual_gpu).name

    # Check if GPU supports lmdeploy (Blackwell/RTX 50xx not supported yet)
    use_lmdeploy = "5090" not in gpu_name and "5080" not in gpu_name and "B100" not in gpu_name

    if use_lmdeploy:
        try:
            dec_batch = 32 if vram_gb >= 80 else (16 if vram_gb >= 40 else 8)
            tts = SopranoTTS(
                backend="lmdeploy",  # Fastest backend
                device="cuda",  # Uses current device set by torch.cuda.set_device()
                cache_size_mb=4000 if vram_gb >= 24 else 2000,  # More cache = faster
                decoder_batch_size=dec_batch,  # Parallel decoding based on VRAM
            )
            print(f"[TTS-GPU{gpu_id}] Using lmdeploy backend (decoder_batch={dec_batch})")
        except Exception as e:
            print(f"[TTS-GPU{gpu_id}] lmdeploy failed ({e}), falling back to transformers")
            tts = SopranoTTS(backend="transformers", device="cuda")
    else:
        print(f"[TTS-GPU{gpu_id}] Blackwell GPU detected ({gpu_name}), using transformers backend")
        tts = SopranoTTS(backend="transformers", device="cuda")

    # Use centralized batch size
    tts_batch = batch_sizes.get("tts", 200)
    print(f"[TTS-GPU{gpu_id}] Using batch size: {tts_batch}")

    status_queue.put(("tts_ready", gpu_id))

    processed = 0
    t_start = time.time()
    last_heartbeat = time.time()

    while True:
        # Send heartbeat every 30 seconds to show worker is alive
        if time.time() - last_heartbeat > 30:
            status_queue.put(("tts_heartbeat", gpu_id, processed))
            last_heartbeat = time.time()

        # Use timeout to allow heartbeat even when queue is empty
        try:
            item = tts_queue.get(timeout=5)
        except:
            continue  # Timeout, send heartbeat and retry

        if item is None:
            break

        batch_idx, pairs = item

        # Process all pairs in one batch (batch=200 for max speed)
        all_results = []
        questions = [p["q"] for p in pairs]
        answers = [p["a"] for p in pairs]

        try:
            # Generate all audio in one batch call
            combined = questions + answers
            all_audios = tts.infer_batch(combined)
            q_audios = all_audios[:len(questions)]
            a_audios = all_audios[len(questions):]

            for j, p in enumerate(pairs):
                q_np = q_audios[j].cpu().numpy() if hasattr(q_audios[j], 'numpy') else q_audios[j]
                a_np = a_audios[j].cpu().numpy() if hasattr(a_audios[j], 'numpy') else a_audios[j]
                all_results.append({
                    "question": p["q"], "answer": p["a"],
                    "q_audio": np.asarray(q_np, dtype=np.float32),
                    "a_audio": np.asarray(a_np, dtype=np.float32),
                })

        except Exception as e:
            import traceback
            status_queue.put(("tts_error", gpu_id, f"{str(e)}\n{traceback.format_exc()}"))
            # Continue processing other batches instead of crashing
            continue

        if all_results:
            feat_queue.put((batch_idx, all_results))
            processed += len(all_results)
            elapsed = time.time() - t_start
            status_queue.put(("tts", gpu_id, processed, processed/elapsed, len(all_results)/(time.time()-t_start) if elapsed > 0 else 0))

    # TTS worker sends None to feat_queue directly (ensures all batches are sent before close signal)
    feat_queue.put(None)
    status_queue.put(("tts_done", gpu_id, processed))


def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: dict, num_gpus: int = 1, gpu_offset: int = 0):
    """Features worker - VRAM-adjusted batch sizes for Whisper + SNAC + NeMo NFA on specific GPU."""
    import torch
    _orig = torch.load
    torch.load = lambda *a, **kw: _orig(*a, **{**kw, 'weights_only': False})

    import torchaudio
    import snac
    import tempfile
    import soundfile as sf
    from transformers import AutoTokenizer, WhisperModel, WhisperFeatureExtractor
    from huggingface_hub import login

    # Set specific GPU for this worker (distribute across available GPUs with offset)
    actual_gpu = gpu_offset + (gpu_id % num_gpus)
    torch.cuda.set_device(actual_gpu)
    device = f"cuda:{actual_gpu}"
    print(f"[Features-GPU{gpu_id}] Assigned to CUDA device {actual_gpu}")

    # HuggingFace auth for gated models
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        try:
            login(token=hf_token)
        except Exception as e:
            print(f"[WARN] HuggingFace login failed: {e}")

    # Use Whisper Large V3 Turbo (50% faster, same 1280-dim output)
    print("[Features] Loading Whisper large-v3-turbo (transformers)...")
    whisper_model = WhisperModel.from_pretrained("openai/whisper-large-v3-turbo", torch_dtype=torch.float16).to(device).eval()
    whisper_feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-large-v3-turbo")

    # Note: torch.compile is skipped for Whisper as it causes issues with conv1d
    # The batched processing already provides significant speedup
    print("[Features] Whisper Turbo loaded successfully")

    snac_model = snac.SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()

    # Load NeMo ASR model for forced alignment (GPU-accelerated)
    nfa_model = None
    try:
        import nemo.collections.asr as nemo_asr
        # Use FastConformer CTC for fast GPU-based alignment
        nfa_model = nemo_asr.models.EncDecCTCModel.from_pretrained("nvidia/stt_en_fastconformer_ctc_large")
        nfa_model = nfa_model.to(device).eval()
        print("[Features] NeMo Forced Aligner loaded (GPU-accelerated)")
    except Exception as e:
        print(f"[WARN] NeMo NFA not available, using proportional alignment: {e}")

    # Load tokenizer for pre-computing text tokens (Orpheus model)
    tokenizer = None
    for model_path in ["canopylabs/orpheus-3b-0.1-pretrained", "meta-llama/Llama-3.2-3B"]:
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_path, token=hf_token)
            print(f"[Features] Loaded tokenizer from {model_path}")
            break
        except Exception as e:
            print(f"[WARN] Failed to load tokenizer from {model_path}: {e}")

    if tokenizer is None:
        print("[WARN] No tokenizer available - text_tokens will not be pre-computed")

    # Use centralized batch sizes
    snac_batch = batch_sizes.get("snac", 50)
    whisper_workers = batch_sizes.get("whisper", 8)
    print(f"[Features-GPU{gpu_id}] Using SNAC batch={snac_batch}, Whisper workers={whisper_workers}")

    status_queue.put(("feat_ready", gpu_id))

    def process_whisper(audio_data):
        """Process single audio with Whisper. GPU-accelerated resampling."""
        # GPU-accelerated resampling
        audio_tensor = torch.from_numpy(audio_data).to(device)
        audio_16k = torchaudio.functional.resample(audio_tensor, 32000, 16000)

        # Truncate to max 30 seconds (480000 samples at 16kHz) - Whisper limit
        max_samples = 480000
        if audio_16k.shape[0] > max_samples:
            audio_16k = audio_16k[:max_samples]

        audio_16k_np = audio_16k.cpu().numpy().astype(np.float32)

        # Extract features using Whisper feature extractor
        inputs = whisper_feature_extractor(audio_16k_np, sampling_rate=16000, return_tensors="pt")
        input_features = inputs.input_features.to(device, dtype=torch.float16)

        # Encode with Whisper
        with torch.no_grad():
            encoder_outputs = whisper_model.encoder(input_features)

        # Return encoder hidden states [seq_len, 1280]
        return encoder_outputs.last_hidden_state.squeeze(0).cpu().float()

    def extract_snac_tokens(codes, idx):
        """Extract SNAC tokens with correct position-based offsets.

        Each frame has 7 tokens (1 from codebook 0, 2 from codebook 1, 4 from codebook 2).
        Each position within a frame needs a unique offset:
        - Position 0: 128266 + 0*4096
        - Position 1: 128266 + 1*4096
        - Position 2: 128266 + 2*4096
        - Position 3: 128266 + 3*4096
        - Position 4: 128266 + 4*4096
        - Position 5: 128266 + 5*4096
        - Position 6: 128266 + 6*4096
        """
        SNAC_BASE = 128266
        tokens = []
        for j in range(codes[0].shape[-1]):
            # Position 0: codebook 0 token
            tokens.append(codes[0][idx, j].item() + SNAC_BASE + 0 * 4096)

            # Positions 1-2: codebook 1 tokens
            if j * 2 + 1 < codes[1].shape[-1]:
                tokens.append(codes[1][idx, j * 2].item() + SNAC_BASE + 1 * 4096)
                tokens.append(codes[1][idx, j * 2 + 1].item() + SNAC_BASE + 2 * 4096)

            # Positions 3-6: codebook 2 tokens
            for k in range(4):
                if j * 4 + k < codes[2].shape[-1]:
                    tokens.append(codes[2][idx, j * 4 + k].item() + SNAC_BASE + (3 + k) * 4096)
        return tokens

    def get_word_alignments_nfa(audio_data, text, sample_rate=32000):
        """Get word-level alignments using NeMo Forced Aligner (GPU-accelerated)."""
        if nfa_model is None:
            return None

        try:
            # Resample to 16kHz for NeMo ASR
            audio_16k = torchaudio.functional.resample(
                torch.from_numpy(audio_data), sample_rate, 16000
            ).numpy().astype(np.float32)

            # Save to temp file (NeMo requires file path)
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                temp_path = f.name
                sf.write(temp_path, audio_16k, 16000)

            # Run CTC alignment using NeMo
            # This uses GPU-accelerated Viterbi decoding
            with torch.no_grad():
                # Transcribe to get log probs
                hypotheses = nfa_model.transcribe(
                    [temp_path],
                    return_hypotheses=True,
                    batch_size=1
                )

            # Clean up temp file
            os.unlink(temp_path)

            # Extract word timestamps from hypothesis
            if hypotheses and len(hypotheses) > 0:
                hyp = hypotheses[0]
                if hasattr(hyp, 'timestep') and hyp.timestep is not None:
                    # Use NeMo's word-level timestamps
                    word_alignments = []
                    words = text.split()
                    timestamps = hyp.timestep.get('word', [])

                    # Map recognized words to reference text
                    for i, word in enumerate(words):
                        if i < len(timestamps):
                            ts = timestamps[i]
                            start_time = ts.get('start', 0)
                            end_time = ts.get('end', start_time + 0.1)
                        else:
                            # Fallback: estimate from audio duration
                            audio_duration = len(audio_data) / sample_rate
                            word_start = (i / len(words)) * audio_duration
                            word_end = ((i + 1) / len(words)) * audio_duration
                            start_time, end_time = word_start, word_end

                        # Convert to SNAC frames (75 frames/second)
                        start_frame = int(start_time * 75)
                        end_frame = int(end_time * 75)

                        word_tokens = []
                        if tokenizer is not None:
                            word_tokens = tokenizer.encode(word, add_special_tokens=False)

                        word_alignments.append({
                            'word': word,
                            'start': start_time,
                            'end': end_time,
                            'start_frame': start_frame,
                            'end_frame': end_frame,
                            'tokens': word_tokens
                        })
                    return word_alignments

        except Exception as e:
            status_queue.put(("nfa_error", gpu_id, str(e)))

        return None

    def get_word_alignments_proportional(audio_data, text, sample_rate=32000):
        """Fallback: proportional word alignment based on character count.

        Calculates frame indices that match actual SNAC output:
        - SNAC operates at 24kHz with ~320 samples per frame (75 fps)
        - Audio is resampled from sample_rate to 24kHz before SNAC
        """
        words = text.split()
        if not words:
            return []

        # Calculate actual SNAC frame count after resampling to 24kHz
        # SNAC uses ~320 samples per frame at 24kHz
        audio_24k_samples = len(audio_data) * 24000 / sample_rate
        total_snac_frames = int(audio_24k_samples / 320)

        if total_snac_frames == 0:
            return []

        total_chars = sum(len(w) for w in words)
        if total_chars == 0:
            return []

        audio_duration = len(audio_data) / sample_rate
        word_alignments = []
        current_frame = 0

        # Pre-tokenize all words in batch for efficiency
        all_word_tokens = []
        if tokenizer is not None:
            for word in words:
                all_word_tokens.append(tokenizer.encode(word, add_special_tokens=False))
        else:
            all_word_tokens = [[] for _ in words]

        for i, word in enumerate(words):
            # Distribute frames proportionally based on character count
            word_frames = int((len(word) / total_chars) * total_snac_frames)

            # Ensure last word gets remaining frames
            if i == len(words) - 1:
                end_frame = total_snac_frames
            else:
                end_frame = min(current_frame + max(1, word_frames), total_snac_frames)

            start_frame = current_frame

            # Calculate time from frames (for compatibility)
            start_time = start_frame / 75.0
            end_time = end_frame / 75.0

            word_alignments.append({
                'word': word,
                'start': start_time,
                'end': end_time,
                'start_frame': start_frame,
                'end_frame': end_frame,
                'tokens': all_word_tokens[i]
            })
            current_frame = end_frame

        return word_alignments

    processed = 0
    t_start = time.time()
    last_heartbeat = time.time()

    while True:
        # Send heartbeat every 30 seconds to show worker is alive
        if time.time() - last_heartbeat > 30:
            status_queue.put(("feat_heartbeat", gpu_id, processed))
            last_heartbeat = time.time()

        # Use timeout to allow heartbeat even when queue is empty
        try:
            item = feat_queue.get(timeout=5)
        except:
            continue  # Timeout, send heartbeat and retry

        if item is None:
            break

        batch_idx, audio_batch = item
        t0 = time.time()

        try:
            # 1. Parallel Whisper encoding with GPU-accelerated resampling
            q_audios = [ad["q_audio"] for ad in audio_batch]
            with ThreadPoolExecutor(max_workers=whisper_workers) as ex:
                whisper_features = list(ex.map(process_whisper, q_audios))

            # 2. SNAC encoding - GPU-batched with GPU resampling
            a_audios = [ad["a_audio"] for ad in audio_batch]
            all_tokens = []

            for start in range(0, len(a_audios), snac_batch):
                end = min(start + snac_batch, len(a_audios))
                mini_audios = a_audios[start:end]

                max_len = max(a.shape[0] for a in mini_audios)
                padded = [np.pad(a, (0, max_len - len(a))) for a in mini_audios]

                # GPU-accelerated resampling for SNAC
                audios_24k = torch.stack([
                    torchaudio.functional.resample(
                        torch.from_numpy(a).to(device), 32000, 24000
                    ).cpu()
                    for a in padded
                ])

                with torch.no_grad():
                    codes = snac_model.encode(audios_24k.unsqueeze(1).to(device))

                for i in range(len(mini_audios)):
                    all_tokens.append(extract_snac_tokens(codes, i))

            torch.cuda.synchronize()

            # Periodic GPU memory cleanup every 100 batches to prevent fragmentation
            if processed > 0 and processed % (100 * len(audio_batch)) == 0:
                torch.cuda.empty_cache()
                gc.collect()

            # 3. Build results with pre-computed text tokens and word alignments
            results = []
            for i, ad in enumerate(audio_batch):
                answer_text = ad["answer"]
                answer_audio = ad["a_audio"]
                result = {
                    "whisper_features": whisper_features[i],
                    "snac_tokens": torch.tensor(all_tokens[i], dtype=torch.long),
                    "text": ad["question"],
                    "answer": answer_text
                }

                # Pre-tokenize answer text for training (if tokenizer available)
                if tokenizer is not None:
                    text_tokens = tokenizer.encode(answer_text, add_special_tokens=False)
                    result["text_tokens"] = torch.tensor(text_tokens, dtype=torch.long)

                # Generate word alignments for IST-LM interleaving
                # Use proportional alignment (NFA disabled for now - causes hangs)
                word_alignments = get_word_alignments_proportional(answer_audio, answer_text)
                if word_alignments:
                    result["word_alignments"] = word_alignments

                results.append(result)

            batch_time = time.time() - t0
            batch_rate = len(results) / batch_time if batch_time > 0 else 0

            # Put results with timeout to prevent indefinite blocking
            put_start = time.time()
            max_put_attempts = 10
            for attempt in range(max_put_attempts):
                try:
                    result_queue.put((batch_idx, results), timeout=30)
                    break
                except Exception as put_err:
                    if attempt < max_put_attempts - 1:
                        status_queue.put(("feat_warn", gpu_id, f"Queue full, retry {attempt+1}"))
                        time.sleep(1)
                    else:
                        status_queue.put(("feat_error", gpu_id, f"Failed to put results after {max_put_attempts} attempts"))
                        raise put_err

            processed += len(results)
            elapsed = time.time() - t_start
            status_queue.put(("feat", gpu_id, processed, processed/elapsed, batch_rate))

            # Clear intermediate tensors to prevent memory accumulation
            del whisper_features, all_tokens, results
            if processed % 500 == 0:  # More aggressive cleanup every 500 items
                torch.cuda.empty_cache()
                gc.collect()

        except Exception as e:
            import traceback
            error_msg = str(e)
            status_queue.put(("feat_error", gpu_id, f"{error_msg}\n{traceback.format_exc()}"))

            # Clear GPU memory aggressively
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            gc.collect()

            # If OOM, try to recover by reducing batch sizes
            if "out of memory" in error_msg.lower() or "OOM" in error_msg:
                status_queue.put(("feat_warn", gpu_id, "OOM detected, clearing memory..."))
                time.sleep(2)  # Give GPU time to recover
                torch.cuda.empty_cache()

            continue

    status_queue.put(("feat_done", gpu_id, processed))


def main():
    mp.set_start_method('spawn', force=True)

    # Increase file descriptor limit
    import resource
    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
    resource.setrlimit(resource.RLIMIT_NOFILE, (min(65536, hard), hard))

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--count", "--num_samples", type=int, default=100, dest="count")
    parser.add_argument("--output", type=str, default="./data/dataset.pt")
    parser.add_argument("--gpus", type=int, default=NUM_GPUS)
    parser.add_argument("--resume", action="store_true", help="Resume from existing checkpoint")
    parser.add_argument("--checkpoint-interval", type=int, default=1000, help="Save checkpoint every N items")
    parser.add_argument("--log-file", type=str, default=None, help="Log file path (default: output.log)")
    # Optional overrides for batch sizes (if not set, auto-calculated from VRAM)
    parser.add_argument("--tts-batch", type=int, default=None)
    parser.add_argument("--snac-batch", type=int, default=None)
    parser.add_argument("--whisper-workers", type=int, default=None)
    args = parser.parse_args()

    Path(args.output).parent.mkdir(parents=True, exist_ok=True)

    # Setup logging (default to output path + .log)
    log_file = args.log_file or (args.output + ".log")
    setup_logging(log_file)
    log(f"[Log] Writing to {log_file}")

    # Calculate batch sizes based on GPU VRAM
    vram_gb = get_gpu_vram_gb()
    batch_sizes = calculate_batch_sizes(vram_gb)

    # Apply command-line overrides if specified
    if args.tts_batch is not None:
        batch_sizes["tts"] = args.tts_batch
    if args.snac_batch is not None:
        batch_sizes["snac"] = args.snac_batch
    if args.whisper_workers is not None:
        batch_sizes["whisper"] = args.whisper_workers

    log("=" * 60)
    log("Dataset Generator - Fully Async Pipeline")
    log(f"Target: {args.count} items, GPUs: {args.gpus}")
    print_batch_config(batch_sizes, vram_gb)
    log("=" * 60)

    # Resume from checkpoint or output file if exists
    checkpoint_path = args.output + ".checkpoint"
    existing_items = []
    start_count = 0

    # Track resume state
    resume_from_path = None

    if args.resume:
        # Check for existing data to resume from
        # Priority: .new checkpoint (partial new items) + output file, or just output file

        new_checkpoint = checkpoint_path + ".new"
        new_items_count = 0
        base_count = 0

        # Check if we have new items checkpoint
        if Path(new_checkpoint).exists():
            try:
                data = torch.load(new_checkpoint, map_location="cpu", weights_only=False, mmap=True)
                new_items_count = len(data)
                del data
                log(f"[Resume] Found {new_items_count} new items in checkpoint")
            except Exception as e:
                log(f"[Resume] Failed to read {new_checkpoint}: {e}")

        # Check output file for base items
        if Path(args.output).exists():
            try:
                data = torch.load(args.output, map_location="cpu", weights_only=False, mmap=True)
                base_count = len(data)
                del data
                resume_from_path = args.output
                log(f"[Resume] Found {base_count} base items in {args.output}")
            except Exception as e:
                log(f"[Resume] Failed to read {args.output}: {e}")

        # Total count is base + new
        start_count = base_count + new_items_count
        if start_count > 0:
            log(f"[Resume] Total: {start_count} items ({base_count} base + {new_items_count} new), need {args.count - start_count} more")
        else:
            log("[Resume] No valid resume file found, starting fresh")

    if start_count >= args.count:
        log(f"[Resume] Already have {start_count} items, saving final...")
        torch.save(existing_items[:args.count], args.output)
        log(f"COMPLETE: {args.count} items saved to {args.output}")
        return

    remaining_count = args.count - start_count
    log(f"[Main] Generating {remaining_count} new items...")

    total_start = time.time()

    # Use regular mp.Queue with NO size limits to prevent deadlocks
    # Workers will block on put() if queue is full, causing stalls
    # Memory is managed by batch sizes instead
    tts_queue = mp.Queue()  # No maxsize - prevents TTS blocking
    feat_queue = mp.Queue()  # No maxsize - prevents Features blocking
    result_queue = mp.Queue()  # No maxsize - prevents result collection blocking
    status_queue = mp.Queue()

    workers = []

    # Get actual GPU count for worker assignment
    actual_num_gpus = get_num_gpus() if torch.cuda.is_available() else 1

    # With 4+ GPUs, separate TTS and Features to avoid OOM
    # TTS uses GPUs 0 to (N/2-1), Features uses GPUs (N/2) to (N-1)
    if actual_num_gpus >= 4:
        tts_gpus = actual_num_gpus // 2  # First half for TTS
        feat_gpus = actual_num_gpus - tts_gpus  # Second half for Features
        tts_gpu_offset = 0
        feat_gpu_offset = tts_gpus
        log(f"\n[Main] Detected {actual_num_gpus} GPUs - Separating: TTS on GPUs 0-{tts_gpus-1}, Features on GPUs {feat_gpu_offset}-{actual_num_gpus-1}")
    else:
        tts_gpus = actual_num_gpus
        feat_gpus = actual_num_gpus
        tts_gpu_offset = 0
        feat_gpu_offset = 0
        log(f"\n[Main] Detected {actual_num_gpus} GPUs, sharing between TTS and Features")

    # Adjust number of workers based on available GPUs
    num_tts_workers = min(args.gpus, tts_gpus)
    num_feat_workers = min(args.gpus, feat_gpus)
    log(f"[Main] Spawning {num_tts_workers} TTS workers, {num_feat_workers} Features workers")

    qa_proc = mp.Process(target=qa_producer, args=(remaining_count, tts_queue, batch_sizes["tts"], status_queue, num_tts_workers))
    qa_proc.start()
    workers.append(qa_proc)

    for gpu_id in range(num_tts_workers):
        p = mp.Process(target=tts_worker, args=(gpu_id, tts_queue, feat_queue, status_queue, batch_sizes, tts_gpus, tts_gpu_offset))
        p.start()
        workers.append(p)

    for gpu_id in range(num_feat_workers):
        p = mp.Process(target=features_worker, args=(gpu_id, feat_queue, result_queue, status_queue, batch_sizes, feat_gpus, feat_gpu_offset))
        p.start()
        workers.append(p)

    log("[Pipeline] All workers started, monitoring...")

    results = {}
    tts_done_count = 0
    feat_done_count = 0
    total_items = 0
    t0 = time.time()
    tts_ready = 0
    feat_ready = 0
    expected_from_feat = {}
    feat_queue_closed = False
    last_checkpoint_count = 0

    # Main loop with improved stall detection
    last_result_time = time.time()
    last_status_time = time.time()
    last_heartbeat_time = {f"tts_{i}": time.time() for i in range(num_tts_workers)}
    last_heartbeat_time.update({f"feat_{i}": time.time() for i in range(num_feat_workers)})
    stall_warning_shown = False
    errors = []

    while True:
        # Process ALL pending status messages
        for _ in range(100):  # Limit to prevent infinite loop
            try:
                msg = status_queue.get_nowait()
                last_status_time = time.time()
                msg_type = msg[0]

                if msg_type == "tts_ready":
                    tts_ready += 1
                    log(f"[TTS-GPU{msg[1]}] Ready ({tts_ready}/{num_tts_workers})")
                elif msg_type == "feat_ready":
                    feat_ready += 1
                    log(f"[Features-GPU{msg[1]}] Ready ({feat_ready}/{num_feat_workers})")
                elif msg_type == "qa":
                    log(f"[Q&A] {msg[1]}/{msg[2]} | {msg[3]:.1f}/s")
                elif msg_type == "qa_done":
                    log(f"[Q&A] Done: {msg[1]} pairs, {msg[2]} batches")
                elif msg_type == "tts":
                    log(f"[TTS-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s | batch {msg[4]:.1f}/s")
                    last_heartbeat_time[f"tts_{msg[1]}"] = time.time()
                elif msg_type == "tts_done":
                    tts_done_count += 1
                    log(f"[TTS-GPU{msg[1]}] Done: {msg[2]} items")
                elif msg_type == "feat":
                    log(f"[Feat-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s | batch {msg[4]:.1f}/s")
                    last_heartbeat_time[f"feat_{msg[1]}"] = time.time()
                elif msg_type == "feat_done":
                    feat_done_count += 1
                    expected_from_feat[msg[1]] = msg[2]
                    log(f"[Features-GPU{msg[1]}] Done: {msg[2]} items")
                elif msg_type == "tts_heartbeat":
                    last_heartbeat_time[f"tts_{msg[1]}"] = time.time()
                elif msg_type == "feat_heartbeat":
                    last_heartbeat_time[f"feat_{msg[1]}"] = time.time()
                elif msg_type == "feat_warn":
                    log(f"[WARN] Features-GPU{msg[1]}: {msg[2]}")
                elif "error" in msg_type:
                    log(f"[Error] {msg}")
                    errors.append(msg)
            except:
                break

        if tts_done_count >= num_tts_workers and not feat_queue_closed:
            log(f"[Main] TTS done, waiting for features to finish...")
            feat_queue_closed = True

        # Collect ALL available results (drain aggressively with non-blocking gets)
        collected_this_round = 0
        drain_start = time.time()
        while time.time() - drain_start < 2.0:  # Spend up to 2s draining
            try:
                batch_idx, items = result_queue.get_nowait()
                results[batch_idx] = items
                total_items += len(items)
                collected_this_round += len(items)
                last_result_time = time.time()
                stall_warning_shown = False
            except:
                # No more items immediately available, wait briefly then check again
                time.sleep(0.05)
                try:
                    batch_idx, items = result_queue.get_nowait()
                    results[batch_idx] = items
                    total_items += len(items)
                    collected_this_round += len(items)
                    last_result_time = time.time()
                    stall_warning_shown = False
                except:
                    break  # Queue truly empty

        if collected_this_round > 0:
            elapsed = time.time() - t0
            log(f"[Results] {total_items}/{remaining_count} | {total_items/elapsed:.1f}/s")

        # Save checkpoint periodically (combine with any existing checkpoint)
        if total_items - last_checkpoint_count >= args.checkpoint_interval:
            log(f"[Checkpoint] Saving {total_items} new items...")
            # Collect items from this run
            items_this_run = []
            for i in sorted(results.keys()):
                items_this_run.extend(results[i])

            # Load and combine with existing checkpoint if present
            checkpoint_new_path = checkpoint_path + ".new"
            all_checkpoint_items = []
            if Path(checkpoint_new_path).exists():
                try:
                    prev_items = torch.load(checkpoint_new_path, map_location="cpu", weights_only=False)
                    all_checkpoint_items = list(prev_items)
                    del prev_items
                except:
                    pass

            # Only add items not already in checkpoint
            items_to_add = items_this_run[len(all_checkpoint_items):]
            all_checkpoint_items.extend(items_to_add)

            torch.save(all_checkpoint_items, checkpoint_new_path)
            last_checkpoint_count = total_items
            log(f"[Checkpoint] Saved {len(all_checkpoint_items)} total to {checkpoint_new_path}")

        # Check exit conditions
        if total_items >= remaining_count:
            log(f"[Main] Target reached: {total_items}/{remaining_count}")
            break

        # If all feature workers done, drain remaining results
        if feat_done_count >= num_feat_workers:
            log(f"[Main] All workers done, draining queue...")
            drain_start = time.time()
            while time.time() - drain_start < 60.0:
                try:
                    batch_idx, items = result_queue.get(timeout=0.5)
                    results[batch_idx] = items
                    total_items += len(items)
                    log(f"[Results] {total_items}/{remaining_count} | (drained)")
                except:
                    time.sleep(0.2)
                    try:
                        batch_idx, items = result_queue.get_nowait()
                        results[batch_idx] = items
                        total_items += len(items)
                        log(f"[Results] {total_items}/{remaining_count} | (drained)")
                    except:
                        break
            break

        # Check if workers are still alive and responding
        alive_workers = sum(1 for p in workers if p.is_alive())
        time_since_result = time.time() - last_result_time
        time_since_status = time.time() - last_status_time

        # Check heartbeats - detect stuck workers
        stuck_workers = []
        for worker_id, last_hb in last_heartbeat_time.items():
            if time.time() - last_hb > 120:  # No heartbeat for 2 minutes
                stuck_workers.append(worker_id)

        if time_since_result > 30 and not stall_warning_shown:
            queue_size = 0
            try:
                queue_size = result_queue.qsize()
            except:
                pass
            log(f"[WARN] No results for 30s, {alive_workers} workers alive, queue ~{queue_size} items")
            if stuck_workers:
                log(f"[WARN] Stuck workers (no heartbeat >2min): {stuck_workers}")
            stall_warning_shown = True

        # Save checkpoint on stall detection (combine with existing checkpoint)
        if time_since_result > 60 and total_items > last_checkpoint_count:
            log(f"[Checkpoint] Stall detected, saving {total_items} items...")
            items_this_run = []
            for i in sorted(results.keys()):
                items_this_run.extend(results[i])

            # Load and combine with existing checkpoint
            checkpoint_new_path = checkpoint_path + ".new"
            all_checkpoint_items = []
            if Path(checkpoint_new_path).exists():
                try:
                    prev_items = torch.load(checkpoint_new_path, map_location="cpu", weights_only=False)
                    all_checkpoint_items = list(prev_items)
                    del prev_items
                except:
                    pass

            items_to_add = items_this_run[len(all_checkpoint_items):]
            all_checkpoint_items.extend(items_to_add)
            torch.save(all_checkpoint_items, checkpoint_new_path)
            last_checkpoint_count = total_items

        # Exit if truly stalled - but be smarter about it
        # Only exit if: no results for 3min AND no heartbeats AND workers dead
        all_workers_stuck = len(stuck_workers) >= (num_tts_workers + num_feat_workers)
        if time_since_result > 180 and all_workers_stuck:
            log(f"[WARN] All workers stuck, stopping with {total_items} items")
            log(f"[WARN] Errors encountered: {len(errors)}")
            for err in errors[-5:]:  # Show last 5 errors
                log(f"  {err}")
            break

        if time_since_result > 600:
            log(f"[WARN] No progress for 10min, stopping with {total_items} items")
            break

    # Save checkpoint before cleanup (combine with existing checkpoint)
    if total_items > last_checkpoint_count:
        log(f"[Checkpoint] Final save before cleanup: {total_items} items...")
        items_this_run = []
        for i in sorted(results.keys()):
            items_this_run.extend(results[i])

        # Load and combine with existing checkpoint
        checkpoint_new_path = checkpoint_path + ".new"
        all_checkpoint_items = []
        if Path(checkpoint_new_path).exists():
            try:
                prev_items = torch.load(checkpoint_new_path, map_location="cpu", weights_only=False)
                all_checkpoint_items = list(prev_items)
                del prev_items
            except:
                pass

        items_to_add = items_this_run[len(all_checkpoint_items):]
        all_checkpoint_items.extend(items_to_add)
        torch.save(all_checkpoint_items, checkpoint_new_path)
        last_checkpoint_count = total_items

    # Wait for workers to finish
    log("[Main] Waiting for workers to finish...")
    for p in workers:
        p.join(timeout=5)
        if p.is_alive():
            p.terminate()

    # Final drain
    for _ in range(100):
        try:
            batch_idx, items = result_queue.get_nowait()
            results[batch_idx] = items
            total_items += len(items)
        except:
            break

    # Collect new results from this run
    new_items_this_run = []
    for i in sorted(results.keys()):
        new_items_this_run.extend(results[i])

    log(f"[Main] Generated {len(new_items_this_run)} new items this run")

    # Load any previously checkpointed new items and combine with this run's items
    new_checkpoint = checkpoint_path + ".new"
    new_items = []

    if Path(new_checkpoint).exists():
        try:
            prev_new = torch.load(new_checkpoint, map_location="cpu", weights_only=False)
            new_items = list(prev_new)
            log(f"[Main] Loaded {len(new_items)} items from previous checkpoint")
            del prev_new
        except Exception as e:
            log(f"[Main] Failed to load previous checkpoint: {e}")

    # Add this run's items to the checkpoint items
    new_items.extend(new_items_this_run)
    log(f"[Main] Total new items: {len(new_items)} (checkpoint: {len(new_items) - len(new_items_this_run)}, this run: {len(new_items_this_run)})")

    total_new = len(new_items)

    if total_new == 0 and (not resume_from_path or start_count == 0):
        log("[ERROR] No items collected! Check worker logs.")
        sys.exit(1)

    # Combine with base items if resuming
    if resume_from_path and Path(resume_from_path).exists():
        log(f"[Main] Loading base items from {resume_from_path}...")
        base_data = torch.load(resume_from_path, map_location="cpu", weights_only=False, mmap=True)
        base_count = len(base_data)

        # Calculate how many base items to keep
        items_needed_from_base = min(base_count, args.count - total_new)
        log(f"[Main] Combining {items_needed_from_base} base + {total_new} new items...")

        final_items = list(base_data[:items_needed_from_base]) + new_items
        del base_data
    else:
        final_items = new_items

    # Trim to target count
    final_items = final_items[:args.count]

    # Save final dataset
    log(f"[Main] Saving {len(final_items)} items to {args.output}...")
    torch.save(final_items, args.output)

    # Remove checkpoint files if complete
    if len(final_items) >= args.count:
        for cp in [checkpoint_path, checkpoint_path + ".new"]:
            if Path(cp).exists():
                Path(cp).unlink()
                log(f"[Cleanup] Removed {cp}")

    total_time = time.time() - total_start
    log("\n" + "=" * 60)
    log(f"COMPLETE: {len(final_items)} items saved to {args.output}")
    if start_count > 0:
        log(f"  (resumed from {start_count}, added {len(final_items) - start_count} new)")
    log(f"Total time: {total_time:.1f}s ({total_time/60:.1f}m)")
    log(f"Throughput: {(len(final_items) - start_count)/total_time:.2f} items/s")
    log("=" * 60)


if __name__ == "__main__":
    main()