#!/usr/bin/env python3
"""
Create dataset from Hindi TTS data with phonemes.
Format: filename.wav|phoneme|language_id|speaker_id|text
"""

import json
import shutil
import multiprocessing as mp
from pathlib import Path
from tqdm import tqdm
import logging
try:
    from phonemizer import phonemize
except ImportError:
    print("Error: 'phonemizer' package is not installed. Please install it using 'pip install phonemizer' and ensure 'espeak-ng' is installed on your system.")
    exit(1)

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
log = logging.getLogger(__name__)
logging.getLogger("phonemizer").setLevel(logging.ERROR)

BASE_DIR = Path("IISc_SYSPIN_Data")
MALE_DIR = BASE_DIR / "IISc_SYSPINProject_Hindi_Male_Spk001_HC"
FEMALE_DIR = BASE_DIR / "IISc_SYSPINProject_Hindi_Female_Spk001_HC"
OUTPUT_WAV_DIR = Path("hindi_wavs")
WAV_REL_PATH = "hindi/hindi_wavs"
OUTPUT_FILE = Path("dataset.txt")
EXISTING_TRAIN_LIST = Path("..") / "list.txt"
MERGED_TRAIN_LIST = Path("..") / "list_merged.txt"
# NUM_WORKERS = 25

# IDs for dataset format
LANGUAGE_ID = 6
SPEAKER_ID_MAP = {
    "male": 9,
    "female": 10,
}


def get_phonemes(text):
    """Convert text to phonemes using phonemizer with espeak-ng."""
    try:
        phonemes_list = phonemize(
            [text],
            language="hi",
            backend="espeak",
            strip=True,
            preserve_punctuation=True,
            with_stress=True,
        )
        if phonemes_list:
            return " ".join(phonemes_list[0].split())
        return ""
    except Exception:
        return ""


def process_entry(entry):
    """Process single entry: (key, text, speaker_id) -> formatted line or None."""
    key, text, sid = entry
    clean_text = text.replace("\n", " ").strip()
    phn = get_phonemes(clean_text)
    if phn:
        return f"{WAV_REL_PATH}/{key}.wav|{phn}|{LANGUAGE_ID}|{sid}|{clean_text}"
    return None


def load_transcripts():
    """Load all transcripts from JSON files."""
    entries = []
    
    male_json = MALE_DIR / "IISc_SYSPINProject_Hindi_Male_Spk001_HC_Transcripts.json"
    female_json = FEMALE_DIR / "IISc_SYSPINProject_Hindi_Female_Spk001_HC_Transcripts.json"
    
    log.info(f"Loading {male_json}")
    with open(male_json, "r", encoding="utf-8") as f:
        data = json.load(f)
        for key, val in data.get("Transcripts", {}).items():
            entries.append((key, val["Transcript"], SPEAKER_ID_MAP["male"]))
    log.info(f"Loaded {len(entries)} male entries")
    
    male_count = len(entries)
    log.info(f"Loading {female_json}")
    with open(female_json, "r", encoding="utf-8") as f:
        data = json.load(f)
        for key, val in data.get("Transcripts", {}).items():
            entries.append((key, val["Transcript"], SPEAKER_ID_MAP["female"]))
    log.info(f"Loaded {len(entries) - male_count} female entries")
    
    return entries


def test_espeak():
    """Test that espeak-ng works."""
    log.info("Testing espeak-ng...")
    result = get_phonemes("नमस्ते")
    if result:
        log.info(f"espeak-ng test passed: नमस्ते -> {result}")
        return True
    else:
        log.error("espeak-ng test FAILED!")
        return False


def phonemize_worker(texts, language, q):
    """Worker for phonemization in a separate process."""
    try:
        phns = phonemize(
            texts,
            language=language,
            backend="espeak",
            strip=True,
            preserve_punctuation=True,
            with_stress=True,
        )
        q.put(phns)
    except Exception:
        q.put(None)


def phonemize_batch_safe(batch_texts, language="hi"):
    """Run phonemizer in a separate process to avoid crashing the main process.

    Returns a list of phoneme strings or None if the worker crashed.
    """
    ctx = mp.get_context("spawn")
    queue = ctx.Queue()

    proc = ctx.Process(target=phonemize_worker, args=(batch_texts, language, queue))
    proc.start()
    proc.join()

    if proc.exitcode != 0:
        return None

    try:
        return queue.get_nowait()
    except Exception:
        return None


def main():
    log.info("Starting dataset creation...")
    
    if not test_espeak():
        return
    
    log.info("Loading transcripts...")
    entries = load_transcripts()
    total = len(entries)
    log.info(f"Total entries to process: {total}")
    
    OUTPUT_WAV_DIR.mkdir(exist_ok=True)
    log.info(f"Output wav dir: {OUTPUT_WAV_DIR}")
    

    log.info(f"Starting IPA conversion in batches...")
    
    results = []
    failed = 0
    BATCH_SIZE = 200
    
    for i in tqdm(range(0, total, BATCH_SIZE), desc="IPA conversion", mininterval=0.5):
        batch = entries[i:i+BATCH_SIZE]
        batch_texts = [e[1].replace("\n", " ").strip() for e in batch]
        
        phonemes_list = phonemize_batch_safe(batch_texts, language="hi")
        if phonemes_list is None:
            log.error(f"Batch crashed at index {i}. Skipping this batch...")
            failed += len(batch)
            continue

        for j, phn in enumerate(phonemes_list):
            if phn:
                key, _, sid = batch[j]
                clean_text = batch_texts[j]
                phn = " ".join(phn.split())
                results.append(f"{WAV_REL_PATH}/{key}.wav|{phn}|{LANGUAGE_ID}|{sid}|{clean_text}")
            else:
                failed += 1

    log.info(f"Conversion complete. Success: {len(results)}, Failed: {failed}")
    log.info(f"Writing {OUTPUT_FILE}...")
    
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        f.write("\n".join(results) + "\n")
    
    log.info(f"Wrote {len(results)} lines to {OUTPUT_FILE}")

    log.info(f"Merging with {EXISTING_TRAIN_LIST}...")
    all_lines = []
    if EXISTING_TRAIN_LIST.exists():
        with open(EXISTING_TRAIN_LIST, "r", encoding="utf-8") as f:
            content = f.read().strip()
            if content:
                all_lines = content.split("\n")
    
    log.info(f"Loaded {len(all_lines)} lines from existing list.")
    all_lines.extend(results)
    
    log.info(f"Writing {len(all_lines)} lines to {MERGED_TRAIN_LIST}...")
    with open(MERGED_TRAIN_LIST, "w", encoding="utf-8") as f:
        f.write("\n".join(all_lines) + "\n")
    log.info(f"Merged file created at {MERGED_TRAIN_LIST}")
    
    # Move wavs
    log.info("Collecting wav files...")
    wav_files = list((MALE_DIR / "wav").glob("*.wav")) + list((FEMALE_DIR / "wav").glob("*.wav"))
    log.info(f"Found {len(wav_files)} wav files to move")
    
    log.info("Moving wav files...")
    for wav in tqdm(wav_files, desc="Moving wavs", mininterval=0.5):
        shutil.move(str(wav), str(OUTPUT_WAV_DIR / wav.name))
    
    log.info(f"Done! Dataset: {len(results)} entries, Wavs: {len(wav_files)} files")


if __name__ == "__main__":
    main()