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

/// Status of a remote desktop session.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum SessionStatus {
    /// Session created, waiting for an agent to connect.
    Waiting,
    /// Agent is streaming display/audio.
    Active,
    /// Session ended or agent disconnected.
    Disconnected,
}

impl std::fmt::Display for SessionStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SessionStatus::Waiting => write!(f, "waiting"),
            SessionStatus::Active => write!(f, "active"),
            SessionStatus::Disconnected => write!(f, "disconnected"),
        }
    }
}

/// A remote desktop session. Each session corresponds to one VM / remote machine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    pub id: String,
    pub status: SessionStatus,
    pub created_at: DateTime<Utc>,
    pub display_active: bool,
    pub audio_active: bool,
    /// Resolution reported by the agent (e.g. "1920x1080").
    pub resolution: Option<String>,
    /// Arbitrary key-value metadata attached to this session.
    pub metadata: HashMap<String, String>,
}

impl Session {
    pub fn new(id: String) -> Self {
        Self {
            id,
            status: SessionStatus::Waiting,
            created_at: Utc::now(),
            display_active: false,
            audio_active: false,
            resolution: None,
            metadata: HashMap::new(),
        }
    }
}

/// Connection metadata for a registered desktop agent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentConnection {
    pub agent_id: String,
    pub session_id: String,
    pub connected_at: DateTime<Utc>,
    pub ip_address: String,
    pub display_active: bool,
    pub audio_active: bool,
}

/// Client type for WebSocket connections.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ClientType {
    /// Browser-based viewer watching the remote desktop.
    Viewer,
    /// Native agent executable running inside the VM.
    Agent,
}

// ── WebSocket message types ──────────────────────────────────────────────────

/// Every WebSocket frame carries a JSON envelope with a `msg_type` discriminator.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "msg_type", rename_all = "snake_case")]
pub enum WsMessage {
    // ── Agent → Server ───────────────────────────────────────────────────────
    /// A single display frame, base64-encoded JPEG or PNG.
    DisplayFrame { session_id: String, data: String, timestamp: Option<f64> },
    /// An audio chunk, base64-encoded PCM / Opus / AAC.
    AudioFrame { session_id: String, data: String, timestamp: Option<f64> },
    /// Agent announces its capabilities and info on connect.
    AgentInfo { session_id: String, agent_id: String, resolution: Option<String> },
    /// Keep-alive ping from agent.
    Heartbeat,

    // ── Viewer → Server ──────────────────────────────────────────────────────
    /// Viewer sends a HUD (heads-up display) command to the agent.
    HudCommand { session_id: String, command: String, params: serde_json::Value },
    /// Viewer requests display resize.
    Resize { session_id: String, width: u32, height: u32 },

    // ── Server → Viewer ──────────────────────────────────────────────────────
    /// Server relays a display frame to viewers.
    FrameBroadcast { data: String, content_type: String },
    /// Server relays an audio chunk to viewers.
    AudioBroadcast { data: String, content_type: String },
    /// Session state update pushed to viewers.
    SessionUpdate { session_id: String, status: SessionStatus, resolution: Option<String> },

    // ── Server → Agent ───────────────────────────────────────────────────────
    /// Server forwards a HUD command from a viewer to the agent.
    ForwardHudCommand { command: String, params: serde_json::Value },
    /// Server forwards a resize request from a viewer.
    ForwardResize { width: u32, height: u32 },
    /// Server tells the agent to start / stop streaming.
    StreamControl { action: String },

    // ── Generic ──────────────────────────────────────────────────────────────
    Error { message: String },
    Ack { message: String },
}

/// Response body for health-check and generic API responses.
#[derive(Debug, Serialize)]
pub struct ApiResponse<T: Serialize> {
    pub ok: bool,
    pub data: Option<T>,
    pub error: Option<String>,
}

impl<T: Serialize> ApiResponse<T> {
    pub fn ok(data: T) -> Self {
        Self { ok: true, data: Some(data), error: None }
    }

    pub fn err(msg: impl Into<String>) -> Self {
        Self { ok: false, data: None, error: Some(msg.into()) }
    }
}

/// Health-check payload.
#[derive(Debug, Serialize)]
pub struct HealthInfo {
    pub status: String,
    pub uptime_secs: u64,
    pub active_sessions: usize,
    pub connected_agents: usize,
    pub connected_viewers: usize,
    pub version: String,
}