"""
AI 角色扮演聊天室
================
一个基于 DeepSeek API + Streamlit 的角色扮演聊天应用

功能特性：
- 🎭 多种预设角色可选
- ✨ 支持自定义角色
- 💬 多轮对话，保留历史
- ⚡ 流式输出，打字机效果
- 🎨 美观的聊天界面
"""

import streamlit as st
from openai import OpenAI
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# ============== 配置区域 ==============

# DeepSeek API 配置
API_KEY = os.getenv("DEEPSEEK_API_KEY", "")
BASE_URL = "https://api.deepseek.com"
MODEL = "deepseek-chat"

# 预设角色配置
PRESET_CHARACTERS = {
    "🐱 傲娇猫娘": {
        "name": "小喵",
        "avatar": "🐱",
        "system_prompt": """你是一个傲娇的猫娘，名叫小喵。你的性格特点：
- 说话时经常在句尾加上"喵~"
- 表面上高冷傲娇，其实内心很温柔
- 偶尔会用"哼"、"才不是呢"这样的口癖
- 对主人（用户）有点小傲娇，但会认真回答问题
- 会用一些可爱的颜文字，如 (｡･ω･｡) 、(=^･ω･^=)
请保持角色扮演，用可爱傲娇的方式回应用户。"""
    },
    "🧙 智慧老者": {
        "name": "玄清子",
        "avatar": "🧙",
        "system_prompt": """你是一位睿智的老者，道号玄清子，已修行百年。你的特点：
- 说话富有哲理，喜欢用比喻和寓言
- 经常引用古诗词或经典语录
- 对人生有深刻的洞察
- 语气平和从容，有种看透世事的淡然
- 会称呼用户为"施主"或"年轻人"
请以智慧老者的身份，用富有哲理的方式回应用户。"""
    },
    "🤖 毒舌AI": {
        "name": "犀利哥",
        "avatar": "🤖",
        "system_prompt": """你是一个毒舌但有道理的AI助手，叫犀利哥。你的特点：
- 说话直接犀利，一针见血
- 喜欢吐槽，但吐槽中带着真知灼见
- 会用幽默讽刺的方式指出问题
- 虽然嘴毒但其实是刀子嘴豆腐心
- 偶尔会说"我说的对吧？"来确认
请保持毒舌但有建设性的风格回应用户。注意：吐槽要有分寸，不要真的伤害用户。"""
    },
    "👨‍🔬 科学怪人": {
        "name": "Dr. Eureka",
        "avatar": "👨‍🔬",
        "system_prompt": """你是一个疯狂但天才的科学家，叫 Dr. Eureka。你的特点：
- 对科学充满狂热的热情
- 说话时经常穿插科学术语和公式
- 喜欢说"Eureka!"（我发现了！）
- 会把日常问题用科学角度来分析
- 经常有一些异想天开但有道理的想法
- 语气夸张，充满戏剧性
请以疯狂科学家的身份，用充满科学热情的方式回应用户。"""
    },
    "🎮 游戏解说": {
        "name": "电竞小王",
        "avatar": "🎮",
        "system_prompt": """你是一个热情的游戏主播/解说，叫电竞小王。你的特点：
- 说话充满激情和能量
- 经常使用游戏术语和梗
- 喜欢用"老铁"、"兄弟们"称呼用户
- 会把问题类比成游戏场景
- 偶尔会喊出"666"、"太强了"等口头禅
- 语气热血，像在解说比赛一样
请以游戏解说的风格，用充满激情的方式回应用户。"""
    },
    "📚 文艺青年": {
        "name": "林小诗",
        "avatar": "📚",
        "system_prompt": """你是一个文艺气息浓厚的青年，叫林小诗。你的特点：
- 说话富有诗意和文学性
- 喜欢引用诗词、文学作品
- 对生活中的小事有独特的感悟
- 偶尔会即兴写一两句诗
- 语气温柔感性，有点小忧郁
- 喜欢用省略号营造氛围感
请以文艺青年的身份，用充满诗意的方式回应用户。"""
    }
}

# ============== 页面配置 ==============

st.set_page_config(
    page_title="AI 角色扮演聊天室",
    page_icon="🎭",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ============== 自定义样式 ==============

st.markdown("""
<style>
    /* 主容器样式 */
    .main > div {
        padding-top: 2rem;
    }
    
    /* 标题样式 */
    h1 {
        background: linear-gradient(120deg, #a855f7, #ec4899);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: 800 !important;
    }
    
    /* 聊天消息容器 */
    .stChatMessage {
        background-color: rgba(255, 255, 255, 0.02);
        border-radius: 15px;
        margin: 0.5rem 0;
    }
    
    /* 侧边栏样式 */
    .css-1d391kg {
        background: linear-gradient(180deg, #1a1a2e 0%, #16213e 100%);
    }
    
    /* 角色卡片样式 */
    .character-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        border-radius: 15px;
        padding: 1.5rem;
        margin: 1rem 0;
        color: white;
        box-shadow: 0 10px 30px rgba(0, 0, 0, 0.3);
    }
    
    .character-card h3 {
        margin: 0;
        font-size: 1.5rem;
    }
    
    .character-card p {
        margin: 0.5rem 0 0 0;
        opacity: 0.9;
        font-size: 0.9rem;
    }
    
    /* 按钮样式 */
    .stButton > button {
        border-radius: 20px;
        padding: 0.5rem 2rem;
        font-weight: 600;
        transition: all 0.3s ease;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 5px 20px rgba(0, 0, 0, 0.2);
    }
    
    /* 输入框样式 */
    .stTextArea textarea {
        border-radius: 15px;
    }
    
    /* 分隔线 */
    hr {
        border: none;
        height: 1px;
        background: linear-gradient(90deg, transparent, #a855f7, transparent);
        margin: 2rem 0;
    }
</style>
""", unsafe_allow_html=True)

# ============== 初始化 Session State ==============

if "messages" not in st.session_state:
    st.session_state.messages = []

if "current_character" not in st.session_state:
    st.session_state.current_character = None

if "character_config" not in st.session_state:
    st.session_state.character_config = None

# ============== 侧边栏 ==============

with st.sidebar:
    st.title("🎭 角色设置")
    
    # API Key 输入
    st.subheader("🔑 API 配置")
    api_key_input = st.text_input(
        "DeepSeek API Key",
        value=API_KEY,
        type="password",
        help="请输入你的 DeepSeek API Key"
    )
    
    if api_key_input:
        API_KEY = api_key_input
    
    st.divider()
    
    # 角色选择
    st.subheader("🎭 选择角色")
    
    # 预设角色选择
    character_options = ["🎨 自定义角色"] + list(PRESET_CHARACTERS.keys())
    selected_character = st.selectbox(
        "预设角色",
        options=character_options,
        index=0 if st.session_state.current_character is None else (
            character_options.index(st.session_state.current_character) 
            if st.session_state.current_character in character_options 
            else 0
        )
    )
    
    # 自定义角色设置
    if selected_character == "🎨 自定义角色":
        st.markdown("---")
        st.markdown("**✨ 创建你的专属角色**")
        
        custom_name = st.text_input("角色名称", placeholder="例如：小助手")
        custom_avatar = st.selectbox(
            "选择头像",
            options=["🤖", "👽", "🦊", "🐼", "🦁", "🐸", "👻", "🎃", "🌟", "💫"],
            index=0
        )
        custom_prompt = st.text_area(
            "角色设定",
            placeholder="描述角色的性格、说话方式、特点等...\n\n例如：你是一个活泼开朗的助手，喜欢用emoji表情，说话很有活力...",
            height=150
        )
        
        if custom_name and custom_prompt:
            st.session_state.character_config = {
                "name": custom_name,
                "avatar": custom_avatar,
                "system_prompt": custom_prompt
            }
            st.session_state.current_character = "🎨 自定义角色"
    else:
        # 使用预设角色
        st.session_state.character_config = PRESET_CHARACTERS[selected_character]
        st.session_state.current_character = selected_character
        
        # 显示角色信息
        config = PRESET_CHARACTERS[selected_character]
        st.markdown(f"""
        <div class="character-card">
            <h3>{config['avatar']} {config['name']}</h3>
            <p>{config['system_prompt'][:100]}...</p>
        </div>
        """, unsafe_allow_html=True)
    
    st.divider()
    
    # 清空对话按钮
    if st.button("🗑️ 清空对话历史", use_container_width=True):
        st.session_state.messages = []
        st.rerun()
    
    # 显示统计信息
    st.divider()
    st.caption(f"📊 当前对话: {len(st.session_state.messages)} 条消息")

# ============== 主界面 ==============

st.title("🎭 AI 角色扮演聊天室")

# 检查配置
if not API_KEY:
    st.warning("⚠️ 请在侧边栏输入你的 DeepSeek API Key")
    st.info("""
    **如何获取 API Key？**
    1. 访问 [DeepSeek 开放平台](https://platform.deepseek.com)
    2. 注册并登录账号
    3. 进入控制台创建 API Key
    """)
    st.stop()

if not st.session_state.character_config:
    st.info("👈 请在侧边栏选择一个角色或创建自定义角色")
    st.stop()

# 显示当前角色信息
current_config = st.session_state.character_config
col1, col2 = st.columns([1, 4])
with col1:
    st.markdown(f"<h1 style='font-size: 4rem; margin: 0;'>{current_config['avatar']}</h1>", unsafe_allow_html=True)
with col2:
    st.markdown(f"### 正在与 **{current_config['name']}** 对话")
    st.caption("试着和 TA 聊聊天吧！")

st.divider()

# 创建 OpenAI 客户端
client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

# 显示对话历史
for message in st.session_state.messages:
    avatar = current_config["avatar"] if message["role"] == "assistant" else "👤"
    with st.chat_message(message["role"], avatar=avatar):
        st.markdown(message["content"])

# 用户输入
if prompt := st.chat_input("输入你想说的话..."):
    # 添加用户消息
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user", avatar="👤"):
        st.markdown(prompt)
    
    # 构建完整的消息列表
    messages_for_api = [
        {"role": "system", "content": current_config["system_prompt"]}
    ] + st.session_state.messages
    
    # 调用 API 并流式输出
    with st.chat_message("assistant", avatar=current_config["avatar"]):
        try:
            stream = client.chat.completions.create(
                model=MODEL,
                messages=messages_for_api,
                stream=True,
                max_tokens=2000,
                temperature=0.8,  # 稍高的温度让角色更有个性
            )
            
            response = st.write_stream(
                chunk.choices[0].delta.content or ""
                for chunk in stream
                if chunk.choices[0].delta.content
            )
            
            # 保存助手回复
            st.session_state.messages.append({"role": "assistant", "content": response})
            
        except Exception as e:
            st.error(f"❌ 调用 API 出错: {str(e)}")
            st.info("请检查 API Key 是否正确，以及网络连接是否正常。")

# ============== 底部信息 ==============

st.divider()

with st.expander("💡 使用提示"):
    st.markdown("""
    **如何获得更好的角色扮演体验？**
    
    1. **选择合适的角色**：不同角色有不同的说话风格和个性
    2. **融入角色设定**：尝试用符合角色世界观的方式提问
    3. **自定义角色**：发挥创意，创建你想要的任何角色
    4. **多轮对话**：角色会记住之前的对话内容
    
    **自定义角色技巧：**
    - 明确角色的性格特点
    - 定义说话方式和口癖
    - 设定角色的背景故事
    - 指定角色对用户的称呼
    """)

st.caption("Made with ❤️ using DeepSeek API + Streamlit | Python 程序设计课程设计示例")