use std::collections::HashMap;
use std::sync::{
    atomic::{AtomicU64, Ordering},
    Arc,
};

use anyhow::{anyhow, Result};
use tokio::sync::{mpsc, oneshot, RwLock};
use tracing::{debug, warn};

use crate::server::ServerManager;

const DEFAULT_QUEUE_CAPACITY: usize = 256;

#[derive(Debug)]
struct CommandJob {
    command: String,
    response_tx: oneshot::Sender<Result<()>>,
}

#[derive(Clone)]
struct WorkerHandle {
    id: u64,
    sender: mpsc::Sender<CommandJob>,
}

pub struct CommandDispatcher {
    server_manager: Arc<ServerManager>,
    workers: Arc<RwLock<HashMap<String, WorkerHandle>>>,
    next_worker_id: Arc<AtomicU64>,
    queue_capacity: usize,
}

impl CommandDispatcher {
    pub fn new(server_manager: Arc<ServerManager>) -> Self {
        Self {
            server_manager,
            workers: Arc::new(RwLock::new(HashMap::new())),
            next_worker_id: Arc::new(AtomicU64::new(1)),
            queue_capacity: DEFAULT_QUEUE_CAPACITY,
        }
    }

    pub async fn send_command(&self, server_uuid: &str, command: &str) -> Result<()> {
        let cmd = command.trim();
        if cmd.is_empty() {
            return Err(anyhow!("Command cannot be empty"));
        }

        // Retry once if the current worker channel is unexpectedly closed.
        for _ in 0..2 {
            let worker = self.get_or_create_worker(server_uuid).await;
            let (response_tx, response_rx) = oneshot::channel();
            let job = CommandJob {
                command: cmd.to_string(),
                response_tx,
            };

            match worker.sender.send(job).await {
                Ok(_) => {
                    return response_rx
                        .await
                        .unwrap_or_else(|_| Err(anyhow!("Command worker dropped response channel")));
                }
                Err(send_err) => {
                    warn!(
                        server_uuid = %server_uuid,
                        worker_id = worker.id,
                        error = %send_err,
                        "Command worker queue send failed, rotating worker",
                    );
                    self.remove_worker_if_matches(server_uuid, worker.id).await;
                }
            }
        }

        Err(anyhow!("Failed to dispatch command after retry"))
    }

    async fn get_or_create_worker(&self, server_uuid: &str) -> WorkerHandle {
        if let Some(existing) = self.workers.read().await.get(server_uuid).cloned() {
            return existing;
        }

        let worker_id = self.next_worker_id.fetch_add(1, Ordering::Relaxed);
        let (sender, receiver) = mpsc::channel::<CommandJob>(self.queue_capacity);
        let handle = WorkerHandle {
            id: worker_id,
            sender: sender.clone(),
        };

        {
            let mut workers = self.workers.write().await;
            if let Some(existing) = workers.get(server_uuid).cloned() {
                return existing;
            }
            workers.insert(server_uuid.to_string(), handle.clone());
        }

        self.spawn_worker(server_uuid.to_string(), worker_id, receiver);
        handle
    }

    fn spawn_worker(
        &self,
        server_uuid: String,
        worker_id: u64,
        mut receiver: mpsc::Receiver<CommandJob>,
    ) {
        let server_manager = self.server_manager.clone();
        let workers = self.workers.clone();

        tokio::spawn(async move {
            debug!(server_uuid = %server_uuid, worker_id, "Command worker started");

            while let Some(job) = receiver.recv().await {
                let result = execute_command(server_manager.clone(), &server_uuid, &job.command).await;
                let _ = job.response_tx.send(result);
            }

            let mut map = workers.write().await;
            if let Some(current) = map.get(&server_uuid) {
                if current.id == worker_id {
                    map.remove(&server_uuid);
                }
            }
            debug!(server_uuid = %server_uuid, worker_id, "Command worker stopped");
        });
    }

    async fn remove_worker_if_matches(&self, server_uuid: &str, worker_id: u64) {
        let mut workers = self.workers.write().await;
        if let Some(current) = workers.get(server_uuid) {
            if current.id == worker_id {
                workers.remove(server_uuid);
            }
        }
    }
}

async fn execute_command(
    server_manager: Arc<ServerManager>,
    server_uuid: &str,
    command: &str,
) -> Result<()> {
    server_manager
        .docker()
        .send_command(server_uuid, command)
        .await?;

    Ok(())
}