use anyhow::Result;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use enact_core::kernel::NodeId;
use redis::AsyncCommands;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeInfo {
    pub node_id: NodeId,
    pub endpoint: String,
    pub capabilities: Vec<String>,
    pub load: f32,
    pub last_heartbeat: DateTime<Utc>,
}

#[async_trait]
pub trait NodeRegistry: Send + Sync {
    async fn register(&self, info: NodeInfo) -> Result<()>;
    async fn heartbeat(&self, node_id: &NodeId) -> Result<()>;
    async fn list_healthy(&self) -> Result<Vec<NodeInfo>>;
    async fn deregister(&self, node_id: &NodeId) -> Result<()>;
}

pub struct RedisNodeRegistry {
    client: redis::Client,
    key_prefix: String,
    ttl_secs: u64,
}

impl RedisNodeRegistry {
    pub fn new(redis_url: &str, ttl_secs: u64) -> Result<Self> {
        Ok(Self {
            client: redis::Client::open(redis_url)?,
            key_prefix: "enact:nodes".to_string(),
            ttl_secs,
        })
    }

    fn node_key(&self, node_id: &NodeId) -> String {
        format!("{}:{}", self.key_prefix, node_id)
    }

    fn index_key(&self) -> String {
        format!("{}:index", self.key_prefix)
    }
}

#[async_trait]
impl NodeRegistry for RedisNodeRegistry {
    async fn register(&self, info: NodeInfo) -> Result<()> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let key = self.node_key(&info.node_id);
        let value = serde_json::to_string(&info)?;

        let _: () = conn.set_ex(&key, value, self.ttl_secs).await?;
        let _: usize = conn
            .sadd(self.index_key(), info.node_id.to_string())
            .await?;
        Ok(())
    }

    async fn heartbeat(&self, node_id: &NodeId) -> Result<()> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let key = self.node_key(node_id);
        let value: Option<String> = conn.get(&key).await?;

        if let Some(raw) = value {
            let mut info: NodeInfo = serde_json::from_str(&raw)?;
            info.last_heartbeat = Utc::now();
            let updated = serde_json::to_string(&info)?;
            let _: () = conn.set_ex(&key, updated, self.ttl_secs).await?;
        }

        Ok(())
    }

    async fn list_healthy(&self) -> Result<Vec<NodeInfo>> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let node_ids: Vec<String> = conn.smembers(self.index_key()).await.unwrap_or_default();

        let mut out = Vec::new();
        for node_id in node_ids {
            let key = format!("{}:{}", self.key_prefix, node_id);
            let raw: Option<String> = conn.get(&key).await?;
            if let Some(raw) = raw {
                if let Ok(info) = serde_json::from_str::<NodeInfo>(&raw) {
                    out.push(info);
                }
            } else {
                let _: usize = conn.srem(self.index_key(), node_id).await.unwrap_or(0);
            }
        }

        Ok(out)
    }

    async fn deregister(&self, node_id: &NodeId) -> Result<()> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let _: usize = conn.del(self.node_key(node_id)).await?;
        let _: usize = conn.srem(self.index_key(), node_id.to_string()).await?;
        Ok(())
    }
}
