#!/usr/bin/env python3
"""Convert a long text file to MP3 using the existing Riva TTS client.

The script splits input text into smaller chunks, synthesizes each chunk with
`python-clients/scripts/tts/talk.py`, joins all WAV chunks, and finally converts
the result to MP3 using ffmpeg.
"""

from __future__ import annotations

import argparse
import re
import shutil
import subprocess
import sys
import tempfile
import wave
from pathlib import Path
from typing import List


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Convert a long TXT file to MP3 using NVIDIA Riva TTS.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--input", type=Path, required=True, help="Path to input .txt file.")
    parser.add_argument("--output", type=Path, required=True, help="Path to output .mp3 file.")
    parser.add_argument("--server", default="0.0.0.0:50051", help="Riva server address.")
    parser.add_argument("--language-code", default="pl-PL", help="Language code for TTS.")
    parser.add_argument(
        "--voice",
        default="Chatterbox-Multilingual.pl-PL.Male",
        help="Riva voice name.",
    )
    parser.add_argument(
        "--talk-script",
        type=Path,
        default=Path("python-clients/scripts/tts/talk.py"),
        help="Path to existing talk.py script.",
    )
    parser.add_argument(
        "--max-chars",
        type=int,
        default=700,
        help="Maximum characters per TTS chunk.",
    )
    parser.add_argument(
        "--sample-rate-hz",
        type=int,
        default=22050,
        help="Output sample rate used by talk.py.",
    )
    parser.add_argument(
        "--keep-wav",
        action="store_true",
        help="Keep final merged WAV next to MP3.",
    )
    return parser.parse_args()


def normalize_whitespace(text: str) -> str:
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    lines = [line.strip() for line in text.split("\n")]
    return "\n".join(lines)


def split_sentence_if_needed(sentence: str, max_chars: int) -> List[str]:
    if len(sentence) <= max_chars:
        return [sentence]

    words = sentence.split()
    parts: List[str] = []
    current = ""
    for word in words:
        candidate = f"{current} {word}".strip() if current else word
        if len(candidate) <= max_chars:
            current = candidate
        else:
            if current:
                parts.append(current)
            # Single very long token fallback.
            if len(word) > max_chars:
                for i in range(0, len(word), max_chars):
                    parts.append(word[i : i + max_chars])
                current = ""
            else:
                current = word
    if current:
        parts.append(current)
    return parts


def chunk_text(text: str, max_chars: int) -> List[str]:
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    sentence_splitter = re.compile(r"(?<=[.!?…])\s+")

    chunks: List[str] = []
    current = ""

    for paragraph in paragraphs:
        sentences = [s.strip() for s in sentence_splitter.split(paragraph) if s.strip()]
        expanded_sentences: List[str] = []
        for s in sentences:
            expanded_sentences.extend(split_sentence_if_needed(s, max_chars))

        for sentence in expanded_sentences:
            candidate = f"{current} {sentence}".strip() if current else sentence
            if len(candidate) <= max_chars:
                current = candidate
            else:
                if current:
                    chunks.append(current)
                current = sentence

        if current:
            chunks.append(current)
            current = ""

    if current:
        chunks.append(current)

    if not chunks:
        raise ValueError("Input text is empty after preprocessing.")

    return chunks


def synthesize_chunk(
    talk_script: Path,
    server: str,
    language_code: str,
    voice: str,
    sample_rate_hz: int,
    text: str,
    output_wav: Path,
) -> None:
    cmd = [
        sys.executable,
        str(talk_script),
        "--server",
        server,
        "--language-code",
        language_code,
        "--voice",
        voice,
        "--text",
        text,
        "--stream",
        "--sample-rate-hz",
        str(sample_rate_hz),
        "--output",
        str(output_wav),
    ]

    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(
            "TTS synthesis failed.\n"
            f"Command: {' '.join(cmd)}\n"
            f"stdout:\n{result.stdout}\n"
            f"stderr:\n{result.stderr}"
        )


def merge_wav_files(wav_files: List[Path], output_wav: Path) -> None:
    if not wav_files:
        raise ValueError("No WAV files to merge.")

    with wave.open(str(wav_files[0]), "rb") as first:
        params = first.getparams()
        frames = [first.readframes(first.getnframes())]

    for wav_path in wav_files[1:]:
        with wave.open(str(wav_path), "rb") as wf:
            if (
                wf.getnchannels() != params.nchannels
                or wf.getsampwidth() != params.sampwidth
                or wf.getframerate() != params.framerate
                or wf.getcomptype() != params.comptype
            ):
                raise RuntimeError(f"Incompatible WAV params in {wav_path}")
            frames.append(wf.readframes(wf.getnframes()))

    with wave.open(str(output_wav), "wb") as out:
        out.setparams(params)
        for chunk in frames:
            out.writeframes(chunk)


def convert_wav_to_mp3(input_wav: Path, output_mp3: Path) -> None:
    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg is None:
        raise RuntimeError(
            "ffmpeg not found in PATH. Install ffmpeg to generate MP3 output."
        )

    cmd = [
        ffmpeg,
        "-y",
        "-i",
        str(input_wav),
        "-codec:a",
        "libmp3lame",
        "-q:a",
        "2",
        str(output_mp3),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(
            "MP3 conversion failed.\n"
            f"Command: {' '.join(cmd)}\n"
            f"stdout:\n{result.stdout}\n"
            f"stderr:\n{result.stderr}"
        )


def main() -> int:
    args = parse_args()

    if args.max_chars < 50:
        raise ValueError("--max-chars should be at least 50.")

    input_path = args.input.expanduser().resolve()
    output_mp3 = args.output.expanduser().resolve()
    talk_script = args.talk_script.expanduser().resolve()

    if not input_path.is_file():
        raise FileNotFoundError(f"Input file not found: {input_path}")
    if not talk_script.is_file():
        raise FileNotFoundError(f"talk.py not found: {talk_script}")

    output_mp3.parent.mkdir(parents=True, exist_ok=True)
    merged_wav = output_mp3.with_suffix(".wav")

    text = input_path.read_text(encoding="utf-8")
    text = normalize_whitespace(text)
    chunks = chunk_text(text, args.max_chars)

    print(f"Input split into {len(chunks)} chunks (max {args.max_chars} chars each).")

    with tempfile.TemporaryDirectory(prefix="tts_chunks_") as tmp_dir:
        tmp_path = Path(tmp_dir)
        wav_paths: List[Path] = []

        for idx, chunk in enumerate(chunks, start=1):
            chunk_wav = tmp_path / f"chunk_{idx:05d}.wav"
            print(f"[{idx}/{len(chunks)}] Synthesizing chunk...")
            synthesize_chunk(
                talk_script=talk_script,
                server=args.server,
                language_code=args.language_code,
                voice=args.voice,
                sample_rate_hz=args.sample_rate_hz,
                text=chunk,
                output_wav=chunk_wav,
            )
            wav_paths.append(chunk_wav)

        print("Merging WAV chunks...")
        merge_wav_files(wav_paths, merged_wav)

    print("Converting WAV to MP3...")
    convert_wav_to_mp3(merged_wav, output_mp3)

    if not args.keep_wav and merged_wav.exists():
        merged_wav.unlink()

    print(f"Done: {output_mp3}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:  # pragma: no cover
        print(f"Error: {exc}", file=sys.stderr)
        raise SystemExit(1)