/// Godot RL Agents protocol client
///
/// Implements the godot-rl-agents WebSocket protocol for communicating with Godot.
/// Protocol: https://github.com/edbeeching/godot_rl_agents
use anyhow::{Context, Result};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tokio::net::TcpStream;
use tokio_tungstenite::{
    connect_async_with_config,
    tungstenite::protocol::{Message, WebSocketConfig},
    WebSocketStream,
};
use tracing::{debug, info, warn};

type WsStream = WebSocketStream<tokio_tungstenite::MaybeTlsStream<TcpStream>>;

// Re-export ImageFormat from video package
pub use mecha10_video::ImageFormat;

/// Create WebSocket configuration with appropriate buffer sizes for camera streaming
///
/// Buffer sizes are calculated to handle high-resolution camera frames without blocking.
/// Formula: frame_size = width × height × 4 (RGBA bytes)
/// - 640×480 frame = 1.2 MB
/// - Buffer = 3 MB (fits 2-3 frames at max resolution)
fn create_websocket_config() -> WebSocketConfig {
    WebSocketConfig {
        // Maximum message size: 64 MB (default, sufficient for any reasonable frame)
        max_message_size: Some(64 * 1024 * 1024),
        // Maximum frame size: 16 MB (default, sufficient for up to 4K frames)
        max_frame_size: Some(16 * 1024 * 1024),
        // Accept unmasked frames from server (Godot sends unmasked)
        accept_unmasked_frames: true,
        // Write buffer size: 3 MB (handles 2-3 frames at 640×480)
        // CRITICAL: This prevents blocking when frames arrive faster than network can send
        write_buffer_size: 3 * 1024 * 1024,
        // Max write buffer size: 6 MB (allows some burst tolerance)
        max_write_buffer_size: 6 * 1024 * 1024,
        ..Default::default()
    }
}

/// Godot RL Agents protocol messages
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum GodotMessage {
    #[serde(rename = "handshake")]
    Handshake {
        major_version: String,
        minor_version: String,
    },
    #[serde(rename = "env_info")]
    EnvInfo {
        n_agents: u32,
        action_space: ActionSpace,
        obs_space: ObservationSpace,
    },
    #[serde(rename = "step")]
    Step { reward: f32, done: bool, obs: Vec<f32> },
    #[serde(rename = "reset")]
    Reset { obs: Vec<f32> },
    #[serde(rename = "action")]
    Action {
        action: Vec<f32>,
        /// Timestamp in microseconds (when command was created)
        #[serde(skip_serializing_if = "Option::is_none")]
        timestamp_us: Option<u64>,
    },
    #[serde(rename = "reset_env")]
    ResetEnv,
    #[serde(rename = "close")]
    Close,
    /// Camera configuration message sent after handshake
    #[serde(rename = "camera_config")]
    CameraConfig {
        /// Camera resolution width
        width: u32,
        /// Camera resolution height
        height: u32,
        /// Camera frames per second
        fps: u32,
    },
    /// Camera data (raw bytes) - BINARY FORMAT
    /// This variant is not serializable and is only created from binary WebSocket messages
    #[serde(skip)]
    CameraDataBinary {
        /// Camera name/ID
        camera_id: String,
        /// Raw image bytes (JPEG or RGB depending on format)
        image_bytes: Vec<u8>,
        /// Image format (JPEG or RGB)
        format: ImageFormat,
        /// Image dimensions
        width: u32,
        height: u32,
        /// Timestamp in microseconds
        timestamp: u64,
    },
    /// IMU data (binary format) - BINARY FORMAT
    /// This variant is not serializable and is only created from binary WebSocket messages
    #[serde(skip)]
    ImuDataBinary {
        /// IMU name/ID
        imu_id: String,
        /// Timestamp in microseconds
        timestamp: u64,
        /// Orientation (Euler angles in radians)
        roll: f32,
        pitch: f32,
        yaw: f32,
        /// Angular velocity (rad/s)
        angular_velocity_x: f32,
        angular_velocity_y: f32,
        angular_velocity_z: f32,
        /// Linear acceleration (m/s²)
        linear_acceleration_x: f32,
        linear_acceleration_y: f32,
        linear_acceleration_z: f32,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActionSpace {
    pub shape: Vec<usize>,
    pub dtype: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ObservationSpace {
    pub shape: Vec<usize>,
    pub dtype: String,
}

/// Godot RL Agents client
pub struct GodotClient {
    ws: WsStream,
    env_info: Option<EnvInfo>,
}

#[allow(dead_code)]
#[derive(Debug, Clone)]
pub struct EnvInfo {
    pub n_agents: u32,
    pub action_space: ActionSpace,
    pub obs_space: ObservationSpace,
}

impl GodotClient {
    /// Connect to Godot RL Agents server
    ///
    /// # Arguments
    /// * `url` - WebSocket URL (e.g., "ws://localhost:11008")
    pub async fn connect(url: &str) -> Result<Self> {
        info!("🔌 Connecting to Godot at {}...", url);

        let config = create_websocket_config();
        let (ws, _) = connect_async_with_config(url, Some(config), false)
            .await
            .context("Failed to connect to Godot")?;

        info!("✅ Connected to Godot (write_buffer=3MB, max_write_buffer=6MB)");

        Ok(Self { ws, env_info: None })
    }

    /// Connect to Godot camera server without handshake
    ///
    /// # Arguments
    /// * `url` - WebSocket URL (e.g., "ws://localhost:11009")
    pub async fn connect_no_handshake(url: &str) -> Result<Self> {
        info!("🔌 Connecting to Godot camera server at {}...", url);

        let config = create_websocket_config();
        let (ws, _) = connect_async_with_config(url, Some(config), false)
            .await
            .context("Failed to connect to Godot camera server")?;

        info!("✅ Connected to Godot camera server (write_buffer=3MB, max_write_buffer=6MB)");

        Ok(Self { ws, env_info: None })
    }

    /// Perform handshake with Godot
    pub async fn handshake(&mut self) -> Result<()> {
        debug!("🤝 Performing handshake...");

        // Send handshake
        let handshake = GodotMessage::Handshake {
            major_version: "1".to_string(),
            minor_version: "0".to_string(),
        };

        self.send_message(&handshake).await?;

        // Wait for env_info response
        match self.receive_message().await? {
            GodotMessage::EnvInfo {
                n_agents,
                action_space,
                obs_space,
            } => {
                self.env_info = Some(EnvInfo {
                    n_agents,
                    action_space: action_space.clone(),
                    obs_space: obs_space.clone(),
                });

                info!(
                    "✅ Handshake complete - {} agents, action shape: {:?}, obs shape: {:?}",
                    n_agents, action_space.shape, obs_space.shape
                );

                Ok(())
            }
            msg => {
                warn!("Expected EnvInfo, got {:?}", msg);
                anyhow::bail!("Unexpected message during handshake")
            }
        }
    }

    /// Send camera configuration to Godot
    ///
    /// Should be called after handshake to configure camera resolution and FPS.
    pub async fn send_camera_config(&mut self, width: u32, height: u32, fps: u32) -> Result<()> {
        info!("📷 Sending camera config: {}x{} @ {} FPS", width, height, fps);

        let config_msg = GodotMessage::CameraConfig { width, height, fps };
        self.send_message(&config_msg).await?;

        info!("✅ Camera config sent");
        Ok(())
    }

    // REMOVED: reset() method - never used
    // REMOVED: step() method - never used (use step_with_timestamp directly)

    /// Send action with timestamp and receive step result
    pub async fn step_with_timestamp(&mut self, action: Vec<f32>, timestamp_us: Option<u64>) -> Result<StepResult> {
        use tokio::time::{timeout, Duration};

        debug!("🎮 Sending action: {:?} (timestamp: {:?})", action, timestamp_us);

        let action_msg = GodotMessage::Action { action, timestamp_us };
        self.send_message(&action_msg).await?;

        // Wait for Step response with 5 second timeout
        // Note: Camera data now comes from dedicated camera WebSocket (port 11009)
        let step_future = async {
            #[allow(clippy::never_loop)]
            loop {
                match self.receive_message().await? {
                    GodotMessage::Step { reward, done, obs } => {
                        debug!(
                            "✅ Step complete - reward: {}, done: {}, obs size: {}",
                            reward,
                            done,
                            obs.len()
                        );

                        return Ok(StepResult {
                            observation: obs,
                            reward,
                            done,
                            info: HashMap::new(),
                        });
                    }
                    msg => {
                        warn!("Expected Step, got unexpected message: {:?}", msg);
                        anyhow::bail!("Unexpected message during step")
                    }
                }
            }
        };

        // Add 5 second timeout to prevent hanging indefinitely
        match timeout(Duration::from_secs(5), step_future).await {
            Ok(result) => result,
            Err(_) => {
                warn!("⏱️  Step timed out after 5 seconds - Godot may be stuck or overloaded");
                anyhow::bail!("Step timeout - no response from Godot within 5 seconds")
            }
        }
    }

    /// Close the WebSocket connection gracefully
    pub async fn close(&mut self) -> Result<()> {
        debug!("🔌 Closing Godot WebSocket connection");
        self.ws.close(None).await?;
        Ok(())
    }

    /// Send message to Godot
    async fn send_message(&mut self, msg: &GodotMessage) -> Result<()> {
        let json = serde_json::to_string(msg)?;
        self.ws.send(Message::Text(json)).await?;
        Ok(())
    }

    /// Receive message from Godot
    async fn receive_message(&mut self) -> Result<GodotMessage> {
        while let Some(msg) = self.ws.next().await {
            match msg? {
                Message::Text(text) => {
                    let godot_msg: GodotMessage = serde_json::from_str(&text)?;
                    return Ok(godot_msg);
                }
                Message::Close(_) => {
                    anyhow::bail!("Godot closed connection");
                }
                _ => {
                    // Skip ping/pong/binary messages
                    continue;
                }
            }
        }

        anyhow::bail!("Godot connection closed unexpectedly")
    }

    /// Receive message with event-driven blocking (await until message arrives)
    ///
    /// This is the preferred method for camera streaming - it properly awaits
    /// the WebSocket without polling, minimizing CPU usage and latency.
    pub async fn receive_message_async(&mut self) -> Result<Option<GodotMessage>> {
        // Directly await on WebSocket stream (no timeout - truly event-driven!)
        match self.ws.next().await {
            Some(Ok(msg)) => match msg {
                Message::Text(text) => {
                    debug!("📥 Received Text message: {} bytes", text.len());
                    let godot_msg: GodotMessage = serde_json::from_str(&text)?;
                    Ok(Some(godot_msg))
                }
                Message::Binary(data) => {
                    debug!(
                        "📥 Received Binary message: {} bytes, first byte: 0x{:02X}",
                        data.len(),
                        data.first().unwrap_or(&0)
                    );
                    // Check version byte to determine message type
                    if data.is_empty() {
                        warn!("Received empty binary message");
                        return Ok(None);
                    }

                    let version = data[0];
                    match version {
                        0x01..=0x04 => {
                            // Camera data (versions 0x01-0x04, including JPEG compression)
                            match Self::parse_binary_camera_message(&data) {
                                Ok(camera_msg) => Ok(Some(camera_msg)),
                                Err(e) => {
                                    warn!("Failed to parse binary camera message: {}", e);
                                    Ok(None)
                                }
                            }
                        }
                        0x10 => {
                            // IMU data (version 0x10)
                            match Self::parse_binary_imu_message(&data) {
                                Ok(imu_msg) => Ok(Some(imu_msg)),
                                Err(e) => {
                                    warn!("Failed to parse binary IMU message: {}", e);
                                    Ok(None)
                                }
                            }
                        }
                        _ => {
                            warn!("Unsupported binary protocol version: 0x{:02X}", version);
                            Ok(None)
                        }
                    }
                }
                Message::Close(frame) => {
                    warn!("🔌 Godot WebSocket closed: {:?}", frame);
                    anyhow::bail!("Godot closed connection");
                }
                Message::Ping(data) => {
                    debug!("🏓 Received WebSocket Ping, sending Pong");
                    // Respond to Ping with Pong to keep connection alive
                    if let Err(e) = self.ws.send(Message::Pong(data)).await {
                        warn!("Failed to send Pong response: {}", e);
                        return Err(e.into());
                    }
                    Ok(None)
                }
                Message::Pong(_) => {
                    debug!("🏓 Received WebSocket Pong");
                    Ok(None)
                }
                msg_type => {
                    debug!("📭 Received other WebSocket message type: {:?}", msg_type);
                    Ok(None)
                }
            },
            Some(Err(e)) => {
                // WebSocket error
                Err(e.into())
            }
            None => {
                // Stream ended
                warn!("WebSocket stream ended");
                anyhow::bail!("Godot connection closed unexpectedly")
            }
        }
    }

    /// Parse binary camera data message
    ///
    /// Binary protocol format:
    ///   1 byte: version (0x01 = JPEG, 0x02 = RGB)
    ///   4 bytes: camera_id length (uint32 little-endian)
    ///   N bytes: camera_id (UTF-8)
    ///   4 bytes: width (uint32 little-endian)
    ///   4 bytes: height (uint32 little-endian)
    ///   8 bytes: timestamp (uint64 little-endian)
    ///   Rest: Image data (JPEG or RGB depending on version)
    fn parse_binary_camera_message(data: &[u8]) -> Result<GodotMessage> {
        use std::convert::TryInto;

        let mut offset = 0;

        // Check minimum size (1 + 4 + 4 + 4 + 8 = 21 bytes minimum)
        if data.len() < 21 {
            anyhow::bail!("Binary message too short: {} bytes", data.len());
        }

        // Parse version and determine format
        let version = data[offset];
        offset += 1;
        let format = match version {
            0x01 => ImageFormat::Jpeg,
            0x02 => ImageFormat::Rgb,
            0x03 => ImageFormat::Rgb, // v3 = RGBA, but we treat as Rgb (encoder handles conversion)
            0x04 => ImageFormat::Jpeg, // v4 = JPEG compression for higher resolutions
            _ => anyhow::bail!("Unsupported binary protocol version: {}", version),
        };

        // Parse camera_id length
        let camera_id_len = u32::from_le_bytes(data[offset..offset + 4].try_into()?) as usize;
        offset += 4;

        // Parse camera_id
        if offset + camera_id_len > data.len() {
            anyhow::bail!("Invalid camera_id length: {}", camera_id_len);
        }
        let camera_id = String::from_utf8(data[offset..offset + camera_id_len].to_vec())?;
        offset += camera_id_len;

        // Parse width
        let width = u32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;

        // Parse height
        let height = u32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;

        // Parse timestamp
        let timestamp = u64::from_le_bytes(data[offset..offset + 8].try_into()?);
        offset += 8;

        // Rest is image data (JPEG or RGB)
        let image_bytes = data[offset..].to_vec();

        debug!(
            "📷 Parsed binary camera message: camera={}, format={:?}, size={}KB, {}x{}",
            camera_id,
            format,
            image_bytes.len() / 1024,
            width,
            height
        );

        Ok(GodotMessage::CameraDataBinary {
            camera_id,
            image_bytes,
            format,
            width,
            height,
            timestamp,
        })
    }

    /// Parse binary IMU data message
    ///
    /// Binary protocol format:
    ///   1 byte: version (0x10 = IMU data)
    ///   4 bytes: imu_id length (uint32 little-endian)
    ///   N bytes: imu_id (UTF-8)
    ///   8 bytes: timestamp (uint64 little-endian)
    ///   36 bytes: 9 floats (roll, pitch, yaw, ang_vel xyz, lin_acc xyz)
    fn parse_binary_imu_message(data: &[u8]) -> Result<GodotMessage> {
        use std::convert::TryInto;

        let mut offset = 0;

        // Check minimum size (1 + 4 + 8 + 36 = 49 bytes minimum)
        if data.len() < 49 {
            anyhow::bail!("IMU binary message too short: {} bytes", data.len());
        }

        // Parse version (already checked by caller, but verify)
        let version = data[offset];
        offset += 1;
        if version != 0x10 {
            anyhow::bail!("Invalid IMU message version: 0x{:02X}", version);
        }

        // Parse imu_id length
        let imu_id_len = u32::from_le_bytes(data[offset..offset + 4].try_into()?) as usize;
        offset += 4;

        // Parse imu_id
        if offset + imu_id_len > data.len() {
            anyhow::bail!("Invalid imu_id length: {}", imu_id_len);
        }
        let imu_id = String::from_utf8(data[offset..offset + imu_id_len].to_vec())?;
        offset += imu_id_len;

        // Parse timestamp (8 bytes)
        let timestamp = u64::from_le_bytes(data[offset..offset + 8].try_into()?);
        offset += 8;

        // Parse 9 floats (36 bytes total)
        if offset + 36 > data.len() {
            anyhow::bail!("IMU message too short for float data");
        }

        let roll = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;
        let pitch = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;
        let yaw = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;

        let angular_velocity_x = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;
        let angular_velocity_y = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;
        let angular_velocity_z = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;

        let linear_acceleration_x = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;
        let linear_acceleration_y = f32::from_le_bytes(data[offset..offset + 4].try_into()?);
        offset += 4;
        let linear_acceleration_z = f32::from_le_bytes(data[offset..offset + 4].try_into()?);

        debug!(
            "📡 Parsed binary IMU message: imu={}, timestamp={}, orientation=[{:.3}, {:.3}, {:.3}]",
            imu_id, timestamp, roll, pitch, yaw
        );

        Ok(GodotMessage::ImuDataBinary {
            imu_id,
            timestamp,
            roll,
            pitch,
            yaw,
            angular_velocity_x,
            angular_velocity_y,
            angular_velocity_z,
            linear_acceleration_x,
            linear_acceleration_y,
            linear_acceleration_z,
        })
    }
}

/// Result from a simulation step
#[derive(Debug, Clone)]
pub struct StepResult {
    pub observation: Vec<f32>,
    pub reward: f32,
    pub done: bool,
    #[allow(dead_code)]
    pub info: HashMap<String, String>,
}

impl Drop for GodotClient {
    fn drop(&mut self) {
        // Explicitly close WebSocket connection when GodotClient is dropped
        // This ensures Godot's ProtocolServer removes the client from connected_peers
        debug!("🔌 GodotClient: Closing WebSocket connection");

        // tokio_tungstenite's WebSocketStream implements Drop which sends Close frame
        // But we'll explicitly close to ensure it's sent before drop
        // Note: We can't use async in Drop, but the underlying stream will close
        // when dropped. The WebSocket protocol will send a Close frame automatically.
    }
}
