use anyhow::{Context, Result};
use eridani_hpc_common::{
    Config, DispatcherMessage, DispatcherResponse, ProgramId, SchedulerStrategy, WorkUnitId,
    WorkUnitResult,
};
use futures::{SinkExt, StreamExt};
use reqwest::Client;
use serde::{de::DeserializeOwned, Serialize};
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use tokio_tungstenite::{connect_async, tungstenite::Message};

/// A typed result that includes metadata and deserialized data
pub struct TypedResult<T> {
    pub work_unit_id: WorkUnitId,
    pub executor_id: String,
    pub data: T,
}

/// Re-export Config for dispatcher convenience
pub use eridani_hpc_common::Config as HpcConfig;

pub struct HpcClient {
    base_url: String,
    ws_url: String,
    token: Option<String>,
    client: Client,
}

impl HpcClient {
    pub fn new(base_url: impl Into<String>) -> Self {
        let base_url = base_url.into().trim().to_string();
        let ws_url = base_url
            .replace("http://", "ws://")
            .replace("https://", "wss://");
        let ws_url = format!("{}/api/dispatcher/ws", ws_url);

        Self {
            base_url,
            ws_url,
            token: None,
            client: Client::new(),
        }
    }

    /// Create client from config file (uses default profile)
    pub fn from_config() -> Self {
        Self::from_profile(None)
    }

    /// Create client from a specific profile in the config file
    pub fn from_profile(profile_name: Option<&str>) -> Self {
        let config = Config::load();
        let profile = config.profile(profile_name);

        let mut client = Self::new(profile.cnc_url());
        if let Some(token) = profile.cli_token() {
            client = client.with_token(token);
        } else {
            eprintln!("Warning: No cli_token found in config profile");
        }
        client
    }

    pub fn with_token(mut self, token: impl Into<String>) -> Self {
        self.token = Some(token.into());
        self
    }

    pub async fn connect(base_url: impl Into<String>) -> Result<Self> {
        let client = Self::new(base_url);
        // TODO: Could add a health check here
        Ok(client)
    }

    /// Log cluster status and check for available agents
    /// Returns Ok(status) if agents are online, Err if no agents available
    pub async fn require_agents(&self) -> Result<ClusterStatus> {
        let status = self.status().await?;

        println!("Cluster: {}", self.base_url);
        println!("  Agents online: {}", status.agents_online);
        println!("  Total cores:   {}", status.total_cores);
        println!(
            "  Total memory:  {:.1} GB",
            status.total_memory_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
        );

        if status.agents_online == 0 {
            anyhow::bail!("No agents online");
        }

        Ok(status)
    }

    /// Create a job builder for submitting work
    pub fn job(&self, program_id: impl Into<ProgramId>) -> JobBuilder<'_> {
        JobBuilder {
            client: self,
            program_id: program_id.into(),
            work_units: Vec::new(),
            scheduler: SchedulerStrategy::default(),
        }
    }

    /// Create a streaming job that connects via WebSocket
    pub async fn streaming_job(&self, program_id: impl Into<ProgramId>) -> Result<StreamingJob> {
        StreamingJob::new(
            &self.ws_url,
            self.token.as_deref(),
            program_id.into(),
            SchedulerStrategy::default(),
        )
        .await
    }

    pub async fn status(&self) -> Result<ClusterStatus> {
        let url = format!("{}/api/status", self.base_url);
        let mut req = self.client.get(&url);

        if let Some(token) = &self.token {
            req = req.bearer_auth(token);
        }

        let resp = req.send().await.context("Failed to get cluster status")?;
        let status = resp
            .json()
            .await
            .context("Failed to parse cluster status")?;

        Ok(status)
    }
}

pub struct JobBuilder<'a> {
    client: &'a HpcClient,
    program_id: ProgramId,
    work_units: Vec<Vec<u8>>,
    scheduler: SchedulerStrategy,
}

impl<'a> JobBuilder<'a> {
    /// Set raw work units (already serialized bytes)
    pub fn work_units(mut self, work_units: Vec<Vec<u8>>) -> Self {
        self.work_units = work_units;
        self
    }

    /// Add a raw work unit (already serialized bytes)
    pub fn work_unit(mut self, work_unit: Vec<u8>) -> Self {
        self.work_units.push(work_unit);
        self
    }

    /// Add a typed work unit (auto-serializes with bincode)
    pub fn add<T: Serialize>(mut self, item: &T) -> Result<Self> {
        let bytes = bincode::serialize(item)?;
        self.work_units.push(bytes);
        Ok(self)
    }

    /// Add multiple typed work units (auto-serializes with bincode)
    pub fn add_all<T: Serialize>(mut self, items: impl IntoIterator<Item = T>) -> Result<Self> {
        for item in items {
            let bytes = bincode::serialize(&item)?;
            self.work_units.push(bytes);
        }
        Ok(self)
    }

    /// Use round-robin scheduler (default)
    pub fn round_robin(mut self) -> Self {
        self.scheduler = SchedulerStrategy::RoundRobin;
        self
    }

    /// Use flow-control scheduler with specified buffer factor
    ///
    /// Flow-control distributes work based on agent worker capacity.
    /// The buffer_factor adds extra work to prevent idle time:
    /// - 1.0 = exactly worker count
    /// - 1.2 = 20% extra buffered (recommended)
    /// - 1.5 = 50% extra buffered
    pub fn flow_control(mut self, buffer_factor: f32) -> Self {
        self.scheduler = SchedulerStrategy::FlowControl { buffer_factor };
        self
    }

    /// Submit the job via WebSocket for streaming results
    pub async fn submit(self) -> Result<StreamingJob> {
        let mut job = StreamingJob::new(
            &self.client.ws_url,
            self.client.token.as_deref(),
            self.program_id,
            self.scheduler,
        )
        .await?;

        // Send all work units
        job.send_work_units(self.work_units).await?;

        // Signal end of work units
        job.end_submission().await?;

        Ok(job)
    }
}

/// Job failure information
#[derive(Debug, Clone)]
pub struct JobError {
    pub job_id: String,
    pub reason: String,
    pub work_units_completed: usize,
    pub work_units_total: usize,
}

impl std::fmt::Display for JobError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Job {} failed: {} ({}/{} completed)",
            self.job_id, self.reason, self.work_units_completed, self.work_units_total
        )
    }
}

impl std::error::Error for JobError {}

/// Internal message for the result receiver
enum JobEvent {
    Result(WorkUnitResult),
    Failed(JobError),
    Cancelled(String),
    Completed,
}

/// A job that streams work units and results over WebSocket
pub struct StreamingJob {
    job_id: Option<String>,
    ws_sender: Arc<
        Mutex<
            futures::stream::SplitSink<
                tokio_tungstenite::WebSocketStream<
                    tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
                >,
                Message,
            >,
        >,
    >,
    event_rx: mpsc::UnboundedReceiver<JobEvent>,
    work_units_submitted: usize,
    results_received: usize,
    /// If the job failed, this holds the error info
    failure: Option<JobError>,
}

impl StreamingJob {
    async fn new(
        ws_url: &str,
        token: Option<&str>,
        program_id: ProgramId,
        scheduler: SchedulerStrategy,
    ) -> Result<Self> {
        // Append token as query param if provided
        let url = match token {
            Some(t) => format!("{}?token={}", ws_url, t),
            None => ws_url.to_string(),
        };

        let (ws_stream, _) = connect_async(&url)
            .await
            .context("Failed to connect to CNC WebSocket")?;

        let (ws_sender, mut ws_receiver) = ws_stream.split();
        let ws_sender = Arc::new(Mutex::new(ws_sender));

        // Channel for job events (results, failures, completion)
        let (event_tx, event_rx) = mpsc::unbounded_channel();

        // Send CreateJob message with scheduler strategy
        let create_msg = DispatcherMessage::CreateJob {
            program_id,
            scheduler,
        };
        {
            let mut sender = ws_sender.lock().await;
            let text = serde_json::to_string(&create_msg)?;
            sender.send(Message::Text(text)).await?;
        }

        // Wait for JobCreated response
        let job_id = loop {
            match ws_receiver.next().await {
                Some(Ok(Message::Text(text))) => {
                    match serde_json::from_str::<DispatcherResponse>(&text) {
                        Ok(DispatcherResponse::JobCreated { job_id }) => {
                            break job_id;
                        }
                        Ok(DispatcherResponse::Error { message }) => {
                            anyhow::bail!("Failed to create job: {}", message);
                        }
                        Ok(_) => continue,
                        Err(e) => {
                            anyhow::bail!("Failed to parse response: {}", e);
                        }
                    }
                }
                Some(Ok(Message::Close(_))) => {
                    anyhow::bail!("Connection closed before job creation");
                }
                Some(Err(e)) => {
                    anyhow::bail!("WebSocket error: {}", e);
                }
                None => {
                    anyhow::bail!("Connection closed unexpectedly");
                }
                _ => continue,
            }
        };

        // Spawn task to receive events
        let ws_sender_clone = ws_sender.clone();
        tokio::spawn(async move {
            while let Some(msg) = ws_receiver.next().await {
                match msg {
                    Ok(Message::Text(text)) => {
                        match serde_json::from_str::<DispatcherResponse>(&text) {
                            Ok(DispatcherResponse::Result { work_unit_id, executor_id, data }) => {
                                let result = WorkUnitResult {
                                    work_unit_id: eridani_hpc_common::WorkUnitId(work_unit_id),
                                    executor_id,
                                    data,
                                };
                                if event_tx.send(JobEvent::Result(result)).is_err() {
                                    break;
                                }
                            }
                            Ok(DispatcherResponse::JobCompleted { .. }) => {
                                let _ = event_tx.send(JobEvent::Completed);
                                break;
                            }
                            Ok(DispatcherResponse::JobFailed {
                                job_id,
                                reason,
                                work_units_completed,
                                work_units_total,
                            }) => {
                                let _ = event_tx.send(JobEvent::Failed(JobError {
                                    job_id,
                                    reason,
                                    work_units_completed,
                                    work_units_total,
                                }));
                                break;
                            }
                            Ok(DispatcherResponse::JobCancelled { job_id }) => {
                                let _ = event_tx.send(JobEvent::Cancelled(job_id));
                                break;
                            }
                            Ok(DispatcherResponse::Error { message }) => {
                                let _ = event_tx.send(JobEvent::Failed(JobError {
                                    job_id: String::new(),
                                    reason: message,
                                    work_units_completed: 0,
                                    work_units_total: 0,
                                }));
                                break;
                            }
                            _ => {}
                        }
                    }
                    Ok(Message::Ping(data)) => {
                        let mut sender = ws_sender_clone.lock().await;
                        let _ = sender.send(Message::Pong(data)).await;
                    }
                    Ok(Message::Close(_)) => break,
                    Err(_) => break,
                    _ => {}
                }
            }
        });

        Ok(Self {
            job_id: Some(job_id),
            ws_sender,
            event_rx,
            work_units_submitted: 0,
            results_received: 0,
            failure: None,
        })
    }

    /// Send a batch of work units
    pub async fn send_work_units(&mut self, units: Vec<Vec<u8>>) -> Result<()> {
        if units.is_empty() {
            return Ok(());
        }

        self.work_units_submitted += units.len();

        let msg = DispatcherMessage::WorkUnits { units };
        let text = serde_json::to_string(&msg)?;

        let mut sender = self.ws_sender.lock().await;
        sender.send(Message::Text(text)).await?;

        Ok(())
    }

    /// Signal that work unit submission is complete
    pub async fn end_submission(&mut self) -> Result<()> {
        let msg = DispatcherMessage::EndJob;
        let text = serde_json::to_string(&msg)?;

        let mut sender = self.ws_sender.lock().await;
        sender.send(Message::Text(text)).await?;

        Ok(())
    }

    /// Get the job ID
    pub fn id(&self) -> Option<&str> {
        self.job_id.as_deref()
    }

    /// Get the next result (blocks until available or job completes/fails)
    /// Returns Ok(Some(result)) for each result, Ok(None) on completion, Err on failure
    pub async fn next_result(&mut self) -> Result<Option<WorkUnitResult>> {
        // If we already know the job failed, return the error
        if let Some(ref err) = self.failure {
            anyhow::bail!(err.clone());
        }

        match self.event_rx.recv().await {
            Some(JobEvent::Result(result)) => {
                self.results_received += 1;
                Ok(Some(result))
            }
            Some(JobEvent::Completed) => Ok(None),
            Some(JobEvent::Failed(err)) => {
                self.failure = Some(err.clone());
                anyhow::bail!(err)
            }
            Some(JobEvent::Cancelled(job_id)) => {
                anyhow::bail!("Job {} was cancelled", job_id)
            }
            None => Ok(None),
        }
    }

    /// Get the next result with deserialized data (blocks until available or job completes/fails)
    /// Returns Ok(Some(result)) for each result, Ok(None) on completion, Err on failure
    pub async fn next<T: DeserializeOwned>(&mut self) -> Result<Option<TypedResult<T>>> {
        // If we already know the job failed, return the error
        if let Some(ref err) = self.failure {
            anyhow::bail!(err.clone());
        }

        match self.event_rx.recv().await {
            Some(JobEvent::Result(result)) => {
                self.results_received += 1;
                let value: T = bincode::deserialize(&result.data)?;
                Ok(Some(TypedResult {
                    work_unit_id: result.work_unit_id,
                    executor_id: result.executor_id,
                    data: value,
                }))
            }
            Some(JobEvent::Completed) => Ok(None),
            Some(JobEvent::Failed(err)) => {
                self.failure = Some(err.clone());
                anyhow::bail!(err)
            }
            Some(JobEvent::Cancelled(job_id)) => {
                anyhow::bail!("Job {} was cancelled", job_id)
            }
            None => Ok(None),
        }
    }

    /// Collect all results into a Vec, deserializing each
    /// Returns Err if the job fails or is cancelled
    pub async fn collect<T: DeserializeOwned>(&mut self) -> Result<Vec<TypedResult<T>>> {
        let mut results = Vec::new();
        while let Some(result) = self.next::<T>().await? {
            results.push(result);
        }
        Ok(results)
    }

    /// Check if the job has failed
    pub fn has_failed(&self) -> bool {
        self.failure.is_some()
    }

    /// Get the failure error if the job has failed
    pub fn failure(&self) -> Option<&JobError> {
        self.failure.as_ref()
    }

    /// Get the number of work units submitted
    pub fn work_units_submitted(&self) -> usize {
        self.work_units_submitted
    }

    /// Get the number of results received so far
    pub fn results_received(&self) -> usize {
        self.results_received
    }

    /// Check if all results have been received
    pub fn is_complete(&self) -> bool {
        self.results_received >= self.work_units_submitted && self.work_units_submitted > 0
    }

    /// Cancel the job
    pub async fn cancel(&mut self) -> Result<()> {
        if let Some(job_id) = &self.job_id {
            let msg = DispatcherMessage::CancelJob {
                job_id: job_id.clone(),
            };
            let text = serde_json::to_string(&msg)?;

            let mut sender = self.ws_sender.lock().await;
            sender.send(Message::Text(text)).await?;
        }
        Ok(())
    }
}

#[derive(Debug, serde::Deserialize)]
pub struct ClusterStatus {
    pub agents_online: usize,
    pub total_cores: u32,
    pub total_memory_bytes: u64,
}

