// session_context.h
#pragma once
#include <cstdint>
#include <atomic>
#include <array>
#include <span>

// 固定大小，便于共享内存/Redis Hash 直接映射，避免序列化开销
// 版本号用于乐观锁 CAS 更新
struct alignas(64) SessionContext { // 缓存行对齐，防止伪共享
    uint64_t version;               // 乐观锁版本号 (Redis INCR 或 etcd ModRevision)
    uint64_t connection_id;         // 全局唯一 Connection ID (UUIDv7 高 64 位)
    uint32_t ssrc_local;            // 本地 SSRC
    uint32_t ssrc_remote;           // 远端 SSRC
    uint16_t rtp_seq_local;         // 本地发送序列号
    uint16_t rtp_seq_remote;        // 远端接收序列号 (用于 NACK/重排)
    uint32_t rtp_ts_offset;         // 时间戳偏移
    
    // SRTP 加密上下文 (AES-GCM / ChaCha20-Poly1305) - 仅保留密钥材料与滚动计数器
    // 实际生产建议使用 libsrtp2 的 srtp_t 指针在共享内存中无法直接共享，
    // 故此处存储 Master Key/Salt + ROC/SEQ 重建参数，迁移时重新 srtp_create()
    std::array<uint8_t, 32> srtp_master_key; 
    std::array<uint8_t, 14> srtp_master_salt;
    uint32_t srtp_roc_tx;           // Roll-over Counter TX
    uint32_t srtp_roc_rx;           // Roll-over Counter RX
    
    // GCC/BWE 状态 (Google Congestion Control / Transport-CC)
    int64_t  bwe_bitrate_bps;       // 当前估计带宽
    uint32_t bwe_rtt_ms;            // 当前 RTT
    uint8_t  bwe_state;             // 0:Startup 1:Drain 2:ProbeBW 3:ProbeRTT
    
    // 迁移专用字段
    uint64_t migration_token;       // 迁移令牌 (防重放)
    uint32_t migration_deadline_ms; // 迁移截止时间戳
    uint8_t  status;                // 0:Active 1:Migrating 2:Draining 3:Closed
    
    // 扩展区 (TLV 格式)，用于未来兼容
    uint16_t ext_len;
    uint8_t  ext_data[128];         
};

// 共享内存池管理器 (基于 memfd_create + HugeTLB)
class SessionShmPool {
public:
    // 从池中分配/获取 SessionContext* (零拷贝映射)
    SessionContext* acquire(uint64_t conn_id); 
    void release(uint64_t conn_id);
    // 批量持久化脏页至 Redis (异步协程)
    void flush_dirty_to_redis(redisAsyncContext* ctx); 
private:
    int shm_fd_;
    SessionContext* base_ptr_;
    size_t capacity_;
    std::atomic<uint32_t>* dirty_bitmap_; // 脏页位图
};

// io_uring_engine.h
#pragma once
#include <liburing.h>
#include <vector>
#include <thread>
#include <atomic>
#include <functional>

class IoUringEngine {
public:
    struct Config {
        int entries = 8192;           // SQ/CQ 条目数
        int sq_poll_cpu = -1;         // SQPOLL 线程绑定 CPU 核心 ID
        int buffer_group_id = 0;      // 缓冲区组 ID
        size_t buffer_size = 2048;    // 单个缓冲区大小 (MTU + 头部预留)
        size_t buffer_count = 65536;  // 缓冲区总数 (巨页内存)
    };

    explicit IoUringEngine(const Config& cfg);
    ~IoUringEngine();

    // 禁止拷贝，仅移动
    IoUringEngine(const IoUringEngine&) = delete;
    IoUringEngine& operator=(const IoUringEngine&) = delete;

    // 启动事件循环 (阻塞当前线程，通常跑在独立绑核线程)
    void run(std::stop_token token);

    // 提交接收请求 (零拷贝模式)
    // 返回值: 0 成功, -EAGAIN SQ 满, -ENOBUFS 缓冲区组耗尽
    int submit_recv(int fd, uint64_t conn_id, uint32_t flags = 0);

    // 提交发送请求 (支持 MSG_ZEROCOPY)
    int submit_send(int fd, const struct iovec* iov, int iovcnt, uint64_t conn_id);

    // 注册/更新文件描述符表 (迁移接管 FD 核心)
    int register_files(const std::vector<int>& fds); // 返回起始 fixed_fd 索引
    int unregister_files(uint32_t start_idx, uint32_t count);

    // 回调注册
    using RecvCallback = std::function<void(uint64_t conn_id, struct io_uring_cqe* cqe)>;
    using SendCallback = std::function<void(uint64_t conn_id, int status)>;
    using MigrationCallback = std::function<void(uint64_t conn_id, MigrationAction action)>;

    void set_recv_cb(RecvCallback cb) { recv_cb_ = std::move(cb); }
    void set_send_cb(SendCallback cb) { send_cb_ = std::move(cb); }
    void set_migration_cb(MigrationCallback cb) { migration_cb_ = std::move(cb); }

private:
    // 内部事件循环核心
    void process_cqes();
    void handle_recv_cqe(struct io_uring_cqe* cqe);
    void handle_send_cqe(struct io_uring_cqe* cqe);
    void handle_migration_cmd(uint64_t conn_id, MigrationAction action); // 处理控制面下发指令

    struct io_uring ring_;
    Config cfg_;
    RecvCallback recv_cb_;
    SendCallback send_cb_;
    MigrationCallback migration_cb_;
    
    // 缓冲区池元数据 (用户态管理，内核仅持有引用)
    struct BufferPool { void* base; std::atomic<uint32_t>* free_bitmap; } buf_pool_;
    
    // FD -> Connection ID 映射 (固定数组或 robin_hood_hashing)
    std::vector<uint64_t> fd_to_conn_; 
};

// migration_handler.cpp
void MigrationHandler::on_migration_start(uint64_t conn_id, const MigrationCmd& cmd) {
    auto* ctx = shm_pool_.acquire(conn_id);
    if (!ctx) return log_error("Session not found");

    // 1. CAS 更新状态为 MIGRATING
    uint64_t expected_ver = ctx->version;
    SessionContext new_ctx = *ctx;
    new_ctx.status = SessionStatus::MIGRATING;
    new_ctx.migration_token = cmd.token;
    new_ctx.migration_deadline_ms = now_ms() + cmd.timeout_ms;
    new_ctx.version++; 

    // 2. 原子写入 Redis (Lua 脚本保证原子性: CHECK version -> SET new -> RETURN ok)
    bool ok = redis_cas_update(conn_id, expected_ver, new_ctx);
    if (!ok) return log_warn("CAS conflict, retry later");

    // 3. 通知 io_uring 引擎停止业务逻辑投递，仅保活
    io_engine_.pause_connection(conn_id); 

    // 4. 异步快照完整上下文至 Redis (Pipeline 批量)
    async_snapshot_to_redis(conn_id, new_ctx);

    // 5. 准备 FD 传递 (通过 Unix Domain Socket 发送给目标节点 Agent)
    // 利用 io_uring 注册的 fixed_fd 索引，避免重复查找
    int fixed_fd = fd_registry_.get_fixed_fd(conn_id);
    send_fd_over_unix_socket(cmd.target_agent_sock, fixed_fd, conn_id, cmd.token);
}

void MigrationHandler::on_fd_received(uint64_t conn_id, int received_fd, uint64_t token) {
    // 目标节点接管入口
    // 1. 校验 Token
    // 2. 注册 FD 至本地 io_uring (IORING_REGISTER_FILES_UPDATE 或重新注册)
    int fixed_idx = io_engine_.register_single_fd(received_fd);
    
    // 3. 从 Redis 拉取最新 Context (带版本号)
    auto ctx = redis_get_session(conn_id);
    
    // 4. 重建 SRTP 上下文 (libsrtp2)
    srtp_session_ = srtp_create(&policy); // policy 基于 ctx->srtp_master_key/salt
    srtp_set_roc(srtp_session_, ctx->srtp_roc_tx, ctx->srtp_roc_rx); // 恢复 ROC
    
    // 5. 预投递 RECV 请求 (零拷贝)
    io_engine_.submit_recv(received_fd, conn_id, IORING_RECV_MULTISHOT); // 5.19+ Multishot
    
    // 6. 通知 Orchestrator "Ready for Traffic Switch"
    orchestrator_client_.notify_ready(conn_id, token);
}

// tenant_quota.h
struct TenantQuota {
    uint32_t max_concurrent_sessions;   // 最大并发路数
    uint64_t max_bandwidth_bps;         // 最大出口带宽
    uint32_t max_cpu_quota_percent;     // CPU 使用率上限 (CGroup v2 绑定)
    uint32_t priority;                  // 调度优先级 (CFS shares / io_uring SQPOLL 亲和性)
    bool     dedicated_numa_node;       // 是否独占 NUMA 节点 (物理隔离)
};

// 调度器决策逻辑 (Orchestrator 侧)
bool Scheduler::can_place_session(const TenantQuota& q, const NodeMetrics& m) {
    // 硬性指标
    if (m.active_sessions >= q.max_concurrent_sessions) return false;
    if (m.current_bandwidth_bps >= q.max_bandwidth_bps) return false;
    if (m.cpu_usage_percent >= q.max_cpu_quota_percent) return false;
    
    // 软性亲和性 (优先填满低优先级节点，保留高优先级节点给大客户)
    if (q.priority > m.node_priority && m.has_lower_priority_space()) return false; 
    
    // NUMA 亲和性校验
    if (q.dedicated_numa_node && !m.has_free_numa_node(q.tenant_id)) return false;
    
    return true;
}

# test_migration_correctness.py (Pytest + Hypothesis)
@given(st.data())
@settings(max_examples=1000, deadline=None)
def test_migration_state_consistency(data):
    # 1. 构造随机会话状态 (SSRC, ROC, SEQ, BWE 状态)
    ctx = random_session_context(data)
    
    # 2. 启动源节点、目标节点 (Testcontainers)
    src = NodeProcess(version="v1.2.0")
    dst = NodeProcess(version="v1.3.0-rc1") # 版本不一致场景
    
    # 3. 建立真实 RTP 流量 (使用 gstreamer/webrtc-driver)
    conn_id = establish_call(src, dst, ctx, duration_sec=30)
    
    # 4. 触发迁移 (模拟 Orchestrator 指令)
    migration_id = trigger_migration(conn_id, src, dst)
    
    # 5. 并发施压：迁移期间持续发包、收包、发送 NACK/PLI
    traffic_gen = TrafficGenerator(conn_id, rate_pps=500)
    traffic_gen.start()
    
    # 6. 等待迁移完成
    wait_migration_done(migration_id, timeout=10)
    
    # 7. 核心断言：状态一致性
    final_ctx_src = src.get_session_state(conn_id) # 应已清理
    final_ctx_dst = dst.get_session_state(conn_id)
    
    # 关键不变量校验
    assert final_ctx_dst.ssrc_local == ctx.ssrc_local
    assert final_ctx_dst.rtp_seq_local >= ctx.rtp_seq_local # 单调递增
    assert final_ctx_dst.srtp_roc_tx == ctx.srtp_roc_tx # ROC 必须精确恢复
    assert final_ctx_dst.bwe_bitrate_bps == ctx.bwe_bitrate_bps # BWE 状态平滑
    
    # 8. 业务不变量：无花屏、无静音、无重连
    assert traffic_gen.stats.packet_loss_rate < 0.001
    assert traffic_gen.stats.max_jitter_ms < 50
    assert not traffic_gen.stats.has_reconnect_event

graph LR
    A[Client] -->|SRTP/UDP| B(io_uring EnginenCPU Core 0-7)
    B -->|Decrypt + Parse| C{Shared Mem RingnHugepages}
    C -->|Frame Ptr| D[GPU DecodernNVDEC/VA-API]
    D -->|YUV Frame| E[AI EnhancenTensorRT/ONNX Runtime]
    E -->|Enhanced Frame| F[GPU EncodernNVENC/QSV]
    F -->|Bitstream| C
    C -->|Encrypt + Send| B
    B -->|SRTP/UDP| A