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

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

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

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

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateJobRequest {
    pub program_id: ProgramId,
    #[serde(with = "serde_bytes_vec")]
    pub work_units: Vec<Vec<u8>>,
}

mod serde_bytes_vec {
    use base64::{engine::general_purpose::STANDARD, Engine as _};
    use serde::{Deserialize, Deserializer, Serializer};

    pub fn serialize<S>(data: &[Vec<u8>], serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        use serde::ser::SerializeSeq;
        let mut seq = serializer.serialize_seq(Some(data.len()))?;
        for item in data {
            seq.serialize_element(&STANDARD.encode(item))?;
        }
        seq.end()
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<Vec<u8>>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let strings: Vec<String> = Vec::deserialize(deserializer)?;
        strings
            .into_iter()
            .map(|s| STANDARD.decode(&s).map_err(serde::de::Error::custom))
            .collect()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateJobResponse {
    pub job_id: JobId,
    pub status: JobState,
    pub work_units_total: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JobStatus {
    pub job_id: JobId,
    pub program_id: ProgramId,
    pub status: JobState,
    pub work_units_total: usize,
    pub work_units_completed: usize,
    pub work_units_pending: usize,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum JobState {
    Queued,
    Running,
    Completed,
    Failed,
    Cancelled,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkUnitResult {
    pub work_unit_id: WorkUnitId,
    pub executor_id: String,
    #[serde(with = "serde_bytes")]
    pub data: Vec<u8>,
}

pub mod serde_bytes {
    use base64::{engine::general_purpose::STANDARD, Engine as _};
    use serde::{Deserialize, Deserializer, Serializer};

    pub fn serialize<S>(data: &[u8], serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&STANDARD.encode(data))
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        STANDARD.decode(&s).map_err(serde::de::Error::custom)
    }
}

// ============ Dispatcher WebSocket Messages ============

/// Scheduling strategy for work unit distribution
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Default)]
#[serde(rename_all = "snake_case")]
pub enum SchedulerStrategy {
    /// Round-robin distribution (default, simple)
    #[default]
    RoundRobin,
    /// Flow-control based on worker capacity with buffer factor
    /// Agents report worker count, CNC sends work based on capacity
    FlowControl {
        /// Buffer factor (e.g., 1.2 = 20% extra buffered work)
        #[serde(default = "default_buffer_factor")]
        buffer_factor: f32,
    },
}

fn default_buffer_factor() -> f32 {
    1.2
}

/// Messages sent from Dispatcher/SDK to CNC over WebSocket
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum DispatcherMessage {
    /// Create a new job with optional scheduler strategy
    CreateJob {
        program_id: ProgramId,
        #[serde(default)]
        scheduler: SchedulerStrategy,
    },
    /// Stream a batch of work units for the current job
    WorkUnits {
        #[serde(with = "serde_bytes_vec")]
        units: Vec<Vec<u8>>,
    },
    /// Signal end of work units for current job
    EndJob,
    /// Cancel a job
    CancelJob { job_id: String },
    /// Query job status
    GetStatus { job_id: String },
}

/// Messages sent from CNC to Dispatcher/SDK over WebSocket
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum DispatcherResponse {
    /// Job created, streaming can begin
    JobCreated { job_id: String },
    /// Work units received acknowledgment
    WorkUnitsReceived { count: usize },
    /// A result is ready
    Result {
        work_unit_id: String,
        executor_id: String,
        #[serde(with = "serde_bytes")]
        data: Vec<u8>,
    },
    /// Job completed successfully
    JobCompleted {
        job_id: String,
        total_work_units: usize,
        total_results: usize,
    },
    /// Job failed - agent died, work unit failed, etc.
    JobFailed {
        job_id: String,
        reason: String,
        work_units_completed: usize,
        work_units_total: usize,
    },
    /// Job was cancelled
    JobCancelled { job_id: String },
    /// Job status response
    Status(JobStatus),
    /// Error occurred
    Error { message: String },
}
