import os
import gradio as gr
from llama_cpp import Llama
from llama_cpp.llama_chat_format import Qwen2VLChatHandler
from huggingface_hub import hf_hub_download
import base64

# 1. 权限设置
token = os.getenv("HF_TOKEN")
model_repo = "edge-physio-ai/rehab_expert_q4"

# 2. 下载模型 (如果报错说找不到文件，请检查文件名是否准确)
print("--- 正在从仓库拉取模型文件 ---")
model_path = hf_hub_download(repo_id=model_repo, filename="model_q4_k_m.gguf", token=token)

# 3. 加载推理引擎
llm = Llama(
    model_path=model_path,
    chat_handler=Qwen2VLChatHandler(),
    n_ctx=1024,
    n_threads=2  # 免费版 CPU 只有 2 核，设为 2 最稳
)

def analyze(image_path):
    if not image_path:
        return "请先上传一张康复动作照片。"
    
    # 编码图片
    with open(image_path, "rb") as f:
        base64_image = base64.b64encode(f.read()).decode("utf-8")
    
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "你是一位专业的康复医学专家。请分析图中患者动作的标准度，并给出改进建议。"},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
            ]
        }
    ]
    
    print("--- 正在生成分析报告 (CPU 推理中) ---")
    response = llm.create_chat_completion(messages=messages, max_tokens=512)
    return response["choices"][0]["message"]["content"]

# 4. Gradio 6.x 界面布局
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏃 具身康复专家 AI (Gradio 6.5)")
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="filepath", label="上传动作图片")
            btn = gr.Button("开始专家评估", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="康复分析报告", lines=10)
    
    btn.click(fn=analyze, inputs=input_img, outputs=output_text)

demo.launch()