use anyhow::Result;
use async_trait::async_trait;
use enact_core::kernel::SpawnMode;
use redis::streams::{StreamReadOptions, StreamReadReply};
use redis::{AsyncCommands, FromRedisValue};
use serde::{Deserialize, Serialize};

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

impl TaskId {
    pub fn new() -> Self {
        Self(format!("task_{}", uuid::Uuid::new_v4()))
    }
}

impl Default for TaskId {
    fn default() -> Self {
        Self::new()
    }
}

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DispatchTask {
    pub task_id: TaskId,
    pub agent_name: String,
    pub input: serde_json::Value,
    pub spawn_mode: SpawnMode,
    pub priority: u8,
}

#[async_trait]
pub trait TaskQueue: Send + Sync {
    async fn enqueue(&self, task: DispatchTask) -> Result<TaskId>;
    async fn claim_next(
        &self,
        worker_id: &enact_core::kernel::NodeId,
    ) -> Result<Option<DispatchTask>>;
    async fn ack(&self, task_id: &TaskId) -> Result<()>;
    async fn fail(&self, task_id: &TaskId, reason: &str) -> Result<()>;
}

pub struct RedisTaskQueue {
    client: redis::Client,
    stream: String,
    group: String,
}

impl RedisTaskQueue {
    pub fn new(redis_url: &str, stream: impl Into<String>) -> Result<Self> {
        Ok(Self {
            client: redis::Client::open(redis_url)?,
            stream: stream.into(),
            group: "enact-workers".to_string(),
        })
    }

    fn id_map_key(&self) -> String {
        format!("{}:task_id_map", self.stream)
    }

    fn fail_key(&self, task_id: &TaskId) -> String {
        format!("{}:failed:{}", self.stream, task_id)
    }

    async fn ensure_group(&self) -> Result<()> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let res: redis::RedisResult<()> = redis::cmd("XGROUP")
            .arg("CREATE")
            .arg(&self.stream)
            .arg(&self.group)
            .arg("0")
            .arg("MKSTREAM")
            .query_async(&mut conn)
            .await;

        if let Err(err) = res {
            let msg = err.to_string();
            if !msg.contains("BUSYGROUP") {
                return Err(err.into());
            }
        }

        Ok(())
    }
}

#[async_trait]
impl TaskQueue for RedisTaskQueue {
    async fn enqueue(&self, mut task: DispatchTask) -> Result<TaskId> {
        self.ensure_group().await?;

        if task.task_id.0.is_empty() {
            task.task_id = TaskId::new();
        }

        let task_id = task.task_id.clone();
        let payload = serde_json::to_string(&task)?;

        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let entry_id: String = redis::cmd("XADD")
            .arg(&self.stream)
            .arg("*")
            .arg("task_id")
            .arg(task_id.to_string())
            .arg("payload")
            .arg(payload)
            .query_async(&mut conn)
            .await?;

        let _: usize = conn
            .hset(self.id_map_key(), task_id.to_string(), entry_id)
            .await?;

        Ok(task_id)
    }

    async fn claim_next(
        &self,
        worker_id: &enact_core::kernel::NodeId,
    ) -> Result<Option<DispatchTask>> {
        self.ensure_group().await?;

        let consumer = worker_id.to_string();
        let opts = StreamReadOptions::default()
            .group(&self.group, &consumer)
            .count(1)
            .block(1000);

        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let reply: StreamReadReply = conn
            .xread_options(&[&self.stream], &[">"], &opts)
            .await
            .unwrap_or(StreamReadReply { keys: Vec::new() });

        let Some(stream_key) = reply.keys.into_iter().next() else {
            return Ok(None);
        };
        let Some(stream_id) = stream_key.ids.into_iter().next() else {
            return Ok(None);
        };

        let payload_val = stream_id.map.get("payload");
        let Some(payload_val) = payload_val else {
            return Ok(None);
        };

        let payload: String = FromRedisValue::from_redis_value(payload_val).unwrap_or_default();
        if payload.is_empty() {
            return Ok(None);
        }

        let task: DispatchTask = serde_json::from_str(&payload)?;

        let _: usize = conn
            .hset(self.id_map_key(), task.task_id.to_string(), stream_id.id)
            .await?;

        Ok(Some(task))
    }

    async fn ack(&self, task_id: &TaskId) -> Result<()> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let entry_id: Option<String> = conn.hget(self.id_map_key(), task_id.to_string()).await?;

        if let Some(entry_id) = entry_id {
            let _: usize = redis::cmd("XACK")
                .arg(&self.stream)
                .arg(&self.group)
                .arg(entry_id)
                .query_async(&mut conn)
                .await?;
            let _: usize = conn.hdel(self.id_map_key(), task_id.to_string()).await?;
        }

        Ok(())
    }

    async fn fail(&self, task_id: &TaskId, reason: &str) -> Result<()> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;

        let _: () = conn.set(self.fail_key(task_id), reason.to_string()).await?;

        self.ack(task_id).await
    }
}
