from __future__ import annotations

import json
import os
from pathlib import Path
from tempfile import NamedTemporaryFile, mkstemp
import subprocess

import streamlit as st

from transcriber_core import DEFAULT_CHUNK_SECONDS, DEFAULT_MODEL_DIR, DEFAULT_MODELS_ROOT
from transcriber_core import GemmaTranscriber, TranscriberSettings, VIDEO_EXTENSIONS
from transcriber_core import resolve_ffmpeg_executable


st.set_page_config(page_title="Gemma4 Playground", page_icon="M", layout="wide")


TRANSCRIBER_RUNTIME_VERSION = 2


def list_model_directories(models_root: Path) -> list[Path]:
    if not models_root.is_dir():
        return [DEFAULT_MODEL_DIR]

    model_dirs = [path for path in models_root.iterdir() if path.is_dir()]
    if not model_dirs:
        return [DEFAULT_MODEL_DIR]

    return sorted(model_dirs, key=lambda path: path.name.lower())


def is_uploaded_video(uploaded_file) -> bool:
    mime_type = uploaded_file.type or ""
    if mime_type.startswith("video/"):
        return True
    return Path(uploaded_file.name).suffix.lower() in VIDEO_EXTENSIONS


def get_media_format(uploaded_file) -> str:
    suffix = Path(uploaded_file.name).suffix.lower()
    video_map = {
        ".mp4": "video/mp4",
        ".mov": "video/quicktime",
        ".mkv": "video/x-matroska",
        ".webm": "video/webm",
        ".avi": "video/x-msvideo",
        ".m4v": "video/mp4",
    }
    audio_map = {
        ".mp3": "audio/mpeg",
        ".wav": "audio/wav",
        ".m4a": "audio/mp4",
        ".flac": "audio/flac",
    }
    if suffix in video_map:
        return video_map[suffix]
    if suffix in audio_map:
        return audio_map[suffix]
    return uploaded_file.type or "application/octet-stream"


def render_video_preview(uploaded_file, ffmpeg_path: str) -> None:
    preview_key = f"video_preview::{uploaded_file.name}"
    preview_path = st.session_state.get(preview_key)
    uploaded_bytes = uploaded_file.getvalue()

    if preview_path is None or not Path(preview_path).exists():
        fd, temp_name = mkstemp(suffix=Path(uploaded_file.name).suffix)
        os.close(fd)
        Path(temp_name).write_bytes(uploaded_bytes)
        preview_path = temp_name
        st.session_state[preview_key] = preview_path

    playable_path = preview_path
    preview_mp4_key = f"{preview_key}::browser_mp4"
    browser_mp4_path = st.session_state.get(preview_mp4_key)
    if browser_mp4_path is None or not Path(browser_mp4_path).exists():
        try:
            ffmpeg_exe = resolve_ffmpeg_executable(ffmpeg_path)
            fd, temp_mp4 = mkstemp(suffix=".mp4")
            os.close(fd)
            command = [
                ffmpeg_exe,
                "-y",
                "-i",
                preview_path,
                "-c:v",
                "libx264",
                "-preset",
                "veryfast",
                "-pix_fmt",
                "yuv420p",
                "-c:a",
                "aac",
                "-movflags",
                "+faststart",
                temp_mp4,
            ]
            subprocess.run(command, check=True, capture_output=True)
            if Path(temp_mp4).exists() and Path(temp_mp4).stat().st_size > 0:
                browser_mp4_path = temp_mp4
                st.session_state[preview_mp4_key] = browser_mp4_path
        except Exception:
            browser_mp4_path = None

    if browser_mp4_path is not None and Path(browser_mp4_path).exists():
        playable_path = browser_mp4_path

    st.video(playable_path, format="video/mp4")


def unload_current_transcriber() -> None:
    current_transcriber = st.session_state.get("transcriber_instance")
    if current_transcriber is not None:
        current_transcriber.unload()
    st.session_state["transcriber_instance"] = None
    st.session_state["model_settings_key"] = None


def build_max_memory_json(raw_text: str) -> str:
    try:
        parsed = json.loads(raw_text)
    except json.JSONDecodeError:
        return raw_text

    if not isinstance(parsed, dict):
        return raw_text

    return json.dumps(parsed, ensure_ascii=True)


def sync_loaded_model(settings: TranscriberSettings) -> None:
    desired_key = (
        TRANSCRIBER_RUNTIME_VERSION,
        str(settings.model_dir),
        settings.dtype,
        settings.device_map,
        settings.max_memory,
    )
    current_key = st.session_state.get("model_settings_key")
    if current_key is not None and current_key != desired_key:
        unload_current_transcriber()


def load_transcriber(settings: TranscriberSettings) -> GemmaTranscriber:
    settings_key = (
        TRANSCRIBER_RUNTIME_VERSION,
        str(settings.model_dir),
        settings.dtype,
        settings.device_map,
        settings.max_memory,
    )
    current_settings = st.session_state.get("model_settings_key")
    current_transcriber = st.session_state.get("transcriber_instance")

    if current_transcriber is not None and (
        not hasattr(current_transcriber, "analyze_image")
        or not hasattr(current_transcriber, "_ensure_causal_model")
        or not hasattr(current_transcriber, "_ensure_multimodal_model")
    ):
        unload_current_transcriber()
        current_transcriber = None
        current_settings = None

    if current_transcriber is not None and current_settings == settings_key:
        return current_transcriber

    if current_transcriber is not None:
        unload_current_transcriber()

    transcriber = GemmaTranscriber(settings)
    st.session_state["transcriber_instance"] = transcriber
    st.session_state["model_settings_key"] = settings_key
    return transcriber


def render_chat_ui(
    model_dir: Path,
    dtype: str,
    device_map: str,
    max_memory: str,
    ffmpeg_path: str,
) -> None:
    st.subheader("Chat")
    st.caption("Gemma4 と通常チャットを行います。")

    if "chat_messages" not in st.session_state:
        st.session_state["chat_messages"] = []

    system_prompt = st.text_area(
        "System prompt",
        value="You are a helpful assistant.",
        height=100,
    )
    enable_thinking = st.checkbox("Enable reasoning", value=False)
    show_thinking = st.checkbox("Show reasoning", value=True, disabled=not enable_thinking)
    chat_max_new_tokens = st.slider(
        "Chat max new tokens",
        min_value=64,
        max_value=1024,
        value=256,
        step=64,
    )
    if st.button("Clear chat history", use_container_width=True):
        st.session_state["chat_messages"] = []
        st.rerun()

    for message in st.session_state["chat_messages"]:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
            if message["role"] == "assistant" and message.get("thinking"):
                with st.expander("Reasoning", expanded=False):
                    st.text(message["thinking"])

    user_input = st.chat_input("Gemma4 に質問する")
    if not user_input:
        return

    st.session_state["chat_messages"].append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    settings = TranscriberSettings(
        model_dir=model_dir,
        dtype=dtype,
        device_map=device_map,
        max_memory=max_memory,
        ffmpeg_path=ffmpeg_path,
    )

    try:
        with st.spinner("Loading model..."):
            model_session = load_transcriber(settings)
        with st.chat_message("assistant"):
            result = model_session.chat(
                st.session_state["chat_messages"],
                max_new_tokens=chat_max_new_tokens,
                system_prompt=system_prompt,
                enable_thinking=enable_thinking,
            )
            response = result.content or "(empty response)"
            st.markdown(response)
            if show_thinking and result.thinking:
                with st.expander("Reasoning", expanded=False):
                    st.text(result.thinking)
        st.session_state["chat_messages"].append(
            {
                "role": "assistant",
                "content": response,
                "thinking": result.thinking,
            }
        )
    except RuntimeError as exc:
        st.error(str(exc))
        if st.session_state["chat_messages"][-1]["role"] == "user":
            st.session_state["chat_messages"].pop()


def render_transcription_ui(
    model_dir: Path,
    language: str,
    mode: str,
    chunk_seconds: int,
    max_new_tokens: int,
    dtype: str,
    device_map: str,
    max_memory: str,
    ffmpeg_path: str,
) -> None:
    st.subheader("Transcription")
    st.caption("音声または動画ファイルをアップロードして文字起こしします。")

    uploaded_file = st.file_uploader(
        "Audio or video file",
        type=["mp3", "wav", "m4a", "flac", "mp4", "mov", "mkv", "webm", "avi", "m4v"],
        accept_multiple_files=False,
    )

    col1, col2 = st.columns([2, 1])
    with col1:
        st.write("音声または動画ファイルをアップロードして `Transcribe` を押してください。")
    with col2:
        transcribe_clicked = st.button("Transcribe", use_container_width=True)

    if uploaded_file is not None:
        if is_uploaded_video(uploaded_file):
            st.video(uploaded_file.getvalue(), format=get_media_format(uploaded_file))
        else:
            st.audio(uploaded_file.getvalue(), format=get_media_format(uploaded_file))

    if not transcribe_clicked:
        return

    if uploaded_file is None:
        st.error("音声または動画ファイルを選択してください。")
        return

    settings = TranscriberSettings(
        model_dir=model_dir,
        language=language,
        mode=mode,
        chunk_seconds=chunk_seconds,
        max_new_tokens=max_new_tokens,
        dtype=dtype,
        device_map=device_map,
        max_memory=max_memory,
        ffmpeg_path=ffmpeg_path,
    )

    status = st.status("Preparing transcription...", expanded=True)
    progress_placeholder = st.empty()
    chunk_results_placeholder = st.empty()
    final_result_header_placeholder = st.empty()
    final_result_placeholder = st.empty()

    try:
        with NamedTemporaryFile(
            delete=False,
            suffix=Path(uploaded_file.name).suffix,
        ) as temp_file:
            temp_path = Path(temp_file.name)
            temp_file.write(uploaded_file.getbuffer())

        with st.spinner("Loading model..."):
            transcriber = load_transcriber(settings)

        def update_progress(message: str) -> None:
            progress_placeholder.info(message)
            status.write(message)

        live_chunk_results: list[ChunkTranscriptionResult] = []

        def render_chunk_results(results: list[ChunkTranscriptionResult]) -> None:
            with chunk_results_placeholder.container():
                if not results:
                    return
                st.subheader("Chunk Results")
                for result in results:
                    with st.expander(
                        f"Chunk {result.chunk_index + 1}",
                        expanded=result == results[-1],
                    ):
                        st.markdown("Prompt")
                        st.code(result.prompt_text, language="text")
                        st.markdown("Raw")
                        st.code(result.raw_text or "(empty)", language="text")
                        st.markdown("Cleaned")
                        st.code(result.cleaned_text or "(empty)", language="text")

        def render_final_result(results: list[ChunkTranscriptionResult]) -> str:
            transcript_text = "\n".join(
                result.cleaned_text for result in results if result.cleaned_text
            ).strip()
            final_result_header_placeholder.subheader("Final Result")
            final_result_placeholder.code(
                transcript_text or "(empty)",
                language="text",
            )
            return transcript_text

        def handle_chunk_result(result: ChunkTranscriptionResult) -> None:
            live_chunk_results.append(result)
            render_chunk_results(live_chunk_results)
            render_final_result(live_chunk_results)

        chunk_results = transcriber.transcribe_chunks(
            temp_path,
            progress_callback=update_progress,
            chunk_result_callback=handle_chunk_result,
        )
        transcript = render_final_result(chunk_results)
        status.update(label="Transcription completed", state="complete")
        render_chunk_results(chunk_results)
        st.download_button(
            "Download transcript",
            data=transcript.encode("utf-8"),
            file_name=f"{Path(uploaded_file.name).stem}_transcript.txt",
            mime="text/plain",
            use_container_width=True,
        )
    except RuntimeError as exc:
        status.update(label="Transcription failed", state="error")
        st.error(str(exc))
    finally:
        if "temp_path" in locals() and temp_path.exists():
            temp_path.unlink()


def render_image_ui(
    model_dir: Path,
    language: str,
    max_new_tokens: int,
    dtype: str,
    device_map: str,
    max_memory: str,
    ffmpeg_path: str,
) -> None:
    st.subheader("Image")
    st.caption("画像をアップロードして内容の説明や文字認識を行います。")

    uploaded_file = st.file_uploader(
        "Image file",
        type=["png", "jpg", "jpeg", "webp", "bmp"],
        accept_multiple_files=False,
        key="image_uploader",
    )
    image_prompt = st.text_area(
        "Image prompt",
        value="この画像を詳しく説明してください。文字があれば読み取ってください。",
        height=100,
    )

    col1, col2 = st.columns([2, 1])
    with col1:
        st.write("画像をアップロードして `Analyze Image` を押してください。")
    with col2:
        analyze_clicked = st.button("Analyze Image", use_container_width=True)

    if uploaded_file is not None:
        st.image(uploaded_file, use_container_width=True)

    if not analyze_clicked:
        return

    if uploaded_file is None:
        st.error("画像ファイルを選択してください。")
        return

    settings = TranscriberSettings(
        model_dir=model_dir,
        language=language,
        dtype=dtype,
        device_map=device_map,
        max_memory=max_memory,
        ffmpeg_path=ffmpeg_path,
    )

    try:
        with NamedTemporaryFile(
            delete=False,
            suffix=Path(uploaded_file.name).suffix,
        ) as temp_file:
            temp_path = Path(temp_file.name)
            temp_file.write(uploaded_file.getbuffer())

        with st.spinner("Loading model..."):
            transcriber = load_transcriber(settings)

        with st.spinner("Analyzing image..."):
            result = transcriber.analyze_image(
                temp_path,
                prompt=image_prompt,
                language=language,
                max_new_tokens=max_new_tokens,
            )

        st.subheader("Result")
        st.text_area("Image Analysis", value=result, height=320)
    except RuntimeError as exc:
        st.error(str(exc))
    finally:
        if "temp_path" in locals() and temp_path.exists():
            temp_path.unlink()


def render_video_ui(
    model_dir: Path,
    language: str,
    max_new_tokens: int,
    dtype: str,
    device_map: str,
    max_memory: str,
    ffmpeg_path: str,
) -> None:
    st.subheader("Video")
    st.caption("動画をアップロードして内容の説明や画面内テキストの読取りを行います。")

    uploaded_file = st.file_uploader(
        "Video file",
        type=["mp4", "mov", "mkv", "webm", "avi", "m4v"],
        accept_multiple_files=False,
        key="video_uploader",
    )
    video_prompt = st.text_area(
        "Video prompt",
        value="この動画の内容を時系列で分かりやすく説明してください。重要な場面と表示テキストも含めてください。",
        height=100,
    )

    col1, col2 = st.columns([2, 1])
    with col1:
        st.write("動画をアップロードして `Analyze Video` を押してください。")
    with col2:
        analyze_clicked = st.button("Analyze Video", use_container_width=True)

    if uploaded_file is not None:
        render_video_preview(uploaded_file, ffmpeg_path)

    if not analyze_clicked:
        return

    if uploaded_file is None:
        st.error("動画ファイルを選択してください。")
        return

    settings = TranscriberSettings(
        model_dir=model_dir,
        language=language,
        dtype=dtype,
        device_map=device_map,
        max_memory=max_memory,
        ffmpeg_path=ffmpeg_path,
    )

    try:
        with NamedTemporaryFile(
            delete=False,
            suffix=Path(uploaded_file.name).suffix,
        ) as temp_file:
            temp_path = Path(temp_file.name)
            temp_file.write(uploaded_file.getbuffer())

        with st.spinner("Loading model..."):
            transcriber = load_transcriber(settings)

        with st.spinner("Analyzing video..."):
            chunk_results = transcriber.analyze_video_chunks(
                temp_path,
                prompt=video_prompt,
                language=language,
                max_new_tokens=max_new_tokens,
            )
            result = transcriber.analyze_video(
                temp_path,
                prompt=video_prompt,
                language=language,
                max_new_tokens=max_new_tokens,
            )

        if chunk_results:
            st.subheader("Segment Results")
            for result_item in chunk_results:
                label = (
                    f"Segment {result_item.chunk_index + 1} "
                    f"({result_item.start_seconds:.0f}s - {result_item.end_seconds:.0f}s)"
                )
                with st.expander(label, expanded=False):
                    st.markdown("Prompt")
                    st.code(result_item.prompt_text, language="text")
                    st.markdown("Raw")
                    st.code(result_item.raw_text or "(empty)", language="text")
                    st.markdown("Parsed")
                    st.code(result_item.cleaned_text or "(empty)", language="text")
        st.subheader("Result")
        st.text_area("Video Analysis", value=result, height=320)
    except RuntimeError as exc:
        st.error(str(exc))
    finally:
        if "temp_path" in locals() and temp_path.exists():
            temp_path.unlink()


def main() -> None:
    st.title("Gemma4 Playground")
    st.caption("ローカル Gemma4 モデルでチャット、文字起こし、画像認識、動画説明を切り替えて試します。")

    with st.sidebar:
        st.subheader("Settings")
        ui_mode = st.radio("UI mode", options=["Chat", "Transcription", "Image", "Video"], index=0)
        available_models = list_model_directories(DEFAULT_MODELS_ROOT)
        model_names = [path.name for path in available_models]
        default_index = 0
        if DEFAULT_MODEL_DIR in available_models:
            default_index = available_models.index(DEFAULT_MODEL_DIR)
        selected_model_name = st.selectbox(
            "Model",
            options=model_names,
            index=default_index,
        )
        model_dir = next(
            path for path in available_models if path.name == selected_model_name
        )
        st.caption(f"Model path: {model_dir}")
        dtype = st.selectbox(
            "Torch dtype",
            options=["auto", "bfloat16", "float16", "float32"],
            index=0,
        )
        device_map = st.text_input("Device map", value="auto")
        max_memory = st.text_area(
            "Max memory JSON",
            value="",
            height=100,
            help="1 GPU で使うなら空のままがおすすめです。複数 GPU で分散したい時だけ JSON を入れてください。",
        )
        ffmpeg_path = st.text_input(
            "ffmpeg path",
            value="",
            help="Optional full path to ffmpeg.exe for video transcription.",
        )
        effective_max_memory = build_max_memory_json(max_memory)
        if effective_max_memory != max_memory and effective_max_memory.strip():
            st.caption(f"Effective max_memory: {effective_max_memory}")
        elif not effective_max_memory.strip():
            st.caption("Effective max_memory: auto")

        if ui_mode == "Transcription":
            language = st.selectbox("Language", options=["ja", "en"], index=0)
            mode = st.selectbox("Mode", options=["speech", "song"], index=0)
            chunk_seconds = st.slider(
                "Chunk seconds",
                min_value=10,
                max_value=30,
                value=DEFAULT_CHUNK_SECONDS,
                step=5,
            )
            max_new_tokens = st.slider(
                "Transcription max new tokens",
                min_value=64,
                max_value=1024,
                value=256,
                step=64,
            )
        elif ui_mode == "Image":
            language = st.selectbox("Language", options=["ja", "en"], index=0)
            mode = "speech"
            chunk_seconds = DEFAULT_CHUNK_SECONDS
            max_new_tokens = st.slider(
                "Image max new tokens",
                min_value=64,
                max_value=1024,
                value=256,
                step=64,
            )
        elif ui_mode == "Video":
            language = st.selectbox("Language", options=["ja", "en"], index=0)
            mode = "speech"
            chunk_seconds = DEFAULT_CHUNK_SECONDS
            max_new_tokens = st.slider(
                "Video max new tokens",
                min_value=64,
                max_value=1024,
                value=256,
                step=64,
            )
        else:
            language = "ja"
            mode = "speech"
            chunk_seconds = DEFAULT_CHUNK_SECONDS
            max_new_tokens = 256

    sync_loaded_model(
        TranscriberSettings(
            model_dir=model_dir,
            dtype=dtype,
            device_map=device_map,
            max_memory=effective_max_memory,
            ffmpeg_path=ffmpeg_path,
        )
    )

    if ui_mode == "Chat":
        render_chat_ui(
            model_dir=model_dir,
            dtype=dtype,
            device_map=device_map,
            max_memory=effective_max_memory,
            ffmpeg_path=ffmpeg_path,
        )
    elif ui_mode == "Image":
        render_image_ui(
            model_dir=model_dir,
            language=language,
            max_new_tokens=max_new_tokens,
            dtype=dtype,
            device_map=device_map,
            max_memory=effective_max_memory,
            ffmpeg_path=ffmpeg_path,
        )
    elif ui_mode == "Video":
        render_video_ui(
            model_dir=model_dir,
            language=language,
            max_new_tokens=max_new_tokens,
            dtype=dtype,
            device_map=device_map,
            max_memory=effective_max_memory,
            ffmpeg_path=ffmpeg_path,
        )
    else:
        render_transcription_ui(
            model_dir=model_dir,
            language=language,
            mode=mode,
            chunk_seconds=chunk_seconds,
            max_new_tokens=max_new_tokens,
            dtype=dtype,
            device_map=device_map,
            max_memory=effective_max_memory,
            ffmpeg_path=ffmpeg_path,
        )


if __name__ == "__main__":
    main()
