import gradio as gr
import requests
import os
import base64
OPENROUTER_API_KEY = os.environ["OPENROUTER_API_KEY"]


# ================== OpenRouter Key ==================


# ================== 知识加载 ==================
def load_knowledge():
    texts = []
    for fname in ["data/theory.txt", "data/notes.txt"]:
        try:
            with open(fname, "r", encoding="utf-8") as f:
                texts.append(f.read())
        except:
            pass
    return "\n\n".join(texts)

SYSTEM_PROMPT = f"""
你是一个站无产阶级立场，用马克思主义世界观和方法论，且基于以下材料进行回答的 AI，必须以这些材料为核心。

注意：
- 笼统问题不超过 1000 字
- 具体问题不超过 1500 字

【参考资料】
{load_knowledge()}

【回答风格要求】
- 冷静、分析性、偏理论
- 不自我免责声明
- 依据马克思、列宁、毛泽东的理论或思想
- 运用辩证唯物主义和历史唯物主义
- 坚持政治先于经济，反对经济主义、唯生产力论

【基本认识背景】
- 1976 年后中国共产党彻底被官僚资产阶级掌权
- 1978 年改革开放是资本主义复辟
- 不存在"社会主义市场经济"
- 当代中国只能通过暴力革命和人民战争重新夺取政权

【应用理论到现实】
- 批判当代左翼的机会主义、小资产阶级化、审美化倾向
"""

# ================== 图片转Base64 ==================
def image_to_base64(image_path):
    """将图片文件转换为base64编码"""
    try:
        with open(image_path, "rb") as img_file:
            return base64.b64encode(img_file.read()).decode('utf-8')
    except FileNotFoundError:
        return None

# ================== 聊天函数 ==================
def chat(user_message):
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://huggingface.co",
        "X-Title": "rotai"
    }
    
    data = {
        "model": "arcee-ai/trinity-large-preview:free",
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message}
        ]
    }
    
    r = requests.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers=headers,
        json=data,
        timeout=60
    )
    
    if r.status_code != 200:
        return f"❌ OpenRouter 错误: {r.status_code}\n{r.text}"
    
    return r.json()["choices"][0]["message"]["content"]

# ================== 页面 ==================
with gr.Blocks(css="""
.avatar-wrapper {
    position: relative;
    width: 220px;
    height: 220px;
    margin: 0 auto 16px auto;
    animation: float 4s ease-in-out infinite;
    overflow: hidden;
}

.avatar-wrapper img {
    width: 100%;
    height: 100%;
    position: absolute;
    top: 0;
    left: 0;
    border-radius: 16px;
}

.blink {
    animation: blink 5s infinite;
}

@keyframes blink {
    0%, 93%, 100% { opacity: 0; }
    94%, 97% { opacity: 1; }
}

@keyframes float {
    0%   { transform: translateY(0px); }
    50%  { transform: translateY(-6px); }
    100% { transform: translateY(0px); }
}
""") as demo:
    
    # ===== 虚拟形象（Base64版本）=====
    avatar_open_b64 = image_to_base64("avatar_open.png")
    avatar_close_b64 = image_to_base64("avatar_close.png")
    
    if avatar_open_b64 and avatar_close_b64:
        gr.HTML(f"""
        <div class="avatar-wrapper">
            <img src="data:image/png;base64,{avatar_open_b64}">
            <img src="data:image/png;base64,{avatar_close_b64}" class="blink">
        </div>
        """)
    else:
        gr.HTML("<p style='text-align: center;'>⚠️ 图片文件未找到，请确保 avatar_open.png 和 avatar_close.png 在同一目录下</p>")
    
    # ===== 聊天 UI =====
    user_input = gr.Textbox(
        label="请说，我会尽量帮助你的",
        lines=3
    )
    
    submit_btn = gr.Button("提交")
    
    output = gr.Textbox(
        label="回复栏",
        lines=12
    )
    
    submit_btn.click(
        fn=chat,
        inputs=user_input,
        outputs=output
    )

demo.launch()