import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from faster_whisper import WhisperModel

# ==========================================
# 1. 初始化大语言模型 (反欺诈分析大脑)
# ==========================================
REPO_ID = "mradermacher/AntiFraud-SFT-GGUF"
MODEL_FILE = "AntiFraud-SFT.Q4_K_M.gguf"

print("正在下载 LLM 模型文件...")
model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE)

print("正在加载大语言模型...")
# 移除之前的 Llava Handler，现在作为纯文本模型加载
llm = Llama(
    model_path=model_path,
    n_ctx=2048,  # 上下文窗口
    n_threads=4  # CPU 线程数
)

# ==========================================
# 2. 初始化语音识别模型 (充当模型的"耳朵")
# ==========================================
print("正在加载 Whisper 语音识别模型...")
# 使用 "base" 模型，体积小巧（约140MB），在免费 CPU 上运行飞快
asr_model = WhisperModel("medium", device="cpu", compute_type="int8")

# ==========================================
# 3. 核心处理逻辑：先听后想
# ==========================================
def predict_fraud(audio_file, text_prompt):
    transcribed_text = ""
    
    # 步骤 A：语音转文字 (ASR)
    if audio_file is not None:
        print("开始提取音频文字...")
        # beam_size 适当调小加快 CPU 推理速度
        segments, info = asr_model.transcribe(audio_file, beam_size=3, language="zh")
        transcribed_text = "".join([segment.text for segment in segments])
    
    # 步骤 B：构建最终给大模型的 Prompt
    if transcribed_text:
        combined_prompt = (
            f"以下是拦截到的通话录音文本：\n「{transcribed_text}」\n\n"
            f"用户指令：{text_prompt if text_prompt else '请分析这段对话是否存在电信诈骗风险，并指出可疑点。'}"
        )
    else:
        combined_prompt = text_prompt if text_prompt else "请分析是否存在电信诈骗风险。"

    # 步骤 C：将拼接好的文本发给大模型进行反诈分析
    messages = [
        {"role": "system", "content": "你是一个专业的反电信诈骗分析专家。你的任务是根据通话文本识别潜在的诈骗风险，并给出客观的分析。"},
        {"role": "user", "content": combined_prompt}
    ]
    
    print("开始进行反诈推理解析...")
    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=512,
        temperature=0.2, 
    )
    
    analysis_result = response["choices"][0]["message"]["content"]
    
    # 返回两个结果：让用户看到转录文字，以及最终的分析报告
    return transcribed_text if transcribed_text else "未检测到音频或上传为空", analysis_result

# ==========================================
# 4. 构建 Gradio 可视化界面
# ==========================================
demo = gr.Interface(
    fn=predict_fraud,
    inputs=[
        gr.Audio(type="filepath", label="1. 上传或录制可疑音频 (WAV/MP3等)"),
        gr.Textbox(label="2. 输入附加提示 (可选)", value="请分析这段对话是否存在电信诈骗。")
    ],
    outputs=[
        gr.Textbox(label="🎧 Whisper 音频识别结果 (校对用)", lines=2),
        gr.Textbox(label="🛡️ 模型反欺诈分析报告", lines=6)
    ],
    title="AntiFraud-SFT 智能反欺诈检测站 (ASR级联版)",
    description="通过 Faster-Whisper 进行极速语音识别，再由 Qwen2-Audio 微调模型进行欺诈逻辑分析。完全兼容免费 CPU 部署。",
    api_name="analyze"
)

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