//! WebSocket 处理

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Path, Query, State,
    },
    response::IntoResponse,
};
use futures::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::mpsc;
use uuid::Uuid;

use crate::services::{
    device::DeviceRepository,
    session::SessionRepository,
    AppState,
};

/// WebSocket 查询参数
#[derive(Debug, Deserialize)]
pub struct WsQuery {
    pub device_id: String,
    pub token: Option<String>,
}

/// 信令消息
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum SignalMessage {
    /// 注册设备
    #[serde(rename = "register")]
    Register {
        device_id: String,
        verification_code: String,
    },
    /// 心跳
    #[serde(rename = "heartbeat")]
    Heartbeat { device_id: String },
    /// 心跳响应
    #[serde(rename = "heartbeat_ack")]
    HeartbeatAck,
    /// 连接请求
    #[serde(rename = "connect_request")]
    ConnectRequest {
        session_id: String,
        from_device: String,
        to_device: String,
        verification_code: String,
    },
    /// 连接响应
    #[serde(rename = "connect_response")]
    ConnectResponse {
        session_id: String,
        accepted: bool,
        reason: Option<String>,
    },
    /// SDP Offer
    #[serde(rename = "offer")]
    Offer {
        session_id: String,
        from_device: String,
        to_device: String,
        sdp: String,
    },
    /// SDP Answer
    #[serde(rename = "answer")]
    Answer {
        session_id: String,
        from_device: String,
        to_device: String,
        sdp: String,
    },
    /// ICE Candidate
    #[serde(rename = "candidate")]
    Candidate {
        session_id: String,
        from_device: String,
        to_device: String,
        candidate: String,
        sdp_mid: Option<String>,
        sdp_mline_index: Option<u32>,
    },
    /// 会话结束
    #[serde(rename = "session_end")]
    SessionEnd { session_id: String },
    /// 强制下线
    #[serde(rename = "force_offline")]
    ForceOffline { device_id: String },
    /// 错误
    #[serde(rename = "error")]
    Error { code: u32, message: String },
    /// 设置允许远程
    #[serde(rename = "set_allow_remote")]
    SetAllowRemote { device_id: String, allow: bool },
    /// 刷新验证码
    #[serde(rename = "refresh_code")]
    RefreshCode { device_id: String },
    /// 验证码已刷新
    #[serde(rename = "code_refreshed")]
    CodeRefreshed {
        device_id: String,
        verification_code: String,
    },
    /// 屏幕帧数据
    #[serde(rename = "screen_frame")]
    ScreenFrame {
        session_id: String,
        from_device: String,
        to_device: String,
        width: u32,
        height: u32,
        data: String,
    },
    /// 鼠标事件
    #[serde(rename = "mouse_event")]
    MouseEvent {
        session_id: String,
        from_device: String,
        to_device: String,
        x: f64,
        y: f64,
        event_type: String,
        button: Option<u8>,
        delta: Option<f64>,
    },
    /// 键盘事件
    #[serde(rename = "keyboard_event")]
    KeyboardEvent {
        session_id: String,
        from_device: String,
        to_device: String,
        key: String,
        event_type: String,
    },
    /// 流媒体设置
    #[serde(rename = "stream_settings")]
    StreamSettings {
        session_id: String,
        from_device: String,
        to_device: String,
        resolution: f64,
        fps: u32,
        quality: u32,
    },
    /// 获取显示器列表
    #[serde(rename = "get_displays")]
    GetDisplays {
        session_id: String,
        from_device: String,
        to_device: String,
    },
    /// 显示器列表响应
    #[serde(rename = "displays_list")]
    DisplaysList {
        session_id: String,
        from_device: String,
        to_device: String,
        displays: Vec<DisplayInfoMsg>,
        current_display: usize,
    },
    /// 切换显示器
    #[serde(rename = "switch_display")]
    SwitchDisplay {
        session_id: String,
        from_device: String,
        to_device: String,
        display_index: usize,
    },
    /// P2P 地址交换
    #[serde(rename = "p2p_exchange")]
    P2PExchange {
        session_id: String,
        from_device: String,
        to_device: String,
        local_addrs: Vec<String>,
        public_addr: Option<String>,
        udp_port: u16,
        turn_server: Option<TurnServerInfoMsg>,
    },
    /// P2P 连接就绪
    #[serde(rename = "p2p_ready")]
    P2PReady {
        session_id: String,
        from_device: String,
        to_device: String,
        connected_addr: String,
    },
    /// P2P 回退到服务器中转
    #[serde(rename = "p2p_fallback")]
    P2PFallback {
        session_id: String,
        from_device: String,
        to_device: String,
        reason: String,
    },
}

/// 显示器信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DisplayInfoMsg {
    pub index: usize,
    pub width: u32,
    pub height: u32,
    pub is_primary: bool,
    pub name: String,
}

/// TURN 服务器信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TurnServerInfoMsg {
    pub url: String,
    pub username: String,
    pub credential: String,
}

/// 信令WebSocket处理器
pub async fn signal_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    Query(query): Query<WsQuery>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_signal_socket(socket, state, query))
}

async fn handle_signal_socket(socket: WebSocket, state: Arc<AppState>, query: WsQuery) {
    let (mut sender, mut receiver) = socket.split();
    let (tx, mut rx) = mpsc::channel::<String>(32);

    let device_id = query.device_id.clone();

    // 注册连接
    state.add_connection(device_id.clone(), None, tx).await;

    // 更新设备在线状态
    {
        let device_repo = DeviceRepository::new(&state.db);
        if let Err(e) = device_repo.set_online(&device_id, true).await {
            tracing::error!("Failed to set device {} online: {}", device_id, e);
        }
    }

    tracing::info!("Device {} connected via WebSocket", device_id);

    // 发送任务
    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if sender.send(Message::Text(msg)).await.is_err() {
                break;
            }
        }
    });

    // 接收任务
    let state_clone = state.clone();
    let device_id_clone = device_id.clone();
    let recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = receiver.next().await {
            if let Message::Text(text) = msg {
                handle_signal_message(&state_clone, &device_id_clone, &text).await;
            }
        }
    });

    // 等待任务完成
    tokio::select! {
        _ = send_task => {},
        _ = recv_task => {},
    }

    // 清理 - 确保断开连接时设置离线状态
    tracing::info!("Device {} WebSocket disconnecting, cleaning up...", device_id);
    
    state.remove_connection(&device_id).await;
    
    // 重新创建 repo 来设置离线状态
    let device_repo = DeviceRepository::new(&state.db);
    if let Err(e) = device_repo.set_online(&device_id, false).await {
        tracing::error!("Failed to set device {} offline: {}", device_id, e);
    } else {
        tracing::info!("Device {} set to offline", device_id);
    }

    tracing::info!("Device {} disconnected", device_id);
}

async fn handle_signal_message(state: &Arc<AppState>, device_id: &str, text: &str) {
    let msg: SignalMessage = match serde_json::from_str(text) {
        Ok(m) => m,
        Err(e) => {
            tracing::warn!("Failed to parse message: {}", e);
            return;
        }
    };

    match msg {
        SignalMessage::Heartbeat { .. } => {
            let device_repo = DeviceRepository::new(&state.db);
            let _ = device_repo.update_last_seen(device_id).await;
            let response = SignalMessage::HeartbeatAck;
            let _ = state
                .send_to_device(device_id, &serde_json::to_string(&response).unwrap())
                .await;
        }

        SignalMessage::ConnectRequest {
            session_id,
            from_device,
            to_device,
            verification_code,
        } => {
            let device_repo = DeviceRepository::new(&state.db);

            // 验证目标设备是否在线且允许远程
            if let Ok(Some(target_device)) = device_repo.find_by_device_id(&to_device).await {
                if !target_device.is_online {
                    let error = SignalMessage::Error {
                        code: 1001,
                        message: "目标设备不在线".to_string(),
                    };
                    let _ = state
                        .send_to_device(&from_device, &serde_json::to_string(&error).unwrap())
                        .await;
                    return;
                }

                if !target_device.allow_remote {
                    let error = SignalMessage::Error {
                        code: 1002,
                        message: "目标设备未开启远程协助".to_string(),
                    };
                    let _ = state
                        .send_to_device(&from_device, &serde_json::to_string(&error).unwrap())
                        .await;
                    return;
                }

                // 验证验证码
                if target_device.verification_code.as_deref() != Some(&verification_code) {
                    let error = SignalMessage::Error {
                        code: 1003,
                        message: "验证码错误".to_string(),
                    };
                    let _ = state
                        .send_to_device(&from_device, &serde_json::to_string(&error).unwrap())
                        .await;
                    return;
                }

                // 创建会话
                let session_repo = SessionRepository::new(&state.db);
                let _ = session_repo
                    .create(&from_device, &to_device, "p2p")
                    .await;

                // 转发连接请求到目标设备
                let _ = state.send_to_device(&to_device, text).await;
            } else {
                let error = SignalMessage::Error {
                    code: 1004,
                    message: "目标设备不存在".to_string(),
                };
                let _ = state
                    .send_to_device(&from_device, &serde_json::to_string(&error).unwrap())
                    .await;
            }
        }

        SignalMessage::ConnectResponse {
            session_id,
            accepted: _,
            ..
        } => {
            // 获取会话信息并转发响应
            if let Some((controller, _controlled)) = state.active_sessions.read().await.get(&session_id) {
                // 转发到控制端
                let _ = state.send_to_device(controller, text).await;
            }
        }

        SignalMessage::Offer {
            to_device, ..
        }
        | SignalMessage::Answer {
            to_device, ..
        }
        | SignalMessage::Candidate {
            to_device, ..
        } => {
            // 转发信令消息
            let _ = state.send_to_device(&to_device, text).await;
        }

        SignalMessage::SessionEnd { session_id } => {
            let session_repo = SessionRepository::new(&state.db);
            let _ = session_repo.end_session(&session_id).await;
        }

        SignalMessage::SetAllowRemote { device_id, allow } => {
            let device_repo = DeviceRepository::new(&state.db);
            let _ = device_repo.set_allow_remote(&device_id, allow).await;
        }

        SignalMessage::RefreshCode { device_id } => {
            let device_repo = DeviceRepository::new(&state.db);
            if let Ok(new_code) = device_repo.update_verification_code(&device_id).await {
                let response = SignalMessage::CodeRefreshed {
                    device_id: device_id.clone(),
                    verification_code: new_code,
                };
                let _ = state
                    .send_to_device(&device_id, &serde_json::to_string(&response).unwrap())
                    .await;
            }
        }

        // 转发屏幕帧到目标设备（浏览器）
        SignalMessage::ScreenFrame { to_device, .. } => {
            let _ = state.send_to_device(&to_device, text).await;
        }

        // 转发鼠标事件到目标设备（客户端）
        SignalMessage::MouseEvent { to_device, .. } => {
            let _ = state.send_to_device(&to_device, text).await;
        }

        // 转发键盘事件到目标设备（客户端）
        SignalMessage::KeyboardEvent { to_device, .. } => {
            let _ = state.send_to_device(&to_device, text).await;
        }

        // 转发流媒体设置到目标设备（客户端）
        SignalMessage::StreamSettings { to_device, .. } => {
            let _ = state.send_to_device(&to_device, text).await;
        }

        // 转发获取显示器列表请求
        SignalMessage::GetDisplays { to_device, .. } => {
            let _ = state.send_to_device(&to_device, text).await;
        }

        // 转发显示器列表响应
        SignalMessage::DisplaysList { to_device, ref displays, .. } => {
            tracing::info!("★★★ 转发显示器列表到 {}, 共 {} 个显示器", to_device, displays.len());
            let result = state.send_to_device(&to_device, text).await;
            tracing::info!("★★★ 转发结果: {}", result);
        }

        // 转发切换显示器请求
        SignalMessage::SwitchDisplay { to_device, .. } => {
            let _ = state.send_to_device(&to_device, text).await;
        }

        // P2P 地址交换 - 转发给对方
        SignalMessage::P2PExchange { to_device, .. } => {
            tracing::info!("★★★ P2P 地址交换: 转发到 {}", to_device);
            let _ = state.send_to_device(&to_device, text).await;
        }

        // P2P 连接就绪 - 转发给对方
        SignalMessage::P2PReady { to_device, .. } => {
            tracing::info!("★★★ P2P 连接就绪: 转发到 {}", to_device);
            let _ = state.send_to_device(&to_device, text).await;
        }

        // P2P 回退 - 转发给对方
        SignalMessage::P2PFallback { to_device, .. } => {
            tracing::info!("★★★ P2P 回退到服务器中转: 转发到 {}", to_device);
            let _ = state.send_to_device(&to_device, text).await;
        }

        _ => {}
    }
}

/// 远程控制WebSocket处理器（浏览器端）
pub async fn remote_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    Path(device_id): Path<String>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_remote_socket(socket, state, device_id))
}

async fn handle_remote_socket(socket: WebSocket, state: Arc<AppState>, device_id: String) {
    let (mut sender, mut receiver) = socket.split();

    // 检查目标设备是否在线
    if !state.is_device_online(&device_id).await {
        let _ = sender
            .send(Message::Text(
                serde_json::json!({
                    "type": "error",
                    "message": "目标设备不在线"
                })
                .to_string(),
            ))
            .await;
        return;
    }

    let session_id = Uuid::new_v4().to_string();
    let browser_device_id = format!("browser_{}", Uuid::new_v4());

    // 创建浏览器端的临时连接
    let (tx, mut rx) = mpsc::channel::<String>(32);
    state
        .add_connection(browser_device_id.clone(), None, tx)
        .await;

    // 发送连接请求
    let connect_req = SignalMessage::ConnectRequest {
        session_id: session_id.clone(),
        from_device: browser_device_id.clone(),
        to_device: device_id.clone(),
        verification_code: String::new(), // 浏览器端需要管理员权限，跳过验证码
    };

    // 打印当前所有在线设备
    {
        let connections = state.connections.read().await;
        let online_devices: Vec<_> = connections.keys().collect();
        tracing::info!("★★★ 当前在线设备列表: {:?}", online_devices);
        tracing::info!("★★★ 尝试发送ConnectRequest到设备: {}", device_id);
    }

    let send_result = state
        .send_to_device(&device_id, &serde_json::to_string(&connect_req).unwrap())
        .await;
    tracing::info!("★★★ 发送ConnectRequest结果: {}", send_result);

    // 发送任务
    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if sender.send(Message::Text(msg)).await.is_err() {
                break;
            }
        }
    });

    // 接收任务
    let state_clone = state.clone();
    let browser_id_for_recv = browser_device_id.clone();
    let target_device_id = device_id.clone();
    let recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = receiver.next().await {
            if let Message::Text(text) = msg {
                // 打印调试信息（仅非移动事件）
                if !text.contains("\"event_type\":\"move\"") {
                    tracing::info!("★★★ [Remote] 收到消息: {}", &text[..text.len().min(200)]);
                }
                
                // 替换 from_device 为正确的浏览器ID
                let modified_text = if let Ok(mut json) = serde_json::from_str::<serde_json::Value>(&text) {
                    if let Some(obj) = json.as_object_mut() {
                        if obj.get("from_device").map(|v| v.as_str()) == Some(Some("browser")) {
                            obj.insert("from_device".to_string(), serde_json::Value::String(browser_id_for_recv.clone()));
                        }
                    }
                    serde_json::to_string(&json).unwrap_or(text)
                } else {
                    text
                };
                
                // 转发消息到目标设备
                let result = state_clone.send_to_device(&target_device_id, &modified_text).await;
                if !text.contains("\"event_type\":\"move\"") {
                    tracing::info!("★★★ [Remote] 转发到 {} 结果: {}", target_device_id, result);
                }
            }
        }
    });

    tokio::select! {
        _ = send_task => {},
        _ = recv_task => {},
    }

    // 发送 SessionEnd 消息给客户端，停止屏幕流
    let session_end = SignalMessage::SessionEnd {
        session_id: session_id.clone(),
    };
    let _ = state.send_to_device(&device_id, &serde_json::to_string(&session_end).unwrap()).await;
    tracing::info!("★★★ 浏览器断开连接，发送SessionEnd到设备: {}", device_id);

    // 清理
    state.remove_connection(&browser_device_id).await;
}
