import reflex as rx
import reflex as rx


# ==============================================
# 全局状态（已修复所有类型错误）
# ==============================================
class State(rx.State):
    # 登录状态
    is_logged_in: bool = False
    username: str = ""
    password: str = ""

    # 当前页面
    current_page: str = "chat"

    # 会话列表
    sessions: list[dict] = [
        {"id": 1, "title": "会话 1"},
        {"id": 2, "title": "会话 2"},
        {"id": 3, "title": "会话 3"},
    ]
    current_session_id: int | None = 1

    # 登录弹窗控制
    show_login_modal: bool = True

    # ------------------------------
    # 显式定义 setter（修复报错）
    # ------------------------------
    @rx.event
    def set_username(self, value: str):
        self.username = value

    @rx.event
    def set_password(self, value: str):
        self.password = value

    # ------------------------------
    # 登录 / 退出
    # ------------------------------
    @rx.event
    def login(self):
        if self.username.strip() and self.password.strip():
            self.is_logged_in = True
            self.show_login_modal = False

    @rx.event
    def logout(self):
        self.is_logged_in = False
        self.show_login_modal = True
        self.username = ""
        self.password = ""

    # ------------------------------
    # 页面切换
    # ------------------------------
    @rx.event
    def set_page(self, page: str):
        self.current_page = page

    # ------------------------------
    # 会话操作
    # ------------------------------
    @rx.event
    def new_session(self):
        new_id = max([s["id"] for s in self.sessions], default=0) + 1
        self.sessions.append({"id": new_id, "title": f"会话 {new_id}"})
        self.current_session_id = new_id

    @rx.event
    def delete_session(self, session_id: int):
        self.sessions = [s for s in self.sessions if s["id"] != session_id]
        if self.current_session_id == session_id:
            if self.sessions:
                self.current_session_id = self.sessions[0]["id"]
            else:
                self.current_session_id = None

    @rx.event
    def set_current_session_id(self, sid: int):
        self.current_session_id = sid


# ==============================================
# 登录弹窗（已修复）
# ==============================================
def login_modal():
    return rx.cond(
        State.show_login_modal,
        rx.box(
            rx.box(
                rx.vstack(
                    rx.heading("登录"),
                    rx.input(
                        placeholder="用户名",
                        value=State.username,
                        on_change=State.set_username,
                        width="100%",
                    ),
                    rx.input(
                        placeholder="密码",
                        type="password",
                        value=State.password,
                        on_change=State.set_password,
                        width="100%",
                    ),
                    rx.button(
                        "登录",
                        on_click=State.login,
                        width="100%",
                        bg="blue",
                        color="white",
                    ),
                    spacing="4",
                    padding="6",
                    width="350px",
                ),
                bg="white",
                padding="8",
                border_radius="lg",
                shadow="lg",
            ),
            position="fixed",
            inset="0",
            bg="rgba(0,0,0,0.5)",
            display="flex",
            align_items="center",
            justify_content="center",
            z_index="1000",
        ),
    )


# ==============================================
# 左侧边栏
# ==============================================
def sidebar():
    return rx.box(
        rx.vstack(
            rx.heading("AI 智能体", font_size="lg"),
            rx.button(
                "💬 聊天",
                on_click=State.set_page("chat"),
                width="100%",
                bg=rx.cond(State.current_page == "chat", "blue", "gray"),
                color="white",
            ),
            rx.button(
                "📚 知识库",
                on_click=State.set_page("knowledge"),
                width="100%",
                bg=rx.cond(State.current_page == "knowledge", "blue", "gray"),
                color="white",
            ),
            rx.button(
                "⚙️ 设置",
                on_click=State.set_page("settings"),
                width="100%",
                bg=rx.cond(State.current_page == "settings", "blue", "gray"),
                color="white",
            ),
            rx.spacer(),
            rx.button(
                "🚪 退出登录",
                on_click=State.logout,
                width="100%",
                bg="red",
                color="white",
            ),
            padding="4",
            height="100vh",
            spacing="3",
        ),
        width="220px",
        bg="white",
        shadow="md",
    )


# ==============================================
# 聊天页面
# ==============================================
def chat_page():
    return rx.hstack(
        rx.box(
            rx.vstack(
                rx.button(
                    "➕ 新建会话",
                    on_click=State.new_session,
                    width="100%",
                    bg="green",
                    color="white",
                ),
                rx.foreach(
                    State.sessions,
                    lambda s: rx.hstack(
                        rx.text(
                            s["title"],
                            flex=1,
                            cursor="pointer",
                            on_click=lambda: State.set_current_session_id(s["id"]),
                        ),
                        rx.icon(
                            "trash",
                            size=16,
                            color="red",
                            cursor="pointer",
                            on_click=lambda: State.delete_session(s["id"]),
                        ),
                        width="100%",
                        padding="2",
                        bg=rx.cond(
                            State.current_session_id == s["id"],
                            "gray.100",
                            "transparent",
                        ),
                        border_radius="sm",
                    ),
                ),
                spacing="2",
                align_items="start",
                width="100%",
            ),
            width="220px",
            padding="4",
            border_right="1px solid #e2e8f0",
        ),
        rx.box(
            rx.heading("聊天内容"),
            flex=1,
            padding="4",
        ),
        flex=1,
    )


# ==============================================
# 内容区域
# ==============================================
def content_area():
    return rx.box(
        rx.cond(
            State.current_page == "chat",
            chat_page(),
            rx.cond(
                State.current_page == "knowledge",
                rx.heading("知识库管理"),
                rx.heading("系统设置"),
            ),
        ),
        flex=1,
    )


# ==============================================
# 主页面
# ==============================================
def index():
    return rx.box(
        login_modal(),
        rx.cond(
            State.is_logged_in,
            rx.hstack(
                sidebar(),
                content_area(),
            ),
            rx.center(
                rx.heading("请登录使用系统"),
                height="100vh",
            ),
        ),
        min_height="100vh",
        bg="gray.50",
    )


# ==============================================
# 启动
# ==============================================
app = rx.App()
app.add_page(index)