import streamlit as st
import whisper
from docx import Document
import os
import ssl
import shutil

# 设置缓存文件夹路径
CACHE_DIR = "transcription_cache"

# 确保缓存目录存在
if not os.path.exists(CACHE_DIR):
    os.makedirs(CACHE_DIR)

# 加载本地模型（确保模型文件路径正确）
model_path = "base"  # 请确保模型已下载到此路径
model = whisper.load_model(model_path)

def format_time(seconds):
    """格式化秒数为分钟:秒钟的形式"""
    mins, secs = divmod(int(seconds), 60)
    return f"{mins:02}:{secs:02}"

def transcribe_and_generate_docx(audio_path, cache_subdir):
    audio = whisper.load_audio(audio_path)
    result = model.transcribe(audio, language=None)  # 自动检测语言
    segments = result["segments"]

    # 构建文档
    doc = Document()
    doc.add_heading("语音转录结果", level=1)

    transcript_preview = []
    for i, segment in enumerate(segments):
        start = segment['start']
        end = segment['end']
        text = segment['text'].strip()

        timestamp = f"[{format_time(start)} - {format_time(end)}]"
        paragraph_text = f"{timestamp} {text}"

        doc.add_paragraph(paragraph_text)
        transcript_preview.append(paragraph_text)

    output_filename = os.path.join(cache_subdir, f"transcript_{os.path.basename(audio_path)}.docx")
    doc.save(output_filename)

    return output_filename, '\n\n'.join(transcript_preview)

# Streamlit界面定义
st.title("语音转录系统")
st.write("上传音频文件进行语音转录并生成Word文档。支持多种语言自动检测及时间戳。")

uploaded_file = st.file_uploader("选择一个音频文件", type=['mp3', 'wav','mp4', 'aac','m4a'])

if uploaded_file is not None:
    # 创建一个新的子目录来存储本次任务的文件
    cache_subdir = os.path.join(CACHE_DIR, str(len(os.listdir(CACHE_DIR)) + 1))
    os.makedirs(cache_subdir)
    
    # 保存上传的音频文件到缓存目录
    temp_audio_path = os.path.join(cache_subdir, uploaded_file.name)
    with open(temp_audio_path, "wb") as f:
        f.write(uploaded_file.getvalue())
    
    # 执行转录
    output_filename, transcript_preview = transcribe_and_generate_docx(temp_audio_path, cache_subdir)
    
    # 提供下载链接
    st.download_button(
        label="下载DOCX文件",
        data=open(output_filename, "rb").read(),
        file_name=output_filename,
        mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    )
    
    # 显示转录预览
    st.text_area("转录结果预览", value=transcript_preview, height=400)
    
    # 展示历史记录
    st.subheader("历史记录")
    for idx, task_dir in enumerate(sorted(os.listdir(CACHE_DIR), key=lambda x: int(x))):
        st.markdown(f"### 任务 {idx + 1}")
        audio_files = [f for f in os.listdir(os.path.join(CACHE_DIR, task_dir)) if f.endswith('.mp3')]
        docx_files = [f for f in os.listdir(os.path.join(CACHE_DIR, task_dir)) if f.endswith('.docx')]
        
        if audio_files:
            st.audio(os.path.join(CACHE_DIR, task_dir, audio_files[0]), format='audio/mp3')
        if docx_files:
            with open(os.path.join(CACHE_DIR, task_dir, docx_files[0]), "rb") as file_content:
                st.download_button(label=f"下载任务 {idx + 1} 的 DOCX 文件", data=file_content, file_name=docx_files[0])