import os
import re
import asyncio
import tempfile
import logging
import threading
import time
import shutil
import subprocess

import gradio as gr
import edge_tts
from pydub import AudioSegment


# =========================================================
# CONFIGURATION
# =========================================================

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("MyanmarEdgeTTS")

APP_TITLE = "Myanmar Multilingual TTS Studio"

WORDS_PER_SUBTITLE = 20
MAX_CHUNK_CHARACTERS = 1800
MAX_CONCURRENT_JOBS = 2
MIN_AUDIO_BYTES = 100

DEFAULT_VOICE_ID = "it-IT-GiuseppeMultilingualNeural"
DEFAULT_SPEED = 30

OUTPUT_DIR = os.path.join(
    tempfile.gettempdir(),
    "myanmar_multilingual_tts",
)
os.makedirs(OUTPUT_DIR, exist_ok=True)


# =========================================================
# VOICE DISPLAY NAMES
# =========================================================

PINNED_VOICES = {
    "my-MM-NilarNeural",
    "my-MM-ThihaNeural",
    "it-IT-GiuseppeMultilingualNeural",
}

VOICE_DISPLAY_NAMES = {
    "it-IT-GiuseppeMultilingualNeural": "နေတိုး",
    "my-MM-NilarNeural": "ဝတ်မှုန်ရွှေရည်",
    "my-MM-ThihaNeural": "ပြေတီဦး",
}

VOICE_ALIAS_NAMES = [
    "အကယ်ဒမီမင်းမော်ကွန်း",
    "ပိုင်တံခွန်",
    "ရဲတိုက်",
    "မြင့်မြတ်",
    "နေမင်း",
    "အောင်ရဲလင်း",
    "ကျော်ကျော်ဗို",
    "လွင်မိုး",
    "ပြေတီဦး",
    "နေတိုး",
    "ဝတ်မှုန်ရွှေရည်",
    "ဖွေးဖွေး",
    "အိန္ဒြာကျော်ဇင်",
    "သက်မွန်မြင့်",
    "မိုးဟေကို",
    "ရွှေမှုံရတီ",
    "သင်ဇာဝင့်ကျော်",
    "ပိုင်ဖြိုးသု",
    "ခင်ဝင့်ဝါ",
    "အေးမြတ်သူ",
    "အိချောပို",
    "နန်းဆုဦး",
    "အေးဝတ်ရည်သောင်း",
    "ဆုပန်ထွာ",
    "မေမြင့်မိုရ်",
    "ဝါဆိုမိုးဦး",
    "သံသာမိုးသိမ့်",
    "မြတ်နိုးအေး",
    "ခင်သဇင်",
    "ယွန်းယွန်း",
    "မင်းဦး",
    "ခန့်စည်သူ",
    "ထွန်းထွန်း",
    "ဒေါင်း",
    "အောင်လေး",
    "ဇေရဲထက်",
    "ဟိန်းဝေယံ",
    "မင်းရာဇာ",
    "ရဲအောင်",
    "ကျော်ဟိန်း",
    "ဇင်ဝိုင်း",
    "လူမင်း",
    "ခိုင်သင်းကြည်",
    "စိုးမြတ်သူဇာ",
    "ထက်ထက်မိုးဦး",
    "မေသန်းနု",
    "နန္ဒာလှိုင်",
    "စိုးပြည့်သဇင်",
    "ဝိုင်းစုခိုင်သိန်း",
    "ချစ်သုဝေ",
    "မေကဗျာ",
    "မိုးပြည့်ပြည့်မောင်",
    "ယုသန္တာတင်",
    "သဲစုညိမ်း",
    "မေသက်ခိုင်",
    "ရွှေသမီး",
    "စန္ဒီမြင့်လွင်",
    "နန်းမြတ်ဖြိုးသင်း",
    "အဆိုတော်အမျိုးသား",
    "အဆိုတော်အမျိုးသမီး",
]

VOICE_CACHE = []
VOICE_CACHE_TIME = 0
VOICE_CACHE_TTL = 3600

VOICE_LOCK = threading.Lock()
JOB_SEMAPHORE = threading.Semaphore(MAX_CONCURRENT_JOBS)


# =========================================================
# VOICE MANAGEMENT
# =========================================================

async def fetch_voices_async():
    return await edge_tts.list_voices()


def get_fallback_voices():
    return [
        {
            "ShortName": "my-MM-NilarNeural",
            "FriendlyName": "Nilar",
            "Gender": "Female",
            "Locale": "my-MM",
        },
        {
            "ShortName": "my-MM-ThihaNeural",
            "FriendlyName": "Thiha",
            "Gender": "Male",
            "Locale": "my-MM",
        },
        {
            "ShortName": "it-IT-GiuseppeMultilingualNeural",
            "FriendlyName": "Giuseppe Multilingual",
            "Gender": "Male",
            "Locale": "it-IT",
        },
    ]


def get_voices(force_refresh=False):
    global VOICE_CACHE, VOICE_CACHE_TIME

    now = time.time()

    with VOICE_LOCK:
        if (
            VOICE_CACHE
            and not force_refresh
            and now - VOICE_CACHE_TIME < VOICE_CACHE_TTL
        ):
            return VOICE_CACHE

        try:
            voices = asyncio.run(fetch_voices_async())
            selected = []

            for voice in voices:
                short_name = voice.get("ShortName", "")
                if not short_name:
                    continue

                is_pinned = short_name in PINNED_VOICES
                is_multilingual = "MultilingualNeural" in short_name

                if is_pinned or is_multilingual:
                    selected.append(voice)

            available_ids = {
                voice["ShortName"] for voice in selected
            }

            fallback_map = {
                voice["ShortName"]: voice
                for voice in get_fallback_voices()
            }

            for short_name in PINNED_VOICES:
                if short_name not in available_ids:
                    selected.append(fallback_map[short_name])

            selected.sort(
                key=lambda voice: (
                    0
                    if voice["ShortName"] == DEFAULT_VOICE_ID
                    else 1
                    if voice["ShortName"] in PINNED_VOICES
                    else 2,
                    voice.get("Locale", ""),
                    voice.get("ShortName", ""),
                )
            )

            VOICE_CACHE = selected
            VOICE_CACHE_TIME = now

            logger.info("Loaded %d selected voices", len(VOICE_CACHE))
            return VOICE_CACHE

        except Exception:
            logger.exception("Unable to retrieve Edge TTS voices")

            if VOICE_CACHE:
                return VOICE_CACHE

            VOICE_CACHE = get_fallback_voices()
            VOICE_CACHE_TIME = now
            return VOICE_CACHE


def build_voice_display_map():
    """Create Myanmar display names; do not show original Voice IDs."""

    voices = get_voices()
    result = {}
    used_names = set()
    alias_index = 0

    for voice in voices:
        short_name = voice.get("ShortName", "")
        if not short_name:
            continue

        if short_name in VOICE_DISPLAY_NAMES:
            label = VOICE_DISPLAY_NAMES[short_name]
        else:
            label = None

            while alias_index < len(VOICE_ALIAS_NAMES):
                candidate = VOICE_ALIAS_NAMES[alias_index]
                alias_index += 1

                if candidate not in used_names:
                    label = candidate
                    break

            if label is None:
                label = f"အသံဇာတ်ကောင် {len(result) + 1}"

        base_label = label
        suffix = 2

        while label in used_names:
            label = f"{base_label} {suffix}"
            suffix += 1

        used_names.add(label)
        result[label] = short_name

    return result


def voice_label(voice):
    short_name = voice.get("ShortName", "")
    return VOICE_DISPLAY_NAMES.get(short_name, "")


def get_voice_map():
    return build_voice_display_map()


def get_voice_choices():
    return list(get_voice_map().keys())


def get_default_voice_label():
    voice_map = get_voice_map()

    for label, voice_id in voice_map.items():
        if voice_id == DEFAULT_VOICE_ID:
            return label

    return next(iter(voice_map), None)


def refresh_voices():
    get_voices(force_refresh=True)
    choices = get_voice_choices()
    default_label = get_default_voice_label()

    return (
        gr.update(choices=choices, value=default_label),
        f"Voice List ပြန်ရယူပြီးပါပြီ။ Voice {len(choices)} ခု ရရှိထားသည်။",
    )


# =========================================================
# SCRIPT PROCESSING
# =========================================================

def normalize_script(text):
    """Normalize whitespace without changing word order."""

    text = (text or "").replace("\r\n", "\n").replace("\r", "\n")
    lines = []

    for line in text.split("\n"):
        line = re.sub(r"[ \t]+", " ", line).strip()
        if line:
            lines.append(line)

    return "\n".join(lines)


def split_into_sentences(text):
    """Split at sentence boundaries where possible."""

    parts = re.split(r"(?<=[.!?။！？၊])\s+", text)
    return [part.strip() for part in parts if part.strip()]


def split_long_text(text, max_chars=MAX_CHUNK_CHARACTERS):
    """Fallback split at whitespace without dropping words."""

    words = text.split()
    chunks = []
    current_words = []
    current_length = 0

    for word in words:
        added_length = len(word)
        if current_words:
            added_length += 1

        if current_words and current_length + added_length > max_chars:
            chunks.append(" ".join(current_words))
            current_words = [word]
            current_length = len(word)
        else:
            current_words.append(word)
            current_length += added_length

    if current_words:
        chunks.append(" ".join(current_words))

    return chunks


def verify_text_integrity(original, subtitle_texts):
    """Verify whitespace-separated tokens are preserved in order."""

    original_words = original.split()
    output_words = []

    for subtitle in subtitle_texts:
        output_words.extend(subtitle.split())

    if original_words != output_words:
        raise RuntimeError(
            "စာသားစစ်ဆေးမှု မအောင်မြင်ပါ။ "
            "စာသားအစဉ် သို့မဟုတ် စကားလုံး ပျောက်နေသည်။"
        )

    return True


def split_for_tts(text):
    """
    TEST MODE: split after every Myanmar sentence-ending mark '။'.
    Each sentence is synthesized as its own Edge TTS chunk.
    Text order and punctuation are preserved. Long sections without
    '။' use whitespace splitting as a safety fallback.
    """

    text = (text or "").strip()
    if not text:
        return []

    # Split after each '။' while retaining the punctuation.
    # Any text between sentence-ending marks stays in its original order.
    parts = re.findall(r"[^။]*။|[^။]+$", text, flags=re.DOTALL)

    chunks = []
    for part in parts:
        part = part.strip()
        if not part:
            continue

        if len(part) > MAX_CHUNK_CHARACTERS:
            chunks.extend(
                split_long_text(part, MAX_CHUNK_CHARACTERS)
            )
        else:
            chunks.append(part)

    verify_text_integrity(text, chunks)
    return chunks

# =========================================================
# SUBTITLE SEGMENTATION
# =========================================================

def split_for_srt(text, words_per_subtitle=WORDS_PER_SUBTITLE):
    """Split SRT text at whitespace boundaries, around 20 tokens."""

    words = text.split()
    result = []

    for index in range(0, len(words), words_per_subtitle):
        result.append(" ".join(words[index:index + words_per_subtitle]))

    verify_text_integrity(text, result)
    return result


# =========================================================
# TTS SPEED
# =========================================================

def speed_to_rate(speed):
    """
    Edge TTS rate format must include an explicit sign.
    In particular, zero must be '+0%', never '0%'.
    """

    speed = int(speed)
    speed = max(-50, min(100, speed))

    if speed >= 0:
        return f"+{speed}%"

    return f"{speed}%"


# =========================================================
# AUDIO GENERATION
# =========================================================

async def synthesize_async(text, voice, rate, output_path):
    communicate = edge_tts.Communicate(
        text=text,
        voice=voice,
        rate=rate,
    )

    await communicate.save(output_path)

    if (
        not os.path.isfile(output_path)
        or os.path.getsize(output_path) < MIN_AUDIO_BYTES
    ):
        raise RuntimeError(
            "Edge TTS က အသံဖိုင်အလွတ် ပြန်ပေးထားသည်။"
        )


def synthesize_chunk(text, voice, rate, output_path):
    asyncio.run(
        synthesize_async(text, voice, rate, output_path)
    )


# =========================================================
# AUDIO JOINING
# Keep the original audio fully; no trimming is applied.
# =========================================================

def load_and_trim_audio_part(part_path):
    """Load audio without trimming the beginning or ending."""

    return AudioSegment.from_file(part_path, format="mp3")


def combine_audio_parts(audio_parts, output_path):
    """Join audio parts in order without trimming them."""

    if not audio_parts:
        raise RuntimeError("ပေါင်းစပ်ရန် Audio အပိုင်း မရှိပါ။")

    combined_audio = AudioSegment.empty()

    for part_path in audio_parts:
        segment = load_and_trim_audio_part(part_path)
        combined_audio += segment

    if len(combined_audio) <= 0:
        raise RuntimeError(
            "ပေါင်းစပ်ထားသော Audio ဖိုင်တွင် အသံမရှိပါ။"
        )

    combined_audio.export(
        output_path,
        format="mp3",
        bitrate="192k",
    )

    if (
        not os.path.isfile(output_path)
        or os.path.getsize(output_path) < MIN_AUDIO_BYTES
    ):
        raise RuntimeError(
            "ပေါင်းစပ်ထားသော Audio ဖိုင်ကို ထုတ်မရပါ။"
        )

    return len(combined_audio)


# =========================================================
# SRT TIMING
# =========================================================

def format_srt_time(milliseconds):
    milliseconds = max(0, int(milliseconds))

    hours = milliseconds // 3600000
    milliseconds %= 3600000

    minutes = milliseconds // 60000
    milliseconds %= 60000

    seconds = milliseconds // 1000
    milliseconds %= 1000

    return (
        f"{hours:02d}:"
        f"{minutes:02d}:"
        f"{seconds:02d},"
        f"{milliseconds:03d}"
    )


def make_srt_entries(subtitle_texts, total_duration_ms):
    """Estimate timing proportionally by subtitle text length."""

    if not subtitle_texts:
        return []

    weights = [
        max(1, len(text.replace(" ", "")))
        for text in subtitle_texts
    ]

    total_weight = sum(weights)
    if total_weight <= 0:
        return []

    entries = []
    elapsed = 0

    for index, (text, weight) in enumerate(
        zip(subtitle_texts, weights)
    ):
        start = elapsed

        if index == len(subtitle_texts) - 1:
            end = total_duration_ms
        else:
            duration = round(total_duration_ms * weight / total_weight)
            elapsed += duration
            end = elapsed

        entries.append(
            {
                "text": text,
                "start": start,
                "end": end,
            }
        )

    return entries


def build_srt(entries):
    blocks = []

    for index, entry in enumerate(entries, start=1):
        start = format_srt_time(entry["start"])
        end = format_srt_time(entry["end"])

        blocks.append(
            f"{index}\n"
            f"{start} --> {end}\n"
            f"{entry['text']}"
        )

    return "\n\n".join(blocks) + "\n"


# =========================================================
# MAIN GENERATION
# =========================================================

def generate_audio_and_srt(
    script,
    selected_voice,
    speed,
    progress=gr.Progress(),
):
    script = normalize_script(script)

    if not script:
        raise gr.Error("Script ထည့်ပေးပါ။")

    voice_map = get_voice_map()
    voice = voice_map.get(selected_voice)

    if not voice:
        raise gr.Error("Voice ကို ပြန်ရွေးပေးပါ။")

    voice_display = selected_voice
    rate = speed_to_rate(speed)
    tts_chunks = split_for_tts(script)

    if not tts_chunks:
        raise gr.Error("အသံထုတ်ရန် Script မရှိပါ။")

    work_dir = tempfile.mkdtemp(prefix="job_", dir=OUTPUT_DIR)
    audio_path = os.path.join(work_dir, "output.mp3")
    srt_path = os.path.join(work_dir, "output.srt")

    try:
        with JOB_SEMAPHORE:
            audio_parts = []
            total_chunks = len(tts_chunks)

            for index, chunk in enumerate(tts_chunks):
                progress(
                    0.75 * index / total_chunks,
                    desc=f"Audio ထုတ်နေသည် {index + 1}/{total_chunks}",
                )

                part_path = os.path.join(
                    work_dir,
                    f"part_{index:05d}.mp3",
                )

                synthesize_chunk(
                    chunk,
                    voice,
                    rate,
                    part_path,
                )
                audio_parts.append(part_path)

            progress(0.80, desc="Audio ပေါင်းနေသည်...")

            duration_ms = combine_audio_parts(
                audio_parts,
                audio_path,
            )

            progress(0.90, desc="SRT ပြုလုပ်နေသည်...")

            subtitle_texts = split_for_srt(script)
            verify_text_integrity(script, subtitle_texts)

            entries = make_srt_entries(
                subtitle_texts,
                duration_ms,
            )

            with open(
                srt_path,
                "w",
                encoding="utf-8-sig",
                newline="\n",
            ) as file:
                file.write(build_srt(entries))

            progress(1.0, desc="ပြီးပါပြီ!")

            logger.info(
                "Generation complete: voice=%s, rate=%s, chars=%d, chunks=%d",
                voice,
                rate,
                len(script),
                total_chunks,
            )

            return (
                audio_path,
                srt_path,
                (
                    "အောင်မြင်ပါသည်!\n"
                    f"Voice Character: {voice_display}\n"
                    f"Speed: {rate}\n"
                    f"Script: {len(script)} characters\n"
                    f"Audio chunks: {total_chunks}\n"
                    f"SRT lines: {len(subtitle_texts)}"
                ),
            )

    except Exception as exc:
        logger.exception("Audio/SRT generation failed")
        shutil.rmtree(work_dir, ignore_errors=True)

        raise gr.Error(
            f"Error: {type(exc).__name__}: {str(exc)[:400]}"
        )


# =========================================================
# USER INTERFACE
# =========================================================

initial_choices = get_voice_choices()
initial_voice = get_default_voice_label()

with gr.Blocks(
    title=APP_TITLE,
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown(
        """
        # 🎙️ Myanmar Multilingual TTS Studio

        မြန်မာအမည်များဖြင့် Voice ရွေးချယ်ပြီး
        MP3 Audio နှင့် SRT Subtitle ကို တစ်ခါတည်း ထုတ်ယူပါ။
        """
    )

    with gr.Row():
        voice_input = gr.Dropdown(
            choices=initial_choices,
            value=initial_voice,
            label="Voice Character",
            interactive=True,
            scale=5,
        )

        refresh_button = gr.Button(
            "🔄 Voice List ပြန်ယူမယ်",
            scale=2,
        )

    voice_status = gr.Markdown(
        "နေတိုးကို Default Voice အဖြစ် ရွေးထားသည်။"
    )

    script_input = gr.Textbox(
        label="Script",
        placeholder="အသံထုတ်လိုသော စာသားကို ဒီနေရာမှာ ထည့်ပါ...",
        lines=12,
        max_lines=30,
    )

    speed_input = gr.Slider(
        minimum=-50,
        maximum=100,
        value=DEFAULT_SPEED,
        step=5,
        label="အသံ Speed (%)",
    )

    gr.Markdown(
        """
        - Default Voice: နေတိုး
        - Default Speed: +30%
        - SRT: Space ခြားထားသော စကားလုံး ၂၀ ဝန်းကျင်
        - Speed: -50% မှ +100%
        """
    )

    generate_button = gr.Button(
        "🎙️ Audio + SRT ထုတ်မယ်",
        variant="primary",
        size="lg",
    )

    status_output = gr.Textbox(
        label="အခြေအနေ",
        interactive=False,
        lines=5,
    )

    audio_output = gr.Audio(
        label="🎧 MP3 Audio",
        type="filepath",
    )

    srt_output = gr.File(
        label="📝 SRT Subtitle Download",
    )

    refresh_button.click(
        fn=refresh_voices,
        inputs=[],
        outputs=[
            voice_input,
            voice_status,
        ],
    )

    generate_button.click(
        fn=generate_audio_and_srt,
        inputs=[
            script_input,
            voice_input,
            speed_input,
        ],
        outputs=[
            audio_output,
            srt_output,
            status_output,
        ],
        api_name="generate",
    )


if __name__ == "__main__":
    demo.queue(
        max_size=10,
        default_concurrency_limit=MAX_CONCURRENT_JOBS,
    ).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )