use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
pub enum SignalMessage {
    /// 加入会议
    Join { room_id: String, user_id: String, metadata: HashMap<String, String> },
    /// 离开会议
    Leave { room_id: String, user_id: String },
    /// SDP Offer/Answer
    Sdp { from: String, to: String, sdp: String },
    /// ICE Candidate
    Ice { from: String, to: String, candidate: String },
    /// 服务端下发：用户列表变更
    PeerList { room_id: String, peers: Vec<String> },
    /// 错误/拒绝
    Error { code: u16, message: String },
}

use bytes::{Bytes, BytesMut};
use tokio::sync::mpsc;

/// 每个 WebSocket 连接对应一个 Sender，用于下行推送
pub type ClientTx = mpsc::UnboundedSender<Bytes>;

#[derive(Debug)]
pub struct ClientCtx {
    pub user_id: String,
    pub room_id: Option<String>,
    pub tx: ClientTx,
    /// 零拷贝写缓冲，避免频繁分配
    pub write_buf: BytesMut,
}

use async_trait::async_trait;
use crate::types::{SignalMessage, ClientCtx};
use std::sync::Arc;

#[async_trait]
pub trait SignalHandler: Send + Sync {
    /// 处理入站消息，返回需广播/单发的出站消息
    async fn handle(&self, ctx: &mut ClientCtx, msg: SignalMessage) -> Vec<SignalMessage>;
    
    /// 连接建立回调
    async fn on_connect(&self, _ctx: &mut ClientCtx) {}
    
    /// 连接断开回调
    async fn on_disconnect(&self, _ctx: &mut ClientCtx) {}
}

pub struct HandlerChain {
    handlers: Vec<Arc<dyn SignalHandler>>,
}

impl HandlerChain {
    pub fn new() -> Self { Self { handlers: Vec::new() } }
    
    pub fn add(mut self, h: Arc<dyn SignalHandler>) -> Self {
        self.handlers.push(h);
        self
    }
}

#[async_trait]
impl SignalHandler for HandlerChain {
    async fn handle(&self, ctx: &mut ClientCtx, msg: SignalMessage) -> Vec<SignalMessage> {
        let mut results = Vec::new();
        for h in &self.handlers {
            let mut out = h.handle(ctx, msg.clone()).await;
            results.append(&mut out);
        }
        results
    }
    
    async fn on_connect(&self, ctx: &mut ClientCtx) {
        for h in &self.handlers { h.on_connect(ctx).await; }
    }
    
    async fn on_disconnect(&self, ctx: &mut ClientCtx) {
        for h in &self.handlers { h.on_disconnect(ctx).await; }
    }
}

use dashmap::DashMap;
use tokio::sync::broadcast;

/// 房间内成员：user_id -> ClientTx
type RoomMembers = DashMap<String, ClientTx>;

/// 全局房间注册表
pub struct RoomRegistry {
    rooms: DashMap<String, RoomState>,
}

struct RoomState {
    members: RoomMembers,
    /// 广播通道，用于零拷贝分发
    broadcast_tx: broadcast::Sender<Bytes>,
}

impl RoomRegistry {
    pub fn new() -> Self {
        Self { rooms: DashMap::new() }
    }
    
    /// 获取或创建房间，返回广播发送端
    pub fn get_or_create(&self, room_id: &str) -> broadcast::Sender<Bytes> {
        self.rooms
            .entry(room_id.to_string())
            .or_insert_with(|| {
                let (tx, _rx) = broadcast::channel(1024);
                RoomState {
                    members: RoomMembers::new(),
                    broadcast_tx: tx,
                }
            })
            .broadcast_tx
            .clone()
    }
    
    /// 成员加入
    pub fn join(&self, room_id: &str, user_id: String, tx: ClientTx) {
        if let Some(mut room) = self.rooms.get_mut(room_id) {
            room.members.insert(user_id, tx);
        }
    }
    
    /// 成员离开
    pub fn leave(&self, room_id: &str, user_id: &str) {
        if let Some(mut room) = self.rooms.get_mut(room_id) {
            room.members.remove(user_id);
            // 房间为空可延迟清理
        }
    }
    
    /// 获取房间在线成员列表
    pub fn list_peers(&self, room_id: &str) -> Vec<String> {
        self.rooms
            .get(room_id)
            .map(|r| r.members.iter().map(|e| e.key().clone()).collect())
            .unwrap_or_default()
    }
}

use crate::types::{SignalMessage, ClientCtx};
use bytes::Bytes;
use std::sync::Arc;

pub struct BroadcastHandler {
    registry: Arc<RoomRegistry>,
}

#[async_trait]
impl SignalHandler for BroadcastHandler {
    async fn handle(&self, ctx: &mut ClientCtx, msg: SignalMessage) -> Vec<SignalMessage> {
        match msg {
            SignalMessage::Join { room_id, user_id, .. } => {
                let tx = self.registry.get_or_create(&room_id);
                ctx.room_id = Some(room_id.clone());
                self.registry.join(&room_id, user_id.clone(), ctx.tx.clone());
                
                // 通知房间内其他成员
                let peer_list = SignalMessage::PeerList {
                    room_id: room_id.clone(),
                    peers: self.registry.list_peers(&room_id),
                };
                let encoded = Bytes::from(serde_json::to_vec(&peer_list).unwrap());
                let _ = tx.send(encoded); // 零拷贝发送
                
                vec![]
            }
            SignalMessage::Sdp { from, to, sdp } | SignalMessage::Ice { from, to, candidate } => {
                // 单播转发：直接写入目标用户的 mpsc 通道
                if let Some(room_id) = &ctx.room_id {
                    if let Some(room) = self.registry.rooms.get(room_id) {
                        if let Some(target_tx) = room.members.get(&to) {
                            let encoded = Bytes::from(serde_json::to_vec(&msg).unwrap());
                            let _ = target_tx.send(encoded);
                        }
                    }
                }
                vec![]
            }
            _ => vec![msg], // 其他消息透传给下一个 Handler
        }
    }
    
    async fn on_disconnect(&self, ctx: &mut ClientCtx) {
        if let (Some(room_id), Some(user_id)) = (&ctx.room_id, Some(&ctx.user_id)) {
            self.registry.leave(room_id, user_id);
            // 广播更新后的成员列表
            let peer_list = SignalMessage::PeerList {
                room_id: room_id.clone(),
                peers: self.registry.list_peers(room_id),
            };
            let tx = self.registry.get_or_create(room_id);
            let _ = tx.send(Bytes::from(serde_json::to_vec(&peer_list).unwrap()));
        }
    }
}

use tokio::net::TcpListener;
use tokio_tungstenite::{accept_async, tungstenite::Message};
use futures_util::{SinkExt, StreamExt};
use tracing::{info, warn, error};

pub async fn run_server(addr: &str, handler_chain: Arc<HandlerChain>, registry: Arc<RoomRegistry>) -> anyhow::Result<()> {
    let listener = TcpListener::bind(addr).await?;
    info!("Signal gateway listening on {}", addr);
    
    loop {
        let (stream, addr) = listener.accept().await?;
        let handler_chain = handler_chain.clone();
        let registry = registry.clone();
        
        tokio::spawn(async move {
            if let Err(e) = handle_connection(stream, addr, handler_chain, registry).await {
                warn!("Connection {} error: {}", addr, e);
            }
        });
    }
}

async fn handle_connection(
    stream: tokio::net::TcpStream,
    addr: std::net::SocketAddr,
    handler_chain: Arc<HandlerChain>,
    registry: Arc<RoomRegistry>,
) -> anyhow::Result<()> {
    let ws_stream = accept_async(stream).await?;
    let (mut ws_sender, mut ws_receiver) = ws_stream.split();
    
    // 为每个连接创建 mpsc 通道，解耦网络发送与业务逻辑
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Bytes>();
    let mut ctx = ClientCtx {
        user_id: String::new(),
        room_id: None,
        tx,
        write_buf: BytesMut::with_capacity(4096),
    };
    
    // 连接建立回调
    handler_chain.on_connect(&mut ctx).await;
    
    // 发送任务：从 mpsc 接收并写入 WebSocket
    let send_task = tokio::spawn(async move {
        while let Some(bytes) = rx.recv().await {
            if ws_sender.send(Message::Binary(bytes.freeze())).await.is_err() {
                break;
            }
        }
    });
    
    // 接收任务：解析消息 -> 交给 Handler 链 -> 产出下行消息
    while let Some(Ok(msg)) = ws_receiver.next().await {
        if let Message::Binary(data) = msg {
            if let Ok(signal_msg) = serde_json::from_slice::<SignalMessage>(&data) {
                let out_msgs = handler_chain.handle(&mut ctx, signal_msg).await;
                for out in out_msgs {
                    if let Ok(bytes) = serde_json::to_vec(&out) {
                        let _ = ctx.tx.send(Bytes::from(bytes));
                    }
                }
            }
        }
    }
    
    // 连接断开清理
    handler_chain.on_disconnect(&mut ctx).await;
    send_task.abort();
    Ok(())
}

use prometheus::{register_int_counter_vec, register_histogram_vec, IntCounterVec, HistogramVec};

lazy_static::lazy_static! {
    static ref MSG_RECEIVED: IntCounterVec = register_int_counter_vec!(
        "signal_msg_received_total", "Total messages received", &["type"]
    ).unwrap();
    
    static ref MSG_LATENCY: HistogramVec = register_histogram_vec!(
        "signal_handler_latency_seconds", "Handler processing latency", &["handler"]
    ).unwrap();
    
    static ref ACTIVE_CONNECTIONS: IntCounterVec = register_int_counter_vec!(
        "signal_active_connections", "Current active connections", &["room_id"]
    ).unwrap();
}

// Cargo.toml 关键依赖
// openraft = { version = "0.9", features = ["tokio", "serde"] }
// rocksdb = { version = "0.21", features = ["backup"] }

use openraft::{Config, Raft, Storage, RaftNetwork, RaftStorage};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

/// 状态机数据：仅存元数据，不存媒体流
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
pub struct MetaState {
    pub rooms: BTreeMap<String, RoomMeta>, // room_id -> RoomMeta
    pub users: BTreeMap<String, UserMeta>, // user_id -> UserMeta
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct RoomMeta {
    pub owner: String,
    pub members: BTreeSet<String>, // 在线 user_id 集合
    pub policy: RoomPolicy,        // 权限/录制/水印策略
    pub version: u64,              // 乐观锁版本号
}

/// 实现 Raft 状态机接口
#[async_trait]
impl RaftStateMachine<TypeConfig> for MetaStore {
    async fn apply(&mut self, entry: &Entry<TypeConfig>) -> Result<Vec<u8>, StorageError> {
        match &entry.payload {
            EntryPayload::Normal(req) => match req {
                ClientRequest::JoinRoom { room_id, user_id } => {
                    self.state.rooms.entry(room_id.clone()).or_default().members.insert(user_id.clone());
                    Ok(vec![])
                }
                ClientRequest::LeaveRoom { room_id, user_id } => { /* ... */ }
                ClientRequest::UpdatePolicy { room_id, policy } => { /* ... */ }
            },
            _ => Ok(vec![]),
        }
    }
    
    async fn snapshot(&mut self) -> Result<Snapshot<TypeConfig>, StorageError> {
        // 定期将 RocksDB 做增量快照上传对象存储，新节点加入时极速追赶
    }
}

use tokio::time::{interval, Duration, MissedTickBehavior};

pub struct HeartbeatManager {
    /// 配置下发：服务端根据 RTT 动态调整
    pub ping_interval: Duration,   // 默认 15s
    pub pong_timeout: Duration,    // 默认 5s
    pub max_missed: u32,           // 连续丢包阈值，默认 3
}

impl HeartbeatManager {
    pub async fn run(&self, mut ws_sink: SplitSink<WebSocketStream<TcpStream>, Message>, 
                     mut pong_rx: mpsc::Receiver<()>) {
        let mut ticker = interval(self.ping_interval);
        ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
        let mut missed = 0u32;
        
        loop {
            tokio::select! {
                _ = ticker.tick() => {
                    if ws_sink.send(Message::Ping(Bytes::new())).await.is_err() { break; }
                    
                    // 等待 Pong，带超时
                    match tokio::time::timeout(self.pong_timeout, pong_rx.recv()).await {
                        Ok(Some(_)) => missed = 0,
                        _ => {
                            missed += 1;
                            if missed >= self.max_missed { break; }
                        }
                    }
                }
                _ = pong_rx.recv() => { /* 收到 Pong 重置计数 */ }
            }
        }
        // 触发连接关闭流程
    }
}

#[derive(Serialize, Deserialize)]
struct AckEnvelope {
    seq: u64,              // 单调递增序列号
    payload: SignalMessage,
    ack_required: bool,    // 关键消息（Join/Sdp）必须 ACK
}

impl ClientCtx {
    pub async fn send_reliable(&self, msg: SignalMessage, need_ack: bool) -> Result<(), SendError> {
        let seq = self.next_seq();
        let env = AckEnvelope { seq, payload: msg, ack_required: need_ack };
        let bytes = Bytes::from(serde_json::to_vec(&env)?);
        
        if need_ack {
            let (tx, rx) = oneshot::channel();
            self.pending_acks.lock().await.insert(seq, tx);
            self.tx.send(bytes)?;
            
            // 超时重传逻辑（指数退避 200ms/400ms/800ms...）
            tokio::spawn(retry_task(seq, bytes, self.tx.clone(), rx));
        } else {
            self.tx.send(bytes)?;
        }
        Ok(())
    }
}

// 1. 网关层终结 mTLS（使用 rustls + tokio-rustls）
let tls_config = rustls::ServerConfig::builder()
    .with_safe_defaults()
    .with_client_cert_verifier(Arc::new(CustomClientCertVerifier)) // 校验客户端证书链
    .with_single_cert(server_certs, server_key)?;

// 2. 握手后提取客户端证书 Subject Alternative Name 作为 device_id
// 3. 业务层校验 JWT（由 Auth Service 下发，有效期 5min，含 room_id/权限/水印策略）

// Handler 链前置：风控拦截器
pub struct RiskControlHandler {
    redis: redis::Client,
    lua_script: Script, // 原子化 Check-And-Set
}

#[async_trait]
impl SignalHandler for RiskControlHandler {
    async fn handle(&self, ctx: &mut ClientCtx, msg: SignalMessage) -> Vec<SignalMessage> {
        let key = format!("risk:{}:{}", ctx.user_id, msg.msg_type());
        // Lua 脚本：INCR + EXPIRE + 判断阈值，单次 RTT 完成
        let count: u32 = self.redis.eval(&self.lua_script, &[&key], &[&ctx.ip]).await.unwrap_or(0);
        if count > THRESHOLD {
            return vec![SignalMessage::Error { code: 429, message: "频率超限".into() }];
        }
        vec![msg]
    }
}

// 替换 serde_json -> rkyv (零拷贝归档)
#[derive(Archive, Deserialize, Serialize, Debug)]
#[archive_attr(derive(CheckBytes))]
pub struct ArchivedSignalMessage { /* 字段同 SignalMessage */ }

// 网络层直接零拷贝解析
fn parse_zero_copy(buf: &[u8]) -> Result<&ArchivedSignalMessage, ArchiveError> {
    unsafe { archived_root::<ArchivedSignalMessage>(buf) } // 无分配、无拷贝、无解析开销
}

# Cargo.toml
tokio = { version = "1.38", features = ["rt-multi-thread", "net", "io-uring"] }
# 或迁移至 monoio/glommio 获得更纯粹的 io_uring 体验

# K8s Ingress / Gateway API 路由规则
apiVersion: gateway.networking.k8s.io/v1beta1
kind: HTTPRoute
metadata:
  name: signal-canary
spec:
  hostnames: ["signal.api.example.com"]
  rules:
  - matches:
    - headers:
        x-canary: "true"  # 内测/压测流量打标
    backendRefs:
    - name: signal-gateway-canary
      weight: 100
  - backendRefs:
    - name: signal-gateway-stable
      weight: 90
    - name: signal-gateway-canary
      weight: 10

# 示例：连接数倾斜自动再平衡
- name: Detect connection skew
  query: "max by (pod) (active_connections) / min by (pod) (active_connections) > 1.5"
  interval: 1m
  
- name: Trigger graceful drain
  action: 
    1. 标记 Pod `draining=true`（停止接收新连接）
    2. 等待 `active_connections == 0` 或超时 5min
    3. 触发 RollingUpdate 替换新 Pod
    4. 校验新 Pod `Ready=True` 后剔除旧 Pod

// Handler 中根据带宽探测结果，下发层订阅指令
SignalMessage::LayerSubscription {
    room_id,
    subscriber: user_id,
    publisher: peer_id,
    spatial_layers: vec![0, 1],      // 仅订阅低/中分辨率
    temporal_layers: vec![0],        // 仅基础帧率
}

pub struct RecordingOrchestrator {
    k8s_client: kube::Client,
    object_storage: s3::Client,
}

impl RecordingOrchestrator {
    pub async fn start(&self, room_id: &str, policy: &RecordingPolicy) -> Result<String, Error> {
        // 1. 创建 Job/Pod 运行 FFmpeg/Rust 录制模块
        // 2. 通过信令网关下发 `RecordingStarted { task_id, layout_config }` 给 SFU
        // 3. SFU 混流推送至录制模块（SRT/RTMP/WebRTC）
        // 4. 完成后回调 `RecordingFinished { file_url, sha256, duration }` 入库审计
    }
}