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

# Load environment variables
load_dotenv()

st.set_page_config(page_title="SmartNotes AI", page_icon="📚", layout="wide")

st.title("📚 SmartNotes: 智能笔记助手")
st.markdown("把碎片化笔记转化为结构化知识，还能自动生成测验！")

# Sidebar configuration
with st.sidebar:
    st.header("⚙️ 设置")
    
    # Model Provider Selection
    provider = st.selectbox("选择服务提供商", ["OpenAI", "DeepSeek", "Custom"], index=0)
    
    # Dynamic Configuration based on Provider
    if provider == "DeepSeek":
        base_url_default = "https://api.deepseek.com"
        model_options = ["deepseek-chat", "deepseek-reasoner"]
        api_key_help = "请前往 https://platform.deepseek.com 获取 API Key"
    elif provider == "OpenAI":
        base_url_default = "https://api.openai.com/v1"
        model_options = ["gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo"]
        api_key_help = "请前往 https://platform.openai.com 获取 API Key"
    else:
        base_url_default = "https://api.openai.com/v1"
        model_options = ["gpt-4o", "gpt-3.5-turbo", "deepseek-chat"]
        api_key_help = "请输入兼容 OpenAI 格式的 API Key"

    api_key = st.text_input("API Key", type="password", value=os.getenv("OPENAI_API_KEY", ""), help=api_key_help)
    base_url = st.text_input("Base URL", value=os.getenv("OPENAI_BASE_URL", base_url_default))
    
    # Allow custom model input if needed, otherwise select from presets
    if provider == "Custom":
        model = st.text_input("模型名称", value="gpt-3.5-turbo")
    else:
        model = st.selectbox("选择模型", model_options, index=0)
    
    st.subheader("关于项目")
    st.info("本项目利用 AI 帮助你快速整理学习笔记并生成自我测试题。")

# Main Interface
col1, col2 = st.columns([1, 1])

with col1:
    st.subheader("📝 原始笔记")
    raw_notes = st.text_area("在此粘贴你的笔记...", height=400, placeholder="例如：\n今天学习了光合作用的过程...\n1. 光反应阶段...\n2. 暗反应阶段...")
    
    action_type = st.radio("选择操作", ["整理 & 摘要", "生成测验"], horizontal=True)
    
    if st.button("🚀 开始处理", type="primary", use_container_width=True):
        if not api_key:
            st.error("请在侧边栏配置 API Key")
        elif not raw_notes:
            st.warning("请先输入笔记内容")
        else:
            client = OpenAI(api_key=api_key, base_url=base_url)
            
            with st.spinner("AI 正在大脑风暴中..."):
                try:
                    if action_type == "整理 & 摘要":
                        prompt = f"""
                        你是一个专业的学习助手。请将以下原始笔记整理为结构清晰的 Markdown 格式，并提供一段简明的总结。
                        
                        【原始笔记】:
                        {raw_notes}
                        
                        【输出要求】:
                        1. **核心摘要**: 3-5句话总结核心内容。
                        2. **结构化笔记**: 使用 # 标题、- 列表、**加粗** 关键词来整理内容。修复原本的语病或模糊之处。
                        """
                        response = client.chat.completions.create(
                            model=model,
                            messages=[{"role": "user", "content": prompt}]
                        )
                        st.session_state.result = response.choices[0].message.content
                        st.session_state.action = "structure"

                    elif action_type == "生成测验":
                        prompt = f"""
                        你是一个出题专家。请基于以下笔记内容，生成 3 道单项选择题来测试用户的理解。
                        
                        【原始笔记】:
                        {raw_notes}
                        
                        【输出要求】:
                        请按以下格式输出每一题：
                        
                        **Q1. [题目内容]**
                        A) [选项A]
                        B) [选项B]
                        C) [选项C]
                        D) [选项D]
                        
                        *点击下方查看答案*
                        <details>
                        <summary>查看答案 & 解析</summary>
                        **答案**: [正确选项]
                        **解析**: [简短解析]
                        </details>
                        
                        ---
                        (继续下一题)
                        """
                        response = client.chat.completions.create(
                            model=model,
                            messages=[{"role": "user", "content": prompt}]
                        )
                        st.session_state.result = response.choices[0].message.content
                        st.session_state.action = "quiz"

                except Exception as e:
                    st.error(f"发生错误: {e}")

with col2:
    st.subheader("✨ 结果输出")
    if "result" in st.session_state:
        st.markdown(st.session_state.result, unsafe_allow_html=True)
    else:
        st.info("此处将显示 AI 处理后的结果...")