"""
Backend 监控接口需要补充的安全验证代码
添加到 app/api/routers/monitor.py
"""

from fastapi import Header

class ServerHeartbeatSecure(BaseModel):
    identifier: str
    challenge_uuid: str  # 新增：必须提供 challenge
    openclaw_version: str | None = None
    agents: List[dict] = []
    cpu_pct: float | None = None
    mem_pct: float | None = None
    disk_pct: float | None = None
    swap_pct: float | None = None


@router.post('/server/heartbeat')
def server_heartbeat(
    payload: ServerHeartbeatSecure,
    x_challenge_uuid: str = Header(..., description='Challenge UUID from registration'),
    db: Session = Depends(get_db)
):
    """
    安全版本的心跳接口，验证 challenge_uuid
    """
    # 1. 验证服务器存在且启用
    server = db.query(MonitoredServer).filter(
        MonitoredServer.identifier == payload.identifier,
        MonitoredServer.is_enabled == True
    ).first()
    if not server:
        raise HTTPException(status_code=404, detail='unknown server identifier')
    
    # 2. 验证 challenge_uuid 存在且有效
    ch = db.query(ServerChallenge).filter(
        ServerChallenge.challenge_uuid == x_challenge_uuid,
        ServerChallenge.server_id == server.id
    ).first()
    
    if not ch:
        raise HTTPException(status_code=401, detail='invalid challenge')
    
    if ch.expires_at < datetime.now(timezone.utc):
        raise HTTPException(status_code=401, detail='challenge expired')
    
    # 3. 可选：检查 challenge 是否已被使用过
    # 如果是首次验证，标记为已使用
    if ch.used_at is None:
        ch.used_at = datetime.now(timezone.utc)
    
    # 4. 存储状态
    st = db.query(ServerState).filter(ServerState.server_id == server.id).first()
    if not st:
        st = ServerState(server_id=server.id)
        db.add(st)
    
    st.openclaw_version = payload.openclaw_version
    st.agents_json = json.dumps(payload.agents, ensure_ascii=False)
    st.cpu_pct = payload.cpu_pct
    st.mem_pct = payload.mem_pct
    st.disk_pct = payload.disk_pct
    st.swap_pct = payload.swap_pct
    st.last_seen_at = datetime.now(timezone.utc)
    
    db.commit()
    
    return {
        'ok': True,
        'server_id': server.id,
        'last_seen_at': st.last_seen_at,
        'challenge_valid_until': ch.expires_at.isoformat()
    }


# 或者，如果需要长期有效的 API Key 方式：

class ServerHeartbeatApiKey(BaseModel):
    identifier: str
    openclaw_version: str | None = None
    agents: List[dict] = []
    cpu_pct: float | None = None
    mem_pct: float | None = None
    disk_pct: float | None = None
    swap_pct: float | None = None


@router.post('/server/heartbeat-v2')
def server_heartbeat_v2(
    payload: ServerHeartbeatApiKey,
    x_api_key: str = Header(..., description='Server API Key'),
    db: Session = Depends(get_db)
):
    """
    使用 API Key 的心跳接口（长期有效，不需要 challenge）
    需要在 MonitoredServer 模型中添加 api_key 字段
    """
    server = db.query(MonitoredServer).filter(
        MonitoredServer.identifier == payload.identifier,
        MonitoredServer.is_enabled == True,
        MonitoredServer.api_key == x_api_key  # 需要添加 api_key 字段
    ).first()
    
    if not server:
        raise HTTPException(status_code=401, detail='invalid identifier or api key')
    
    # ... 存储状态 ...