import logging
import os
import re
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Any

import gradio as gr
import torch
from transformers import (
    AutoModelForSpeechSeq2Seq,
    AutoTokenizer,
    WhisperFeatureExtractor,
    pipeline,
)


logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")

MODEL_NAME = os.getenv("MODEL_NAME", "nineninesix/kyrgyz-whisper-medium")
ASR_LANGUAGE = os.getenv("ASR_LANGUAGE", "ky").strip().lower()
AUDIO_FILTER_PRESET = os.getenv("AUDIO_FILTER_PRESET", "balanced").strip().lower()
AUDIO_FILTER = os.getenv("AUDIO_FILTER", "").strip()
MAX_DURATION_SECONDS = 60 * 60
NO_SPEECH_TEXT = "Кеп табылган жок."
AUDIO_FILTER_PRESETS = {
    "off": "",
    "balanced": (
        "highpass=f=80,"
        "lowpass=f=7800,"
        "afftdn=nr=10:nf=-25,"
        "dynaudnorm=f=150:g=15:p=0.95:m=8,"
        "acompressor=threshold=-18dB:ratio=2.5:attack=20:release=250,"
        "loudnorm=I=-16:TP=-1.5:LRA=11"
    ),
    "aggressive": (
        "highpass=f=100,"
        "lowpass=f=6500,"
        "afftdn=nr=18:nf=-30:tn=1:gs=12,"
        "dynaudnorm=f=100:g=25:p=0.90:m=12,"
        "acompressor=threshold=-24dB:ratio=4:attack=10:release=200,"
        "loudnorm=I=-16:TP=-1.5:LRA=8"
    ),
}

if ASR_LANGUAGE not in {"ky", "ru", "auto"}:
    logging.warning(
        (
            "ASR_LANGUAGE=%r колдоого алынбайт. "
            "'ky', 'ru' же 'auto' колдонуңуз. Auto режимине өттүм."
        ),
        ASR_LANGUAGE,
    )
    ASR_LANGUAGE = "auto"

# Whisper does not have an official Kyrgyz language token, so ASR_LANGUAGE="ky"
# uses the Kyrgyz fine-tuned model without passing a language kwarg. "auto" may
# work better when Kyrgyz and Russian are mixed in one recording.
logging.info("ASR тили: %s", ASR_LANGUAGE)
logging.info("Аудио фильтр: %s", "custom" if AUDIO_FILTER else AUDIO_FILTER_PRESET)

torch.set_num_threads(min(4, os.cpu_count() or 1))
logging.info("CPU threads: %s", torch.get_num_threads())

logging.info("Модель жүктөлүп жатат: %s", MODEL_NAME)
torch_dtype = torch.float32
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
)
model.to("cpu")
model.eval()

feature_extractor = WhisperFeatureExtractor.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
)
asr_pipeline: Any = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=tokenizer,
    feature_extractor=feature_extractor,
    torch_dtype=torch_dtype,
    device=-1,
    chunk_length_s=20,
    stride_length_s=(4, 2),
)
logging.info("Модель даяр")


CUSTOM_CSS = """
body {
    background:
        radial-gradient(circle at top left, rgba(20, 184, 166, 0.12), transparent 30rem),
        linear-gradient(180deg, #f7fbfb 0%, #eef4f5 100%);
}

.gradio-container {
    max-width: 900px !important;
    margin: 0 auto !important;
    padding: 18px 16px 24px !important;
    font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}

.hero {
    margin: 0 auto 16px;
}

.hero h1 {
    margin: 0 0 10px;
    font-size: clamp(2.35rem, 5vw, 3.6rem);
    line-height: 1;
    letter-spacing: 0;
    color: #102026;
}

.hero p {
    max-width: 760px;
    margin: 0;
    color: #3f5661;
    font-size: clamp(1.08rem, 2vw, 1.22rem);
    line-height: 1.45;
}

.workspace {
    background: rgba(255, 255, 255, 0.86);
    border: 1px solid rgba(127, 151, 160, 0.28);
    border-radius: 16px;
    box-shadow: 0 16px 44px rgba(15, 35, 42, 0.09);
    padding: 16px;
}

.instructions {
    margin: 0 0 16px;
    color: #314852;
    font-size: 1.05rem;
    line-height: 1.45;
}

.instructions strong {
    color: #102026;
}

.upload-panel {
    margin-bottom: 10px;
}

.primary-button button {
    min-height: 52px;
    font-size: 1.08rem !important;
    font-weight: 700 !important;
    border-radius: 12px !important;
}

.status-message {
    margin: 10px 0 12px;
    padding: 12px 14px;
    border-radius: 12px;
    font-weight: 700;
    line-height: 1.4;
}

.status-idle {
    background: #eef8f6;
    color: #155e57;
    border: 1px solid #b8e4dc;
}

.status-loading {
    background: #eef4ff;
    color: #1d4ed8;
    border: 1px solid #bdd3ff;
}

.status-success {
    background: #edfdf3;
    color: #166534;
    border: 1px solid #bbf7d0;
}

.status-error {
    background: #fff1f2;
    color: #be123c;
    border: 1px solid #fecdd3;
}

.transcript-box textarea {
    height: 260px !important;
    min-height: 260px !important;
    max-height: 260px !important;
    overflow-y: auto !important;
    resize: none !important;
    font-size: 1rem !important;
    line-height: 1.55 !important;
    white-space: pre-wrap !important;
}

.action-row {
    margin-top: 8px;
    display: grid !important;
    grid-template-columns: repeat(2, minmax(0, 1fr));
    gap: 10px;
}

.action-button button {
    min-height: 46px;
    width: 100%;
    border-radius: 12px !important;
    font-weight: 700 !important;
}

footer {
    display: none !important;
}
"""


def build_generate_kwargs() -> dict[str, Any]:
    generate_kwargs: dict[str, Any] = {
        "task": "transcribe",
        "num_beams": 1,
        "temperature": 0.0,
        "condition_on_prev_tokens": False,
        "compression_ratio_threshold": 2.4,
        "logprob_threshold": -1.0,
        "no_speech_threshold": 0.6,
    }

    if ASR_LANGUAGE == "ru":
        generate_kwargs["language"] = ASR_LANGUAGE

    return generate_kwargs


def build_audio_filter() -> str:
    if AUDIO_FILTER:
        return AUDIO_FILTER

    if AUDIO_FILTER_PRESET not in AUDIO_FILTER_PRESETS:
        logging.warning(
            (
                "AUDIO_FILTER_PRESET=%r колдоого алынбайт. "
                "'balanced', 'aggressive' же 'off' колдонуңуз. Balanced режимине өттүм."
            ),
            AUDIO_FILTER_PRESET,
        )
        return AUDIO_FILTER_PRESETS["balanced"]

    return AUDIO_FILTER_PRESETS[AUDIO_FILTER_PRESET]


def post_process_transcript(text: str) -> str:
    def normalize_segment(value: str) -> str:
        return re.sub(r"\s+", " ", value).strip().casefold()

    def remove_consecutive_duplicate_sentences(line: str) -> str:
        pieces = re.split(r"(?<=[.!?。！？…])\s+", line)
        deduped: list[str] = []
        previous_normalized = ""

        for piece in pieces:
            sentence = piece.strip()
            if not sentence:
                continue

            normalized = normalize_segment(sentence)
            if normalized == previous_normalized:
                continue

            deduped.append(sentence)
            previous_normalized = normalized

        return " ".join(deduped)

    def token_key(token: str) -> str:
        return re.sub(r"^[^\w]+|[^\w]+$", "", token).casefold()

    def phrase_at(tokens: list[str], start: int, phrase_length: int) -> list[str]:
        return [token_key(token) for token in tokens[start : start + phrase_length]]

    def collapse_extreme_repeated_short_phrases(line: str) -> str:
        tokens = line.split()
        if not tokens:
            return ""

        result: list[str] = []
        index = 0
        thresholds = {1: 6, 2: 5, 3: 4}
        keep_repetitions = {1: 3, 2: 2, 3: 2}

        while index < len(tokens):
            collapsed = False

            for phrase_length in (3, 2, 1):
                if index + phrase_length > len(tokens):
                    continue

                phrase = phrase_at(tokens, index, phrase_length)
                if not all(phrase):
                    continue

                repetitions = 1
                next_index = index + phrase_length
                while (
                    next_index + phrase_length <= len(tokens)
                    and phrase_at(tokens, next_index, phrase_length) == phrase
                ):
                    repetitions += 1
                    next_index += phrase_length

                if repetitions >= thresholds[phrase_length]:
                    kept = keep_repetitions[phrase_length]
                    for _ in range(kept):
                        result.extend(tokens[index : index + phrase_length])
                    index = next_index
                    collapsed = True
                    break

            if not collapsed:
                result.append(tokens[index])
                index += 1

        return " ".join(result)

    cleaned = text.strip()
    if not cleaned:
        return NO_SPEECH_TEXT

    lines: list[str] = []
    previous_normalized = ""

    for raw_line in cleaned.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        line = remove_consecutive_duplicate_sentences(line)
        line = collapse_extreme_repeated_short_phrases(line)
        normalized = normalize_segment(line)

        if not normalized or normalized == previous_normalized:
            continue

        lines.append(line)
        previous_normalized = normalized

    cleaned = "\n".join(lines).strip()
    return cleaned or NO_SPEECH_TEXT


def run_command(command: list[str], tool_name: str) -> subprocess.CompletedProcess[str]:
    try:
        result = subprocess.run(
            command,
            capture_output=True,
            check=False,
            text=True,
        )
    except FileNotFoundError as exc:
        logging.exception("%s табылган жок. Command: %s", tool_name, command)
        raise RuntimeError(f"{tool_name} орнотулган эмес. Сервер конфигурациясын текшериңиз.") from exc

    if result.returncode != 0:
        details = (result.stderr or result.stdout or "").strip()
        logging.error(
            "%s катасы. Return code: %s. Command: %s. Details: %s",
            tool_name,
            result.returncode,
            command,
            details,
        )
        raise RuntimeError(
            f"{tool_name} файлды иштете алган жок. Файл форматын текшерип, кайра аракет кылыңыз."
        )

    return result


def media_duration_seconds(input_path: Path) -> float:
    result = run_command(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=noprint_wrappers=1:nokey=1",
            str(input_path),
        ],
        "ffprobe",
    )

    try:
        return float(result.stdout.strip())
    except ValueError as exc:
        logging.error("ffprobe duration output окулбай калды: %r", result.stdout)
        raise RuntimeError("Файлдын узактыгын окуй алган жокмун.") from exc


def extract_audio(input_path: Path, output_path: Path) -> None:
    command = [
        "ffmpeg",
        "-y",
        "-i",
        str(input_path),
        "-vn",
    ]

    audio_filter = build_audio_filter()
    if audio_filter:
        command.extend(["-af", audio_filter])

    command.extend(
        [
            "-ac",
            "1",
            "-ar",
            "16000",
            str(output_path),
        ]
    )

    run_command(command, "ffmpeg")


def transcribe_audio(audio_path: Path) -> str:
    with torch.inference_mode():
        result = asr_pipeline(
            str(audio_path),
            return_timestamps=False,
            generate_kwargs=build_generate_kwargs(),
        )

    if not isinstance(result, dict):
        raise RuntimeError("Модель күтүлбөгөн жооп кайтарды.")

    return post_process_transcript(str(result.get("text", "")))


def write_transcript_file(text: str) -> str:
    transcript_file = tempfile.NamedTemporaryFile(
        mode="w",
        encoding="utf-8",
        suffix=".txt",
        prefix="transcript-",
        delete=False,
    )

    with transcript_file:
        transcript_file.write(text)

    return transcript_file.name


def status_html(message: str, status_type: str = "idle") -> str:
    return f'<div class="status-message status-{status_type}">{message}</div>'


def uploaded_file_path(uploaded_file: Any) -> Path:
    if isinstance(uploaded_file, (str, Path)):
        return Path(uploaded_file)

    if hasattr(uploaded_file, "name"):
        return Path(uploaded_file.name)

    if isinstance(uploaded_file, dict):
        for key in ("path", "name"):
            if uploaded_file.get(key):
                return Path(uploaded_file[key])

    raise RuntimeError("Жүктөлгөн файлдын жолун окуй алган жокмун.")


def transcribe_file(uploaded_file: Any | None) -> tuple[str, str, str | None]:
    if not uploaded_file:
        return status_html("Файлды тандаңыз.", "error"), "", None

    try:
        with tempfile.TemporaryDirectory(prefix="synchy-") as temp_dir:
            temp_path = Path(temp_dir)
            input_path = uploaded_file_path(uploaded_file)
            source_path = temp_path / input_path.name
            audio_path = temp_path / "audio.wav"

            shutil.copy(input_path, source_path)

            duration = media_duration_seconds(source_path)
            if duration > MAX_DURATION_SECONDS:
                return (
                    status_html("Ката: файл 1 сааттан узун. Кыскараак файл жүктөңүз.", "error"),
                    "",
                    None,
                )

            extract_audio(source_path, audio_path)
            transcript = transcribe_audio(audio_path)
            transcript_path = write_transcript_file(transcript)

        return status_html("Даяр. Текст төмөндө көрсөтүлдү.", "success"), transcript, transcript_path

    except RuntimeError as exc:
        logging.exception("Иштетүү катасы: %s", exc)
        return status_html(f"Ката: {exc}", "error"), "", None
    except Exception as exc:
        logging.exception("Күтүлбөгөн ката: %s", exc)
        return status_html(f"Ката: {exc}", "error"), "", None


def loading_status() -> str:
    return status_html("Иштетилип жатат... CPU режиминде бул бир аз убакыт алышы мүмкүн.", "loading")


def transcript_file_for_download(transcript_path: str | None) -> str | None:
    return transcript_path


with gr.Blocks(title="Synchy", css=CUSTOM_CSS) as demo:
    gr.Markdown(
        """
        """
    )

    with gr.Group(elem_classes=["workspace"]):
        gr.Markdown(
            """
            <p class="instructions"><strong>Кантип колдонулат:</strong> файлды тандаңыз, андан кийин баскычты басыңыз. Текст даяр болгондо аны көчүрүп же <code>transcript.txt</code> файл катары жүктөп алсаңыз болот.</p>
            """
        )

        file_input = gr.File(
            label="Видео же аудио файл",
            file_types=[
                ".aac",
                ".flac",
                ".m4a",
                ".mp3",
                ".mp4",
                ".mov",
                ".ogg",
                ".opus",
                ".wav",
                ".webm",
            ],
            type="filepath",
            elem_classes=["upload-panel"],
        )
        transcribe_button = gr.Button(
            "Текстке айландыруу",
            variant="primary",
            elem_classes=["primary-button"],
        )
        status_output = gr.HTML(
            value=status_html("Файл жүктөп алууга даяр.", "idle"),
        )

        transcript_output = gr.Textbox(
            label="Транскрипция",
            placeholder="Текст ушул жерге чыгат.",
            lines=8,
            interactive=False,
            elem_classes=["transcript-box"],
        )

        with gr.Row(elem_classes=["action-row"]):
            copy_button = gr.Button("Көчүрүү", elem_classes=["action-button"])
            download_button = gr.DownloadButton(
                "Жүктөп алуу",
                value=None,
                elem_classes=["action-button"],
            )
            transcript_file_state = gr.State(value=None)

    transcribe_button.click(
        fn=loading_status,
        inputs=None,
        outputs=status_output,
        show_progress="hidden",
    ).then(
        fn=transcribe_file,
        inputs=file_input,
        outputs=[status_output, transcript_output, transcript_file_state],
    )

    transcript_file_state.change(
        fn=transcript_file_for_download,
        inputs=transcript_file_state,
        outputs=download_button,
    )

    copy_button.click(
        fn=None,
        inputs=transcript_output,
        outputs=None,
        js="""
        (text) => {
            if (text) {
                navigator.clipboard.writeText(text);
            }
            return [];
        }
        """,
    )


if __name__ == "__main__":
    demo.queue().launch()