app.py

import streamlit as st
import os
from utils import VideoUtils, SettingsManager
from stt_module import STTManager
from translator import Translator
import tempfile
from dotenv import load_dotenv

# 加载 .env 文件中的配置
load_dotenv()

# 加载本地保存的设置
saved_settings = SettingsManager.load_settings()

# 取消 Streamlit 上传限制（虽然主要通过命令行配置，但在脚本中提醒用户）
# 实际上 Streamlit 的服务器配置需要在运行命令时指定，或者写在 config.toml 中
# 我们这里先通过 UI 提醒用户

st.set_page_config(page_title="AI 媒体翻译专家", layout="wide")

# 初始化 session_state
if "process_results" not in st.session_state:
    st.session_state.process_results = {}
if "srt_results" not in st.session_state:
    st.session_state.srt_results = {}
if "awaiting_synthesis" not in st.session_state:
    st.session_state.awaiting_synthesis = False
if "processing" not in st.session_state:
    st.session_state.processing = False
if "last_uploaded_file" not in st.session_state:
    st.session_state.last_uploaded_file = None
if "last_uploaded_srt" not in st.session_state:
    st.session_state.last_uploaded_srt = None

# 回调函数：开始处理
def start_processing_callback():
    st.session_state.processing = True
    st.session_state.process_results = {}
    st.session_state.awaiting_synthesis = False

# 回调函数：开始合成
def start_synthesis_callback():
    st.session_state.processing = True

# 回调函数：开始翻译字幕
def start_translation_callback():
    st.session_state.processing = True
    st.session_state.srt_results = {}

st.title("🎬 AI 媒体字幕提取与翻译")

# 侧边栏配置
with st.sidebar:
    st.header("⚙️ 配置参数")
    
    # 从本地保存的设置或环境变量读取默认值
    default_api_key = saved_settings.get("API_KEY", os.getenv("API_KEY", ""))
    default_base_url = saved_settings.get("BASE_URL", os.getenv("BASE_URL", "https://api.siliconflow.cn/v1"))
    default_model = saved_settings.get("MODEL_NAME", os.getenv("MODEL_NAME", "THUDM/glm-4-9b-chat"))
    default_stt_size = saved_settings.get("STT_MODEL_SIZE", os.getenv("STT_MODEL_SIZE", "base"))
    default_lang = saved_settings.get("TARGET_LANG", os.getenv("TARGET_LANG", "中文"))
    default_device = saved_settings.get("DEVICE", os.getenv("DEVICE", "cpu"))

    api_key = st.text_input("API Key", value=default_api_key, type="password", help="输入 API Key", disabled=st.session_state.processing)
    base_url = st.text_input("API Base URL", value=default_base_url, disabled=st.session_state.processing)
    model_name = st.text_input("模型名称", value=default_model, disabled=st.session_state.processing)
    
    st.divider()
    
    # 设备选择
    cuda_available = STTManager.is_cuda_available()
    device_options = ["cpu", "cuda"] if cuda_available else ["cpu"]
    device_index = device_options.index(default_device) if default_device in device_options else 0
    device = st.selectbox(
        "计算设备 (Device)", 
        device_options, 
        index=device_index,
        help="如果有 NVIDIA 显卡且安装了 CUDA，选择 'cuda' 会大幅提升速度。",
        disabled=st.session_state.processing
    )
    
    # 根据设备推荐精度
    default_compute_type = "float16" if device == "cuda" else "int8"
    compute_type = st.selectbox(
        "计算精度 (Compute Type)",
        ["int8", "float16", "int8_float16"],
        index=1 if device == "cuda" else 0,
        help="CPU 推荐 int8，GPU 推荐 float16。",
        disabled=st.session_state.processing
    )

    stt_options = ["tiny", "base", "small", "medium", "large-v3"]
    # 获取已下载模型
    downloaded_models = STTManager.get_downloaded_models()
    
    # 构建选项显示名称
    option_labels = []
    for opt in stt_options:
        label = f"{opt} (已下载)" if opt in downloaded_models else opt
        option_labels.append(label)

    stt_index = stt_options.index(default_stt_size) if default_stt_size in stt_options else 1
    selected_option = st.selectbox(
        "本地 STT 模型大小", 
        option_labels, 
        index=stt_index, 
        help="越大越准，但速度越慢。首次使用未下载的模型时会自动下载。",
        disabled=st.session_state.processing
    )
    # 从选项中提取真实模型名
    model_size = selected_option.split(" ")[0]
    
    lang_options = ["中文", "English", "日本語", "Français"]
    lang_index = lang_options.index(default_lang) if default_lang in lang_options else 0
    target_lang = st.selectbox("目标语言", lang_options, index=lang_index, disabled=st.session_state.processing)

    # 保存当前配置到本地文件的辅助函数
    def save_current_settings():
        SettingsManager.save_settings({
            "API_KEY": api_key,
            "BASE_URL": base_url,
            "MODEL_NAME": model_name,
            "STT_MODEL_SIZE": model_size,
            "TARGET_LANG": target_lang,
            "DEVICE": device
        })

    if st.session_state.processing:
        st.divider()
        if st.button("⏹️ 强制停止任务", type="secondary", use_container_width=True):
            st.session_state.processing = False
            st.session_state.awaiting_synthesis = False
            st.rerun()

# 主界面
tab1, tab2 = st.tabs(["📁 媒体处理", "📜 字幕翻译"])

with tab1:
    st.info("💡 提示：支持视频和音频文件。如果文件非常大，处理会很慢请耐心等待。")
    uploaded_file = st.file_uploader("选择视频或音频文件", type=["mp4", "mkv", "avi", "mov", "mp3", "wav", "m4a", "flac", "aac"], disabled=st.session_state.processing)

    if uploaded_file:
        # 检测文件类型
        video_extensions = [".mp4", ".mkv", ".avi", ".mov"]
        file_ext = os.path.splitext(uploaded_file.name)[1].lower()
        is_video = file_ext in video_extensions
        
        # 如果上传了新文件且与 session_state 中记录的不同，则清除旧结果
        if "last_uploaded_file" not in st.session_state or st.session_state.last_uploaded_file != uploaded_file.name:
            st.session_state.process_results = {}
            st.session_state.awaiting_synthesis = False
            st.session_state.last_uploaded_file = uploaded_file.name

        # 保存上传的文件到临时目录
        tfile = tempfile.NamedTemporaryFile(delete=False, suffix=file_ext)
        tfile.write(uploaded_file.read())
        file_path = tfile.name

        if is_video:
            st.video(file_path)
        else:
            st.audio(file_path)

        if st.button("开始处理", 
                   disabled=not api_key or st.session_state.awaiting_synthesis or st.session_state.processing,
                   on_click=start_processing_callback):
            
            try:
                with st.status("正在处理中...", expanded=True) as status:
                    # 1. 提取/准备音频
                    st.write("🎵 正在准备音频...")
                    audio_path = os.path.splitext(file_path)[0] + '.wav'
                    try:
                        # 无论视频还是音频，都通过 prepare_audio (ffmpeg) 转换为标准格式，确保 STT 兼容性
                        VideoUtils.prepare_audio(file_path, audio_path)
                    except Exception as e:
                        st.error(str(e))
                        status.update(label="处理出错", state="error", expanded=True)
                        st.session_state.processing = False
                        st.stop()

                    # 2. 本地 STT
                    if model_size not in downloaded_models:
                        st.write(f"📥 正在下载 {model_size} 模型，请稍候...")
                    
                    st.write(f"✍️ 正在识别语音 (使用 {model_size} 模型，设备: {device})...")
                    stt_manager = STTManager(model_size=model_size, device=device, compute_type=compute_type)
                    stt_manager.load_model()
                    
                    segments_gen, info = stt_manager.transcribe(audio_path)
                    st.write(f"检测到语言: {info.language} (置信度: {info.language_probability:.2f})")

                    # 增量处理与展示
                    st.write("---")
                    st.write("实时识别与翻译预览：")
                    preview_container = st.empty()
                    all_segments = []
                    all_translated_segments = []
                    
                    translator = Translator(api_key, base_url, model_name)
                    
                    # 用于展示的表格数据
                    display_data = []
                    
                    for segment in segments_gen:
                        # 1. 翻译当前段落
                        trans_text = translator.translate_text(segment.text, target_lang)
                        
                        # 2. 保存原始和翻译后的段落
                        all_segments.append(segment)
                        new_trans_seg = type('Segment', (), {
                            'start': segment.start,
                            'end': segment.end,
                            'text': trans_text
                        })
                        all_translated_segments.append(new_trans_seg)
                        
                        # 3. 更新预览界面
                        time_str = f"{VideoUtils.format_timestamp(segment.start)} -> {VideoUtils.format_timestamp(segment.end)}"
                        display_data.append({
                            "时间轴": time_str,
                            "原文": segment.text,
                            "翻译": trans_text
                        })
                        # 仅显示最后 5 条，避免页面过长，但提供滚动查看全部的可能
                        preview_container.table(display_data[-5:])

                    # 生成原始字幕
                    orig_srt_path = os.path.splitext(file_path)[0] + '_orig.srt'
                    VideoUtils.write_srt(all_segments, orig_srt_path)
                    
                    # 生成翻译字幕
                    trans_srt_path = os.path.splitext(file_path)[0] + '_trans.srt'
                    VideoUtils.write_srt(all_translated_segments, trans_srt_path)

                    status.update(label="处理完成！", state="complete", expanded=False)
                    
                    # 成功处理后自动保存配置
                    save_current_settings()

                    # 保存结果到 session_state
                    st.session_state.process_results = {
                        "orig_srt": orig_srt_path,
                        "trans_srt": trans_srt_path,
                        "output_video": None,
                        "is_video": is_video
                    }

                    # 保存中间结果用于后续合成 (仅针对视频)
                    if is_video:
                        st.session_state.temp_video_path = file_path
                        st.session_state.temp_trans_srt_path = trans_srt_path
                        st.session_state.temp_orig_srt_path = orig_srt_path
                        st.session_state.awaiting_synthesis = True
                    
                    st.session_state.processing = False
                    st.rerun()
            except Exception as e:
                st.session_state.processing = False
                st.error(f"处理过程中发生未预期的错误: {e}")
                st.stop()

        # 4. 嵌入视频 (仅视频且在等待状态)
        if st.session_state.awaiting_synthesis and st.session_state.process_results.get("is_video"):
            st.success("✅ 字幕翻译已完成！")
            col_synth_1, col_synth_2 = st.columns(2)
            with col_synth_1:
                if st.button("🚀 开始合成视频字幕", 
                           type="primary", 
                           disabled=st.session_state.processing,
                           on_click=start_synthesis_callback):
                    try:
                        with st.status("🎬 正在合成视频字幕...", expanded=True) as status:
                            v_path = st.session_state.temp_video_path
                            s_path = st.session_state.temp_trans_srt_path
                            output_video_path = os.path.splitext(v_path)[0] + '_translated.mp4'
                            
                            video_ready = False
                            try:
                                VideoUtils.embed_subtitles(v_path, s_path, output_video_path)
                                video_ready = True
                                st.write("✨ 视频合成成功！")
                            except Exception as e:
                                st.error(f"视频合成失败 (请确保已安装 FFmpeg): {e}")

                            status.update(label="全部处理完成！", state="complete", expanded=False)

                            # 保存最终结果到 session_state
                            st.session_state.process_results["output_video"] = output_video_path if video_ready else None
                            st.session_state.awaiting_synthesis = False
                            st.session_state.processing = False
                            st.rerun()
                    except Exception as e:
                        st.session_state.processing = False
                        st.error(f"合成过程中发生未预期的错误: {e}")
                        st.stop()
            
            with col_synth_2:
                if st.button("📂 仅保存字幕"):
                    st.session_state.awaiting_synthesis = False
                    st.rerun()

    # 结果展示与下载 (移出 button 缩进块，始终根据 session_state 显示)
    if st.session_state.process_results:
        st.divider()
        col_title, col_clear = st.columns([5, 1])
        with col_title:
            st.subheader("🎉 处理结果")
        with col_clear:
            if st.button("🗑️ 清除结果"):
                st.session_state.process_results = {}
                st.session_state.awaiting_synthesis = False
                st.rerun()
        
        col1, col2, col3 = st.columns(3)
        
        results = st.session_state.process_results
        
        if os.path.exists(results.get("orig_srt", "")):
            with col1:
                with open(results["orig_srt"], "rb") as f:
                    st.download_button("⬇️ 下载原始字幕", f, file_name="original.srt", key="dl_orig")
        
        if os.path.exists(results.get("trans_srt", "")):
            with col2:
                with open(results["trans_srt"], "rb") as f:
                    st.download_button("⬇️ 下载翻译字幕", f, file_name="translated.srt", key="dl_trans")
        
        if results.get("output_video") and os.path.exists(results["output_video"]):
            with col3:
                with open(results["output_video"], "rb") as f:
                    st.download_button("⬇️ 下载翻译视频", f, file_name="video_with_subtitles.mp4", key="dl_video")

with tab2:
    st.info("如果你已经有原始语言的字幕文件（SRT），可以在这里直接进行翻译。")
    uploaded_srt = st.file_uploader("上传原始 SRT 字幕", type=["srt"], disabled=st.session_state.processing)
    
    if uploaded_srt:
        # 如果上传了新字幕且与 session_state 中记录的不同，则清除旧结果
        if "last_uploaded_srt" not in st.session_state or st.session_state.last_uploaded_srt != uploaded_srt.name:
            st.session_state.srt_results = {}
            st.session_state.last_uploaded_srt = uploaded_srt.name

        srt_content = uploaded_srt.read().decode("utf-8")
        st.text_area("字幕预览", srt_content, height=200)
        
        if st.button("开始翻译字幕", 
                   disabled=not api_key or st.session_state.processing,
                   on_click=start_translation_callback):
            try:
                with st.spinner("正在翻译字幕..."):
                    segments = VideoUtils.parse_srt(srt_content)
                    translator = Translator(api_key, base_url, model_name)
                    translated_segments = translator.translate_segments(segments, target_lang)
                    
                    # 保存到临时文件
                    temp_trans_srt = tempfile.NamedTemporaryFile(delete=False, suffix='.srt')
                    VideoUtils.write_srt(translated_segments, temp_trans_srt.name)
                    
                    st.session_state.srt_results = {
                        "trans_srt": temp_trans_srt.name
                    }
                    st.success("翻译完成！")
                    
                    # 成功处理后自动保存配置
                    save_current_settings()
                    
                    st.session_state.processing = False
                    st.rerun()
            except Exception as e:
                st.session_state.processing = False
                st.error(f"翻译过程中发生错误: {e}")
                st.stop()

    # 结果下载
    if st.session_state.srt_results:
        st.divider()
        res = st.session_state.srt_results
        if os.path.exists(res.get("trans_srt", "")):
            with open(res["trans_srt"], "rb") as f:
                st.download_button("⬇️ 下载翻译后的字幕", f, file_name="translated_only.srt", key="dl_srt_tab")

if not api_key:
    st.warning("请在左侧边栏输入 API Key 后开始。")