use crate::{LogLine, ProgramId, ProgramStatus};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub struct AgentId(pub String);

impl std::fmt::Display for AgentId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentMetadata {
    pub hostname: String,
    pub cores: u32,
    pub memory_bytes: u64,
    pub cpu_features: Vec<String>,
    pub os: String,
    pub arch: String,
    pub agent_version: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentHeartbeat {
    pub agent_id: AgentId,
    pub timestamp: u64,
    pub metadata: Option<AgentMetadata>,
    pub programs: Vec<ProgramStatus>,
    pub logs: Vec<LogLine>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeartbeatResponse {
    pub config: crate::AgentConfig,
    pub commands: Vec<AgentCommand>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentCommand {
    Deploy {
        program_id: ProgramId,
        version_hash: String,
    },
    Start {
        program_id: ProgramId,
    },
    Stop {
        program_id: ProgramId,
    },
    Rebuild {
        program_id: ProgramId,
    },
    /// Send work unit to a specific program
    SendWork {
        program_id: ProgramId,
        job_id: String,
        work_unit_id: String,
        #[serde(with = "crate::job::serde_bytes")]
        data: Vec<u8>,
    },
    /// Cancel a job - stop processing work units for this job
    CancelJob {
        job_id: String,
    },
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AgentStatus {
    Healthy,
    Unhealthy,
    Dead,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentInfo {
    pub agent_id: AgentId,
    pub hostname: String,
    pub cores: u32,
    pub memory_bytes: u64,
    pub cpu_features: Vec<String>,
    pub os: String,
    pub arch: String,
    pub agent_version: String,
    pub status: AgentStatus,
    pub last_heartbeat: u64,
    pub programs: Vec<ProgramStatus>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentListItem {
    pub agent_id: AgentId,
    pub hostname: String,
    pub cores: u32,
    pub memory_bytes: u64,
    pub status: AgentStatus,
    pub last_heartbeat: u64,
    pub programs: Vec<String>,
}

/// Messages sent from Agent to CNC over WebSocket
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentMessage {
    /// Initial registration/authentication
    Register {
        agent_id: AgentId,
        metadata: AgentMetadata,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        token: Option<String>,
    },
    /// Periodic status update
    Status { programs: Vec<ProgramStatus> },
    /// Work unit result
    Result {
        job_id: String,
        work_unit_id: String,
        #[serde(with = "crate::job::serde_bytes")]
        data: Vec<u8>,
    },
    /// Log lines from agent/programs
    Logs { lines: Vec<LogLine> },
    /// Deploy progress update
    DeployProgress {
        program_id: ProgramId,
        status: DeployStatus,
        message: Option<String>,
    },
}

/// Status of a deploy operation on an agent
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum DeployStatus {
    /// Downloading program archive
    Downloading,
    /// Extracting archive
    Extracting,
    /// Building program
    Building,
    /// Build completed successfully
    Completed,
    /// Build failed
    Failed,
}

/// Messages sent from CNC to Agent over WebSocket
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CncMessage {
    /// Registration accepted, here's your config
    Registered { config: crate::AgentConfig },
    /// Command to execute (nested, not flattened)
    Command { command: AgentCommand },
    /// Ping to keep connection alive
    Ping,
}

// ============ CLI WebSocket Messages ============

/// Messages sent from CLI to CNC over WebSocket
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CliMessage {
    /// Subscribe to deploy updates for a program
    SubscribeDeploy { program_id: ProgramId },
    /// Unsubscribe from deploy updates
    UnsubscribeDeploy { program_id: ProgramId },
    /// Subscribe to all agent status updates
    SubscribeAgents,
}

/// Messages sent from CNC to CLI over WebSocket
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CliUpdate {
    /// Deploy progress update from an agent
    DeployProgress {
        agent_id: AgentId,
        program_id: ProgramId,
        status: DeployStatus,
        message: Option<String>,
    },
    /// Agent status update
    AgentStatus {
        agent_id: AgentId,
        status: AgentStatus,
        programs: Vec<ProgramStatus>,
    },
    /// Agent connected
    AgentConnected {
        agent_id: AgentId,
        metadata: AgentMetadata,
    },
    /// Agent disconnected
    AgentDisconnected { agent_id: AgentId },
    /// Deploy completed on all agents
    DeployCompleted {
        program_id: ProgramId,
        succeeded: Vec<AgentId>,
        failed: Vec<(AgentId, String)>,
    },
}
