"""
Preprocess an audio+text dataset into the durations.jsonl format required
by Inflect-Nano training scripts.

Supports:
  - HuggingFace datasets (CMU ARCTIC, LJSpeech, etc.)
  - Local directory of .wav files with matching .txt transcriptions
  - LJSpeech-format metadata.csv

Output: a .jsonl file with one JSON object per line containing:
  phone_ids, tone_ids, lang_ids, hifigan_durations, target_audio, speaker_id, voice_id
"""
from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torchaudio

REPO_ROOT = Path(__file__).resolve().parent
VENDORED_FRONTEND = REPO_ROOT / "third_party" / "tiny_tts_frontend"
sys.path.insert(0, str(REPO_ROOT))
sys.path.insert(0, str(VENDORED_FRONTEND))

from tiny_tts.nn import commons
from tiny_tts.text import phonemes_to_ids
from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
from tiny_tts.utils import ADD_BLANK

from inflect_nano.text_cleaning import clean_tinytts_text
from inflect_nano.vocoder import HifiGanConfig, MelFrontend


def text_to_ids(text: str) -> tuple[list[int], list[int], list[int]]:
    """Convert English text to TinyTTS phone/tone/lang ID lists."""
    cleaned = clean_tinytts_text(text)
    normalized = normalize_text(cleaned)
    phones, tones, _ = grapheme_to_phoneme(normalized)
    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
    if ADD_BLANK:
        phone_ids = commons.insert_blanks(phone_ids, 0)
        tone_ids = commons.insert_blanks(tone_ids, 0)
        lang_ids = commons.insert_blanks(lang_ids, 0)
    return phone_ids, tone_ids, lang_ids


def estimate_durations_uniform(
    phone_ids: list[int],
    audio_path: str,
    mel_frontend: MelFrontend,
    sample_rate: int = 24000,
) -> list[int]:
    """Estimate mel-frame durations uniformly across all phones.

    Loads the audio, extracts mel spectrogram, then distributes the total
    frame count evenly across phones.  This is a coarse estimate suitable
    for decoder-only fine-tuning where the duration predictor is frozen.
    """
    wav, sr = torchaudio.load(audio_path)
    if wav.shape[0] > 1:
        wav = wav.mean(dim=0, keepdim=True)
    if sr != sample_rate:
        wav = torchaudio.functional.resample(wav, sr, sample_rate)
    wav = wav.clamp(-1, 1)

    with torch.no_grad():
        mel = mel_frontend(wav)
    total_frames = mel.shape[-1]

    # Count non-blank phones (id != 0)
    visible_phones = [i for i, pid in enumerate(phone_ids) if pid != 0]

    if not visible_phones:
        return [0] * len(phone_ids)

    # Distribute frames: each blank gets 1 frame, rest split evenly
    blank_count = len(phone_ids) - len(visible_phones)
    remaining = max(0, total_frames - blank_count)
    base_dur = max(1, remaining // len(visible_phones))
    remainder = remaining - base_dur * len(visible_phones)

    durations = []
    for pid in phone_ids:
        if pid == 0:
            durations.append(1)
        else:
            extra = 1 if remainder > 0 else 0
            remainder = max(0, remainder - 1)
            durations.append(base_dur + extra)

    # Scale to match total_frames
    current_sum = sum(durations)
    if current_sum > 0 and current_sum != total_frames:
        scale = total_frames / current_sum
        scaled = [max(1, round(d * scale)) for d in durations]
        # Fix rounding errors
        diff = total_frames - sum(scaled)
        for i in range(abs(diff)):
            if diff > 0:
                scaled[i % len(scaled)] += 1
            else:
                if scaled[i % len(scaled)] > 1:
                    scaled[i % len(scaled)] -= 1
        durations = scaled

    # Ensure no zero durations for non-blank phones
    for i, pid in enumerate(phone_ids):
        if pid != 0 and durations[i] < 1:
            durations[i] = 1

    return durations


def process_hf_dataset(
    dataset_path: str,
    output_jsonl: Path,
    audio_dir: Path | None,
    subset: str = "default",
    split: str | None = None,
    text_key: str = "text",
    speaker_key: str = "speaker",
    max_rows: int = 0,
    voice_id: str = "speaker",
) -> int:
    """Process a HuggingFace dataset (loaded via `datasets` library)."""
    from datasets import load_dataset

    print(f"Loading HF dataset: {dataset_path}  subset={subset}  split={split}")
    if split:
        ds = load_dataset(dataset_path, subset, split=split, trust_remote_code=True)
    else:
        ds_dict = load_dataset(dataset_path, subset, trust_remote_code=True)
        splits = list(ds_dict.keys())
        print(f"Available splits: {splits}")
        # Use first train split, or first available
        preferred = [s for s in splits if "train" in s.lower()]
        ds = ds_dict[preferred[0] if preferred else splits[0]]

    mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
    output_jsonl.parent.mkdir(parents=True, exist_ok=True)

    count = 0
    with output_jsonl.open("w", encoding="utf-8") as f:
        for i, row in enumerate(ds):
            text = str(row.get(text_key, "")).strip()
            if not text:
                continue

            # Get audio path or array
            audio_info = row.get("audio", row.get("file", None))
            if audio_info is None:
                continue

            if isinstance(audio_info, dict):
                # Audio is already loaded as array
                audio_path = None
                audio_array = audio_info.get("array")
                sample_rate = audio_info.get("sampling_rate", 24000)
                if audio_array is None:
                    continue
            elif isinstance(audio_info, str):
                audio_path = audio_info
                if audio_dir:
                    audio_path = str(audio_dir / Path(audio_info).name)
                if not Path(audio_path).is_file():
                    continue
                audio_array = None
                sample_rate = 24000  # will be detected on load
            else:
                continue

            try:
                phone_ids, tone_ids, lang_ids = text_to_ids(text)
            except Exception as e:
                print(f"  Skipping row {i}: text-to-ids failed: {e}")
                continue

            if not phone_ids:
                continue

            # Estimate durations
            if audio_path:
                durations = estimate_durations_uniform(phone_ids, audio_path, mel_frontend)
            else:
                # Uniform fallback: 8 frames per phone
                durations = [8] * len(phone_ids)

            speaker = str(row.get(speaker_key, voice_id))

            row_out = {
                "phone_ids": phone_ids,
                "tone_ids": tone_ids,
                "lang_ids": lang_ids,
                "hifigan_durations": durations,
                "target_audio": audio_path or "",
                "speaker_id": hash(speaker) % 256,
                "voice_id": speaker,
            }
            f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
            count += 1

            if max_rows > 0 and count >= max_rows:
                break

            if count % 100 == 0:
                print(f"  Processed {count} rows...")

    print(f"Wrote {count} rows to {output_jsonl}")
    return count


def process_local_dir(
    audio_dir: Path,
    output_jsonl: Path,
    ext: str = ".wav",
    text_ext: str = ".txt",
    speaker: str = "speaker",
    max_rows: int = 0,
) -> int:
    """Process a local directory of .wav files with matching .txt files."""
    audio_files = sorted(audio_dir.glob(f"*{ext}"))
    if not audio_files:
        audio_files = sorted(audio_dir.rglob(f"*{ext}"))

    mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
    output_jsonl.parent.mkdir(parents=True, exist_ok=True)

    count = 0
    with output_jsonl.open("w", encoding="utf-8") as f:
        for wav_path in audio_files:
            txt_path = wav_path.with_suffix(text_ext)
            if not txt_path.is_file():
                continue

            text = txt_path.read_text(encoding="utf-8").strip()
            if not text:
                continue

            try:
                phone_ids, tone_ids, lang_ids = text_to_ids(text)
            except Exception as e:
                print(f"  Skipping {wav_path.name}: {e}")
                continue

            durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)

            row_out = {
                "phone_ids": phone_ids,
                "tone_ids": tone_ids,
                "lang_ids": lang_ids,
                "hifigan_durations": durations,
                "target_audio": str(wav_path.resolve()),
                "speaker_id": hash(speaker) % 256,
                "voice_id": speaker,
            }
            f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
            count += 1

            if max_rows > 0 and count >= max_rows:
                break

            if count % 100 == 0:
                print(f"  Processed {count} rows...")

    print(f"Wrote {count} rows to {output_jsonl}")
    return count


def process_metadata_csv(
    csv_path: Path,
    audio_dir: Path,
    output_jsonl: Path,
    sep: str = "|",
    header: bool = False,
    text_col: int = 1,
    file_col: int = 0,
    speaker: str = "speaker",
    max_rows: int = 0,
) -> int:
    """Process an LJSpeech-style metadata.csv."""
    mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
    output_jsonl.parent.mkdir(parents=True, exist_ok=True)

    count = 0
    with csv_path.open("r", encoding="utf-8") as csv_file:
        lines = csv_file.readlines()
        if header:
            lines = lines[1:]

        with output_jsonl.open("w", encoding="utf-8") as out_f:
            for line in lines:
                line = line.strip()
                if not line:
                    continue

                parts = line.split(sep)
                if len(parts) < max(text_col, file_col) + 1:
                    continue

                filename = parts[file_col].strip()
                if not filename.endswith(".wav"):
                    filename += ".wav"
                wav_path = audio_dir / filename
                if not wav_path.is_file():
                    print(f"  Missing: {wav_path}")
                    continue

                text = parts[text_col].strip()
                if not text:
                    continue

                try:
                    phone_ids, tone_ids, lang_ids = text_to_ids(text)
                except Exception as e:
                    print(f"  Skipping {filename}: {e}")
                    continue

                durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)

                row_out = {
                    "phone_ids": phone_ids,
                    "tone_ids": tone_ids,
                    "lang_ids": lang_ids,
                    "hifigan_durations": durations,
                    "target_audio": str(wav_path.resolve()),
                    "speaker_id": hash(speaker) % 256,
                    "voice_id": speaker,
                }
                out_f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
                count += 1

                if max_rows > 0 and count >= max_rows:
                    break

                if count % 100 == 0:
                    print(f"  Processed {count} rows...")

    print(f"Wrote {count} rows to {output_jsonl}")
    return count


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Preprocess a speech dataset for Inflect-Nano training.",
    )

    sub = ap.add_subparsers(dest="mode", required=True)

    # HuggingFace mode
    p_hf = sub.add_parser("hf", help="Process a HuggingFace dataset")
    p_hf.add_argument("--dataset", required=True, help="HF dataset path, e.g. MikhailT/cmu-arctic")
    p_hf.add_argument("--subset", default="default")
    p_hf.add_argument("--split", default=None, help="e.g. 'rms' for CMU ARCTIC single speaker")
    p_hf.add_argument("--text-key", default="text")
    p_hf.add_argument("--speaker-key", default="speaker")
    p_hf.add_argument("--audio-dir", type=Path, default=None, help="Local dir for audio if HF paths are remote")
    p_hf.add_argument("--out", type=Path, required=True, help="Output .jsonl path")
    p_hf.add_argument("--max-rows", type=int, default=0)
    p_hf.add_argument("--voice-id", default="speaker")

    # Local dir mode
    p_dir = sub.add_parser("local", help="Process local directory of .wav+.txt pairs")
    p_dir.add_argument("--audio-dir", type=Path, required=True)
    p_dir.add_argument("--ext", default=".wav")
    p_dir.add_argument("--text-ext", default=".txt")
    p_dir.add_argument("--speaker", default="speaker")
    p_dir.add_argument("--out", type=Path, required=True)
    p_dir.add_argument("--max-rows", type=int, default=0)

    # Metadata CSV mode (LJSpeech-style)
    p_csv = sub.add_parser("csv", help="Process a metadata.csv file")
    p_csv.add_argument("--csv", type=Path, required=True)
    p_csv.add_argument("--audio-dir", type=Path, required=True)
    p_csv.add_argument("--sep", default="|")
    p_csv.add_argument("--header", action="store_true")
    p_csv.add_argument("--file-col", type=int, default=0)
    p_csv.add_argument("--text-col", type=int, default=1)
    p_csv.add_argument("--speaker", default="speaker")
    p_csv.add_argument("--out", type=Path, required=True)
    p_csv.add_argument("--max-rows", type=int, default=0)

    args = ap.parse_args()

    if args.mode == "hf":
        process_hf_dataset(
            dataset_path=args.dataset,
            output_jsonl=args.out,
            audio_dir=args.audio_dir,
            subset=args.subset,
            split=args.split,
            text_key=args.text_key,
            speaker_key=args.speaker_key,
            max_rows=args.max_rows,
            voice_id=args.voice_id,
        )
    elif args.mode == "local":
        process_local_dir(
            audio_dir=args.audio_dir,
            output_jsonl=args.out,
            ext=args.ext,
            text_ext=args.text_ext,
            speaker=args.speaker,
            max_rows=args.max_rows,
        )
    elif args.mode == "csv":
        process_metadata_csv(
            csv_path=args.csv,
            audio_dir=args.audio_dir,
            output_jsonl=args.out,
            sep=args.sep,
            header=args.header,
            file_col=args.file_col,
            text_col=args.text_col,
            speaker=args.speaker,
            max_rows=args.max_rows,
        )


if __name__ == "__main__":
    main()