import asyncssh
from contextlib import asynccontextmanager
from collections.abc import AsyncIterator
from dataclasses import dataclass
from starlette.applications import Starlette
# --- 1. SSH 连接参数 ---
REMOTE_HOST = '202.121.182.208'
REMOTE_USER = 'koko125'
# 确保路径正确，Windows 路径建议使用正斜杠 '/'
PRIVATE_KEY_PATH = 'D:/tool/tool/id_rsa.txt'
INITIAL_WORKING_DIRECTORY = f'/cluster/home/{REMOTE_USER}/sandbox'


# --- 2. 定义共享上下文 ---
# 这个数据类将持有所有模块需要共享的资源
@dataclass
class SharedAppContext:
    """在整个服务器生命周期内持有的共享资源"""
    ssh_connection: asyncssh.SSHClientConnection
    sandbox_path: str


# --- 3. 创建生命周期管理器 ---
# 这是整个架构的核心，负责在服务器启动时连接SSH，在关闭时断开
@asynccontextmanager
async def shared_lifespan(app: Starlette) -> AsyncIterator[SharedAppContext]:
    """
    为整个应用管理共享资源的生命周期。
    """
    print("主应用启动，正在建立共享 SSH 连接...")
    conn = None
    try:
        # 建立 SSH 连接
        conn = await asyncssh.connect(
            REMOTE_HOST,
            username=REMOTE_USER,
            client_keys=[PRIVATE_KEY_PATH]
        )
        print(f"SSH 连接到 {REMOTE_HOST} 成功！")

        # 使用 yield 将创建好的共享资源上下文传递给 Starlette 应用
        # 服务器会在此处暂停并开始处理请求
        yield {"shared_context": SharedAppContext(ssh_connection=conn, sandbox_path=INITIAL_WORKING_DIRECTORY)}

    finally:
        # 当服务器关闭时，yield 之后的代码会被执行
        if conn:
            conn.close()
            await conn.wait_closed()
            print("共享 SSH 连接已关闭。")