#!/usr/bin/env python3
"""
VoxCPM 1.5 Dialogue Generator — Voice Cloning via Continuation Mode

VoxCPM 1.5 supports voice cloning in "Continuation" mode: you provide
a reference wav AND its exact transcript, and the model continues from
that audio, faithfully preserving timbre, rhythm, emotion, and style.

This is different from VoxCPM2's "Isolated Reference" mode which only
needs reference_wav_path. VoxCPM1.5 requires prompt_wav_path + prompt_text.

API comparison:
  VoxCPM 1.5 (Continuation):
    model.generate(
        text="New text to synthesize.",
        prompt_wav_path="speaker.wav",        # reference audio
        prompt_text="Exact transcript of wav", # REQUIRED
    )

  VoxCPM2 (Isolated Reference):
    model.generate(
        text="New text to synthesize.",
        reference_wav_path="speaker.wav",     # only this needed
    )

Usage (single file):
    python dialog.py dialog.txt                    # Output: dialogue_final.mp3
    python dialog.py dialog.txt -o my_audio.mp3    # Custom output
    python dialog.py dialog.txt --timesteps 5      # Faster, slightly lower quality
    python dialog.py dialog.txt --cfg 3.0          # Tighter voice cloning
    python dialog.py dialog.txt --speaker-speeds bruin=1.2,alan=0.9  # Per-speaker speed
    python dialog.py dialog.txt --pause 0.3        # Shorter pause between parts
    python dialog.py dialog.txt --dry-run          # Show parsed lines without generating

Usage (batch / folder):
    python dialog.py ai-chat/                      # TTS all .txt files in folder
    python dialog.py ai-chat/ --samples-dir ~/voices/  # Voice samples in separate dir
    python dialog.py ai-chat/ --dry-run            # Preview all files without generating

Dialogue file format:
    [speaker1]: First line of dialogue.
    [speaker2]: Response to the first line.

Voice samples & transcripts:
    Place <speaker>.wav AND <speaker>.txt files in the same directory as the
    dialogue file, or specify a different directory with --samples-dir.
    The .txt file must contain the exact transcript of the .wav file.
"""

import argparse
import os
import re
import subprocess
import sys
import tempfile
import numpy as np
import soundfile as sf

# VoxCPM import
try:
    from voxcpm import VoxCPM
except ImportError:
    print("Error: voxcpm not installed. Run: pip install voxcpm")
    sys.exit(1)


def parse_dialogue(path):
    """Parse dialogue file into list of (speaker, text) tuples.

    Lines starting with [speaker]: start a new entry.
    Continuation lines (no [speaker]: prefix) are appended to the previous
    speaker's text.  Blank lines are preserved as paragraph breaks.
    """
    lines = []
    with open(path) as f:
        for raw_line in f:
            stripped = raw_line.strip()
            m = re.match(r'\[(.+?)\]:\s*(.*)', stripped)
            if m:
                speaker = m.group(1).strip().lower()
                text = m.group(2).strip()
                lines.append((speaker, text))
            elif lines and stripped:
                # Continuation line — append to previous speaker's text
                speaker, prev_text = lines[-1]
                lines[-1] = (speaker, prev_text + "\n" + stripped)
    return lines

def split_into_chunks(text, max_chars=500):
    """Split text into chunks for VRAM-friendly generation.

    Tries to break on sentence boundaries (periods, question marks, newlines).
    Falls back to hard split if a single sentence exceeds max_chars.
    """
    chunks = []
    current = []
    current_len = 0

    for sentence in re.split(r'([.!?]\s+|\n)', text):
        if sentence in (' ', '  ', '\n', '\n ', ' \n'):
            # whitespace-only delimiter, carry with current
            current.append(sentence)
            continue

        seg_len = len(sentence)
        if current_len + seg_len > max_chars and current:
            chunks.append(''.join(current))
            current = []
            current_len = 0

        current.append(sentence)
        current_len += seg_len

    if current:
        chunks.append(''.join(current))

    return [c.strip() for c in chunks if c.strip()]

def generate_voice(model, text, speaker_name, samples_dir, cfg_value, inference_timesteps):
    """
    Generate speech using VoxCPM 1.5 Continuation voice cloning.

    Requires:
        prompt_wav_path: reference audio file
        prompt_text: exact transcript of the reference audio (REQUIRED for 1.5)
    """
    wav_path = os.path.join(samples_dir, f"{speaker_name}.wav")
    txt_path = os.path.join(samples_dir, f"{speaker_name}.txt")

    with open(txt_path, 'r') as f:
        prompt_text = f.read().strip()

    # Split long text into chunks to avoid cuda-OOM
    chunks = split_into_chunks(text, max_chars=500)
    sample_rate = model.tts_model.sample_rate

    wav_parts = []
    for chunk in chunks:
        wav = model.generate(
            text=chunk,
            prompt_wav_path=wav_path,
            prompt_text=prompt_text,
            cfg_value=cfg_value,
            inference_timesteps=inference_timesteps,
        )
        wav_parts.append(wav)

    # Concatenate all parts for this dialogue line
    if len(wav_parts) > 1:
        wav = np.concatenate(wav_parts)
    else:
        wav = wav_parts[0]

    return wav, sample_rate

def adjust_speed(input_wav, output_wav, speed):
    """Apply ffmpeg atempo filter to adjust playback speed.

    speed: 1.0 = normal, 0.5 = half, 2.0 = double.
    ffmpeg atempo is limited to 0.5–2.0 per filter; values outside
    this range are chained (e.g., 4.0 → atempo=2.0,atempo=2.0).
    """
    if speed == 1.0:
        return  # no change needed

    # Build atempo chain for arbitrary speed
    filters = []
    remaining = speed
    while remaining > 2.0:
        filters.append("atempo=2.0")
        remaining /= 2.0
    while remaining < 0.5:
        filters.append("atempo=0.5")
        remaining /= 0.5
    filters.append(f"atempo={remaining:.4f}")

    filter_str = ",".join(filters)
    subprocess.run(
        ["ffmpeg", "-y", "-i", input_wav, "-filter:a", filter_str, output_wav],
        check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )


def generate_silence(output_wav, duration_sec, sample_rate):
    """Generate a silent WAV file of given duration."""
    samples = int(duration_sec * sample_rate)
    silence = np.zeros(samples, dtype=np.float32)
    sf.write(output_wav, silence, sample_rate)


def main():
    parser = argparse.ArgumentParser(description="VoxCPM 1.5 Dialogue Generator (Continuation cloning)")
    parser.add_argument("input_path", help="Path to a dialogue text file OR a directory containing .txt files")
    parser.add_argument("-o", "--output", default=None,
                        help="Output MP3 file (single-file mode only, default: dialogue_final.mp3)")
    parser.add_argument("--samples-dir", default=None,
                        help="Directory containing voice sample WAV files and transcript TXT files (default: same as input file/directory)")
    parser.add_argument("--timesteps", type=int, default=20,
                        help="Inference timesteps: higher = better quality, slower (5-20, default: 10)")
    parser.add_argument("--cfg", type=float, default=2.0,
                        help="CFG value: higher = tighter to reference voice (1.0-4.0, default: 2.0)")
    parser.add_argument("--speaker-speeds", default="bruin=1.2,alan=1.0",
                        help="Comma-separated speaker=speed pairs for playback speed adjustment (e.g. bruin=1.2,alan=0.9). Speed 1.0 = normal, 0.5 = half, 2.0 = double.")
    parser.add_argument("--pause", type=float, default=0.5,
                        help="Pause duration in seconds between dialogue parts (default: 0.5)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Parse and show dialogue lines without generating audio")
    args = parser.parse_args()

    # Parse speaker speeds
    speaker_speeds = {}
    if args.speaker_speeds:
        for pair in args.speaker_speeds.split(","):
            pair = pair.strip()
            if "=" in pair:
                k, v = pair.split("=", 1)
                speaker_speeds[k.strip().lower()] = float(v.strip())

    # Determine mode: single file or batch directory
    input_path = os.path.abspath(args.input_path)
    if os.path.isdir(input_path):
        mode = "batch"
        # Collect and sort all .txt files in the directory
        txt_files = sorted(
            os.path.join(input_path, f)
            for f in os.listdir(input_path)
            if f.endswith(".txt")
        )
        if not txt_files:
            print(f"Error: No .txt files found in {input_path}")
            sys.exit(1)
        print(f"Batch mode: {len(txt_files)} files in {input_path}")
        # samples-dir defaults to the input directory itself
        samples_dir = os.path.abspath(args.samples_dir) if args.samples_dir else input_path
    else:
        mode = "single"
        txt_files = [input_path]
        dialog_dir = os.path.dirname(input_path)
        samples_dir = os.path.abspath(args.samples_dir) if args.samples_dir else dialog_dir

    # Collect all speakers across all files (for batch validation)
    all_speakers = set()
    file_lines = {}
    for fpath in txt_files:
        lines = parse_dialogue(fpath)
        if not lines:
            print(f"Warning: No dialogue lines in {fpath}, skipping.")
            continue
        file_lines[fpath] = lines
        all_speakers.update(s for s, _ in lines)

    if not file_lines:
        print("No valid dialogue files found.")
        sys.exit(1)

    print(f"Speakers: {', '.join(sorted(all_speakers))}")

    # Verify voice samples AND transcripts exist
    missing = []
    for speaker in all_speakers:
        wav_path = os.path.join(samples_dir, f"{speaker}.wav")
        txt_path = os.path.join(samples_dir, f"{speaker}.txt")
        if not os.path.exists(wav_path):
            missing.append(f"{speaker}.wav")
        if not os.path.exists(txt_path):
            missing.append(f"{speaker}.txt")
    if missing:
        print(f"Error: Missing files: {', '.join(missing)}")
        print(f"Looking in: {samples_dir}")
        print("Note: VoxCPM 1.5 Continuation cloning requires BOTH .wav and .txt (transcript) files.")
        sys.exit(1)

    if args.dry_run:
        for fpath, lines in file_lines.items():
            print(f"\n{'='*60}")
            print(f"File: {fpath}")
            print(f"{'='*60}")
            for i, (speaker, text) in enumerate(lines):
                preview = text[:80] + "..." if len(text) > 80 else text
                print(f"  [{speaker}]: {preview}")
            # Show transcript for first speaker found
            first_speaker = lines[0][0]
            txt_path = os.path.join(samples_dir, f"{first_speaker}.txt")
            with open(txt_path) as f:
                transcript = f.read().strip()
            print(f"    [{first_speaker}] transcript: {transcript[:80]}...")
        return

    # Load model once
    print("\nLoading VoxCPM 1.5 model... (first run downloads ~1.8 GB)")
    model = VoxCPM.from_pretrained("openbmb/VoxCPM1.5")
    print("Model loaded.\n")

    # Process each file
    for fpath in txt_files:
        if fpath not in file_lines:
            continue  # skipped earlier (empty)

        lines = file_lines[fpath]
        basename = os.path.splitext(os.path.basename(fpath))[0]
        file_dir = os.path.dirname(fpath)

        # Output: same directory as the .txt file
        if mode == "single" and args.output:
            output_mp3 = args.output
        else:
            output_mp3 = os.path.join(file_dir, f"{basename}.mp3")

        print(f"\n{'='*60}")
        print(f"Processing: {fpath}")
        print(f"Output: {output_mp3}")
        print(f"{'='*60}")

        wav_files = []
        sample_rate = None  # will be set by first generate_voice call
        try:
            with tempfile.TemporaryDirectory() as tmpdir:
                for i, (speaker, text) in enumerate(lines):
                    print(f"  [{i+1}/{len(lines)}] [{speaker}]: {text[:60]}...")

                    wav, sample_rate = generate_voice(
                        model, text, speaker, samples_dir, args.cfg, args.timesteps
                    )

                    out_wav = os.path.join(tmpdir, f"part_{i:03d}.wav")
                    sf.write(out_wav, wav, sample_rate)

                    # Apply per-speaker speed adjustment
                    speed = speaker_speeds.get(speaker, 1.0)
                    if speed != 1.0:
                        sped_wav = os.path.join(tmpdir, f"part_{i:03d}_speed.wav")
                        print(f"    -> adjusting speed: {speed:.2f}x")
                        adjust_speed(out_wav, sped_wav, speed)
                        out_wav = sped_wav

                    wav_files.append(out_wav)

                # Build concat list with silence between parts
                concat_file = os.path.join(tmpdir, "concat_list.txt")
                with open(concat_file, "w") as f:
                    for idx, wf in enumerate(wav_files):
                        f.write(f"file '{wf}'\n")
                        if idx < len(wav_files) - 1 and args.pause > 0:
                            silence_wav = os.path.join(tmpdir, f"silence_{idx:03d}.wav")
                            generate_silence(silence_wav, args.pause, sample_rate)
                            f.write(f"file '{silence_wav}'\n")

                # Convert to MP3
                print(f"\n  Concatenating {len(wav_files)} parts -> {output_mp3}")
                subprocess.run(
                    [
                        "ffmpeg", "-y",
                        "-f", "concat", "-safe", "0",
                        "-i", concat_file,
                        "-c:a", "libmp3lame", "-q:a", "2",
                        output_mp3,
                    ],
                    check=True,
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )

                file_size = os.path.getsize(output_mp3)
                print(f"  Done! {output_mp3} ({file_size / 1024 / 1024:.1f} MB)")

        except KeyboardInterrupt:
            print("\n  Interrupted. Partial WAV parts will be cleaned up.")
            raise
        except Exception as e:
            print(f"  Error processing {fpath}: {e}")
            continue  # continue with next file in batch

    print(f"\n{'='*60}")
    print("All done!")
    if mode == "batch":
        completed = sum(1 for f in txt_files if f in file_lines and os.path.exists(os.path.join(os.path.dirname(f), os.path.splitext(os.path.basename(f))[0] + ".mp3")))
        print(f"Generated {completed}/{len(txt_files)} MP3 files.")


if __name__ == "__main__":
    main()