use std::{
    collections::{BTreeMap, HashMap},
    net::SocketAddr,
    str::FromStr,
    sync::Arc,
    time::Duration,
};

use anyhow::Result;
use parking_lot::RwLock;
use tonic::{
    transport::{ClientTlsConfig, Endpoint},
    Request,
};
use tracing as log;

pub mod gossip {
    #![allow(unused_qualifications)]
    tonic::include_proto!("mesh.gossip");
}
use gossip::{
    gossip_client, gossip_message, GossipMessage, NodeState, NodeStatus, NodeUpdate, Ping,
    StateSync,
};

use crate::{
    controller::MeshController,
    crdt::SKey,
    mtls::{MTLSConfig, MTLSManager},
    node_state_machine::{ConvergenceConfig, NodeStateMachine},
    partition::PartitionDetector,
    ping_server::GossipService,
    stores::{AppState, StateStores},
    sync::MeshSyncManager,
};

pub type ClusterState = Arc<RwLock<BTreeMap<String, NodeState>>>;

pub struct MeshServerConfig {
    pub self_name: String,
    pub self_addr: SocketAddr,
    pub init_peer: Option<SocketAddr>,
    pub mtls_config: Option<MTLSConfig>,
}

/// MeshServerHandler
/// It is the handler for the mesh server, which is responsible for the node management.
/// Includes some basic node management logic, like shutdown,
/// node discovery(TODO), node status update(TODO), etc.
pub struct MeshServerHandler {
    pub state: ClusterState,
    pub stores: Arc<StateStores>,
    pub sync_manager: Arc<MeshSyncManager>,
    pub self_name: String,
    _self_addr: SocketAddr,
    signal_tx: tokio::sync::watch::Sender<bool>,
    partition_detector: Option<Arc<PartitionDetector>>,
    state_machine: Option<Arc<NodeStateMachine>>,
    rate_limit_task_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
}

impl MeshServerHandler {
    /// Get partition detector
    pub fn partition_detector(&self) -> Option<&Arc<PartitionDetector>> {
        self.partition_detector.as_ref()
    }

    /// Get state machine
    pub fn state_machine(&self) -> Option<&Arc<NodeStateMachine>> {
        self.state_machine.as_ref()
    }

    /// Check if node is ready
    pub fn is_ready(&self) -> bool {
        self.state_machine
            .as_ref()
            .map(|sm| sm.is_ready())
            .unwrap_or(true) // If no state machine, consider ready
    }

    /// Check if we should serve (have quorum)
    pub fn should_serve(&self) -> bool {
        self.partition_detector
            .as_ref()
            .map(|pd| pd.should_serve())
            .unwrap_or(true) // If no partition detector, consider should serve
    }

    /// Start rate limit window reset task
    /// This task will periodically reset the global rate limit counter
    pub fn start_rate_limit_task(&self, window_seconds: u64) {
        use crate::rate_limit_window::RateLimitWindow;

        let window_manager = RateLimitWindow::new(self.sync_manager.clone(), window_seconds);
        let shutdown_rx = self.signal_tx.subscribe();

        let handle = tokio::spawn(async move {
            window_manager.start_reset_task(shutdown_rx).await;
        });

        if let Ok(mut task_handle) = self.rate_limit_task_handle.lock() {
            *task_handle = Some(handle);
        }
    }

    /// Stop rate limit window reset task
    pub fn stop_rate_limit_task(&self) {
        self.signal_tx.send(true).ok();
        if let Ok(mut task_handle) = self.rate_limit_task_handle.lock() {
            if let Some(handle) = task_handle.take() {
                tokio::spawn(async move {
                    if let Err(err) = handle.await {
                        log::warn!("Rate limit task shutdown failed: {}", err);
                    }
                });
            }
        }
    }

    /// Shutdown immediately without graceful shutdown
    pub fn shutdown(&self) {
        self.stop_rate_limit_task();
    }

    /// Graceful shutdown: broadcast LEAVING status to all alive nodes,
    /// wait for propagation, then shutdown
    pub async fn graceful_shutdown(&self) -> Result<()> {
        log::info!("Graceful shutdown for node {}", self.self_name);

        let maybe_leaving = {
            let state = self.state.read();

            if let Some(self_node) = state.get(&self.self_name) {
                let mut self_node = self_node.clone();
                if self_node.status == NodeStatus::Leaving as i32 {
                    None
                } else {
                    self_node.status = NodeStatus::Leaving as i32;
                    self_node.version += 1;

                    let alive_nodes = state
                        .values()
                        .filter(|node| {
                            node.status == NodeStatus::Alive as i32 && node.name != self.self_name
                            // exclude self from broadcast targets
                        })
                        .cloned()
                        .collect::<Vec<NodeState>>();

                    Some((self_node, alive_nodes))
                }
            } else {
                None
            }
        };
        let (leaving_node, alive_nodes) = match maybe_leaving {
            Some(values) => values,
            None => {
                self.stop_rate_limit_task();
                return Ok(());
            }
        };

        log::info!(
            "Broadcasting LEAVING status to {} alive nodes",
            alive_nodes.len()
        );

        // Broadcast LEAVING status to all alive nodes
        let (success_count, total_count) = broadcast_node_states(
            vec![leaving_node],
            alive_nodes,
            Some(Duration::from_secs(3)),
        )
        .await;

        log::info!(
            "Broadcast LEAVING status: {}/{} successful",
            success_count,
            total_count
        );

        // Wait a bit more for state propagation
        let propagation_delay = Duration::from_secs(1);
        log::info!(
            "Waiting {} seconds for LEAVING status propagation",
            propagation_delay.as_secs()
        );
        tokio::time::sleep(propagation_delay).await;

        log::info!("Stopping rate limit task and signaling shutdown");
        self.stop_rate_limit_task();
        Ok(())
    }

    /// Calculate the next version for a key
    /// If the key exists, increment its version by 1
    /// If the key doesn't exist, start with version 1
    fn next_version(&self, key: &str) -> u64 {
        self.stores
            .app
            .get(&SKey(key.to_string()))
            .map(|app_state| app_state.version + 1)
            .unwrap_or(1)
    }

    pub fn write_data(&self, key: String, value: Vec<u8>) -> Result<()> {
        // Keep app store write and metadata/version update in one lock scope.
        let mut state = self.state.write();
        let node = state.get_mut(&self.self_name).ok_or_else(|| {
            anyhow::anyhow!(
                "Node {} not found in cluster state during write_data",
                self.self_name
            )
        })?;

        let version = self.next_version(&key);
        let app_state = AppState {
            key: key.clone(),
            value: value.clone(),
            version,
        };
        self.stores
            .app
            .insert(SKey(key.clone()), app_state, self.self_name.clone());

        node.metadata.insert(key, value);
        node.version += 1;
        Ok(())
    }

    pub fn read_data(&self, key: String) -> Option<Vec<u8>> {
        // Read from the app store
        self.stores
            .app
            .get(&SKey(key))
            .map(|app_state| app_state.value.clone())
    }

    /// Get a snapshot of the app store for synchronization
    /// Returns a CRDT snapshot that can be merged into other nodes
    pub fn snapshot(&self) -> crate::crdt::CRDTMap<AppState> {
        self.stores.app.snapshot()
    }

    /// Sync app store data from a snapshot (for testing and manual sync)
    /// This will be replaced by automatic sync stream in the future
    pub fn sync_app_from_snapshot(&self, snapshot: &crate::crdt::CRDTMap<AppState>) {
        // Merge snapshot into our app store using CRDT merge
        self.stores.app.merge(snapshot);
    }
}

pub struct MeshServerBuilder {
    state: ClusterState,
    stores: Arc<StateStores>,
    self_name: String,
    self_addr: SocketAddr,
    init_peer: Option<SocketAddr>,
    mtls_manager: Option<Arc<MTLSManager>>,
}

impl MeshServerBuilder {
    pub fn new(self_name: String, self_addr: SocketAddr, init_peer: Option<SocketAddr>) -> Self {
        let state = Arc::new(RwLock::new(BTreeMap::from([(
            self_name.clone(),
            NodeState {
                name: self_name.clone(),
                address: self_addr.to_string(),
                status: NodeStatus::Alive as i32,
                version: 1,
                metadata: HashMap::new(),
            },
        )])));
        let stores = Arc::new(StateStores::with_self_name(self_name.clone()));
        Self {
            state,
            stores,
            self_name,
            self_addr,
            init_peer,
            mtls_manager: None,
        }
    }

    pub fn with_mtls(mut self, mtls_config: MTLSConfig) -> Self {
        self.mtls_manager = Some(Arc::new(MTLSManager::new(mtls_config)));
        self
    }

    pub fn build(&self) -> (MeshServer, MeshServerHandler) {
        let (signal_tx, signal_rx) = tokio::sync::watch::channel(false);
        let partition_detector = Arc::new(PartitionDetector::default());
        let sync_manager = Arc::new(MeshSyncManager::new(
            self.stores.clone(),
            self.self_name.clone(),
        ));
        let state_machine = Arc::new(NodeStateMachine::new(
            self.stores.clone(),
            ConvergenceConfig::default(),
        ));
        // Initialize rate-limit hash ring with current membership
        sync_manager.update_rate_limit_membership();
        (
            MeshServer {
                state: self.state.clone(),
                stores: self.stores.clone(),
                sync_manager: sync_manager.clone(),
                self_name: self.self_name.clone(),
                self_addr: self.self_addr,
                init_peer: self.init_peer,
                signal_rx,
                partition_detector: Some(partition_detector.clone()),
                mtls_manager: self.mtls_manager.clone(),
            },
            MeshServerHandler {
                state: self.state.clone(),
                stores: self.stores.clone(),
                sync_manager,
                self_name: self.self_name.clone(),
                _self_addr: self.self_addr,
                signal_tx,
                partition_detector: Some(partition_detector),
                state_machine: Some(state_machine),
                rate_limit_task_handle: std::sync::Mutex::new(None),
            },
        )
    }
}

impl From<&MeshServerConfig> for MeshServerBuilder {
    fn from(value: &MeshServerConfig) -> Self {
        let mut builder =
            MeshServerBuilder::new(value.self_name.clone(), value.self_addr, value.init_peer);
        if let Some(mtls_config) = &value.mtls_config {
            builder = builder.with_mtls(mtls_config.clone());
        }
        builder
    }
}

pub struct MeshServer {
    state: ClusterState,
    stores: Arc<StateStores>,
    sync_manager: Arc<MeshSyncManager>,
    self_name: String,
    self_addr: SocketAddr,
    init_peer: Option<SocketAddr>,
    signal_rx: tokio::sync::watch::Receiver<bool>,
    partition_detector: Option<Arc<PartitionDetector>>,
    mtls_manager: Option<Arc<MTLSManager>>,
}

impl MeshServer {
    fn build_ping_server(&self) -> GossipService {
        GossipService::new(self.state.clone(), self.self_addr, &self.self_name)
    }

    fn build_controller(&self) -> MeshController {
        MeshController::new(
            self.state.clone(),
            self.self_addr,
            &self.self_name,
            self.init_peer,
            self.stores.clone(),
            self.sync_manager.clone(),
            self.mtls_manager.clone(),
        )
    }

    pub async fn start(self) -> Result<()> {
        log::info!("Mesh server listening on {}", self.self_addr);
        let self_name = self.self_name.clone();
        let self_address = self.self_addr;

        let partition_detector = self
            .partition_detector
            .clone()
            .expect("partition detector missing");

        let mut service = self.build_ping_server();
        service = service.with_stores(self.stores.clone());

        service = service.with_sync_manager(self.sync_manager.clone());

        service = service.with_partition_detector(partition_detector);

        // Add mTLS support if configured
        if let Some(mtls_manager) = self.mtls_manager.clone() {
            service = service.with_mtls_manager(mtls_manager);
        }

        let controller = self.build_controller();

        let mut service_shutdown = self.signal_rx.clone();

        let listener = tokio::spawn(service.serve_ping_with_shutdown(async move {
            _ = service_shutdown.changed().await;
        }));
        tokio::time::sleep(Duration::from_secs(1)).await;
        let app_handle = tokio::spawn(controller.event_loop(self.signal_rx.clone()));

        tokio::select! {
            res = listener => res??,
            res = app_handle => res??,
        }

        log::info!(
            "Mesh server {} at {} is shutting down",
            self_name,
            self_address
        );
        Ok(())
    }
}

/// Broadcast node state updates to target nodes
/// Returns (success_count, total_count)
pub async fn broadcast_node_states(
    nodes_to_broadcast: Vec<NodeState>,
    target_nodes: Vec<NodeState>,
    timeout: Option<Duration>,
) -> (usize, usize) {
    if nodes_to_broadcast.is_empty() || target_nodes.is_empty() {
        log::debug!(
            "Nothing to broadcast: nodes_to_broadcast={}, target_nodes={}",
            nodes_to_broadcast.len(),
            target_nodes.len()
        );
        return (0, target_nodes.len());
    }

    let mut broadcast_tasks = Vec::new();
    for target_node in &target_nodes {
        let target_node_clone = target_node.clone();
        let nodes_for_task = nodes_to_broadcast.clone();
        let task = tokio::spawn(async move {
            let state_sync = StateSync {
                nodes: nodes_for_task,
            };
            let ping_payload = gossip_message::Payload::Ping(Ping {
                state_sync: Some(state_sync),
            });
            match try_ping(&target_node_clone, Some(ping_payload), None).await {
                Ok(_) => {
                    log::debug!("Successfully broadcasted to {}", target_node_clone.name);
                    Ok(())
                }
                Err(e) => {
                    log::warn!("Failed to broadcast to {}: {}", target_node_clone.name, e);
                    Err(e)
                }
            }
        });
        broadcast_tasks.push(task);
    }

    let timeout_duration = timeout.unwrap_or(Duration::from_secs(3));
    let broadcast_result = tokio::time::timeout(timeout_duration, async {
        futures::future::join_all(broadcast_tasks).await
    })
    .await;

    match broadcast_result {
        Ok(results) => {
            let success_count = results.iter().filter(|r| matches!(r, Ok(Ok(_)))).count();
            let total_count = target_nodes.len();
            log::info!(
                "Broadcast completed: {}/{} successful",
                success_count,
                total_count
            );
            (success_count, total_count)
        }
        Err(_) => {
            log::warn!(
                "Broadcast timeout after {} seconds",
                timeout_duration.as_secs()
            );
            (0, target_nodes.len())
        }
    }
}

pub async fn try_ping(
    peer_node: &NodeState,
    payload: Option<gossip_message::Payload>,
    mtls_manager: Option<Arc<MTLSManager>>,
) -> Result<NodeUpdate, tonic::Status> {
    let peer_name = peer_node.name.clone();

    let peer_addr = SocketAddr::from_str(&peer_node.address).map_err(|e| {
        tonic::Status::invalid_argument(format!(
            "Invalid address for node {}: {}, {}",
            peer_name, peer_node.address, e
        ))
    })?;

    let connect_url = if mtls_manager.is_some() {
        format!("https://{}", peer_addr)
    } else {
        format!("http://{}", peer_addr)
    };

    let mut endpoint = Endpoint::from_shared(connect_url.clone()).map_err(|e| {
        tonic::Status::invalid_argument(format!(
            "Invalid endpoint for node {}: {}, {}",
            peer_name, connect_url, e
        ))
    })?;

    if let Some(mtls_manager) = mtls_manager {
        mtls_manager.load_client_config().await.map_err(|e| {
            tonic::Status::unavailable(format!(
                "Failed to load mTLS client config for {}: {}",
                peer_name, e
            ))
        })?;

        let tls_domain = endpoint
            .uri()
            .host()
            .map(str::to_owned)
            .unwrap_or_else(|| peer_name.clone());
        let ca_certificate = mtls_manager.load_ca_certificate().await.map_err(|e| {
            tonic::Status::unavailable(format!(
                "Failed to load mTLS CA certificate for {}: {}",
                peer_name, e
            ))
        })?;

        endpoint = endpoint
            .tls_config(
                ClientTlsConfig::new()
                    .domain_name(tls_domain)
                    .ca_certificate(ca_certificate),
            )
            .map_err(|e| {
                tonic::Status::unavailable(format!(
                    "Failed to configure TLS endpoint for {}: {}",
                    peer_name, e
                ))
            })?;
    }

    let channel = endpoint.connect().await.map_err(|e| {
        log::warn!(
            "Failed to connect to peer {} {}: {}.",
            peer_name,
            peer_addr,
            e
        );
        tonic::Status::unavailable("Failed to connect to peer")
    })?;
    let mut client = gossip_client::GossipClient::new(channel);

    let ping_message = GossipMessage { payload };
    let response = client.ping_server(Request::new(ping_message)).await?;

    Ok(response.into_inner())
}

#[macro_export]
macro_rules! mesh_run {
    ($addr:expr, $init_peer:expr) => {{
        mesh_run!($addr.to_string(), $addr, $init_peer)
    }};

    ($name:expr, $addr:expr, $init_peer:expr) => {{
        tracing::info!("Starting mesh server : {}", $addr);
        use $crate::MeshServerBuilder;
        let (server, handler) =
            MeshServerBuilder::new($name.to_string(), $addr, $init_peer).build();
        tokio::spawn(async move {
            if let Err(e) = server.start().await {
                tracing::error!("Mesh server failed: {}", e);
            }
        });
        handler
    }};
}

#[cfg(test)]
mod tests {
    use std::sync::Once;

    use tokio::net::TcpListener;
    use tracing as log;
    use tracing_subscriber::{
        filter::LevelFilter, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter,
    };

    use super::*;
    static INIT: Once = Once::new();
    fn init() {
        INIT.call_once(|| {
            let _ = tracing_subscriber::registry()
                .with(tracing_subscriber::fmt::layer())
                .with(
                    EnvFilter::builder()
                        .with_default_directive(LevelFilter::INFO.into())
                        .from_env_lossy(),
                )
                .try_init();
        });
    }
    async fn find_free_port() -> (TcpListener, u16) {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        log::info!("Found free port: {}", port);
        (listener, port)
    }

    async fn get_node() -> SocketAddr {
        let (_listener, port) = find_free_port().await;
        format!("127.0.0.1:{}", port).parse().unwrap()
    }

    fn print_state(handler: &MeshServerHandler) -> String {
        let state = handler.state.read();
        let mut res = vec![];
        for (k, v) in state.iter() {
            res.push(format!(
                "{}: {:?} - {:?}",
                k,
                NodeStatus::try_from(v.status).unwrap(),
                v.metadata
            ));
        }
        res.join(", ")
    }

    #[tokio::test]
    #[ignore = "reason: only run manually due to flakiness"]
    async fn test_state_synchronization() {
        init();
        log::info!("Starting test_state_synchronization");

        // 1. setup node A and B for initial cluster
        let addr_a = get_node().await;
        let handler_a = mesh_run!("A", addr_a, None);
        let addr_b = get_node().await;
        let handler_b = mesh_run!("B", addr_b, Some(addr_a));

        // 2. wait for node A and B to sync and write some data
        tokio::time::sleep(Duration::from_secs(2)).await;
        handler_a
            .write_data("hello".into(), "world".into())
            .unwrap();
        log::info!("================================================");

        // 3. add node C and D and wait for them to sync
        let addr_c = get_node().await;
        let handler_c = mesh_run!("C", addr_c, Some(addr_a));
        let addr_d = get_node().await;
        let handler_d = mesh_run!("D", addr_d, Some(addr_c));
        tokio::time::sleep(Duration::from_secs(2)).await;
        log::info!("================================================");

        // 4. add node E and wait for it to sync and kill it
        {
            let addr_e = get_node().await;
            let handler_e = mesh_run!("E", addr_e, Some(addr_d));
            tokio::time::sleep(Duration::from_secs(3)).await;
            log::info!("State E: {:?}", print_state(&handler_e));
            // killing_button.send(()).unwrap();
            handler_e.shutdown();
        }

        handler_d.graceful_shutdown().await.unwrap();
        tokio::time::sleep(Duration::from_secs(2)).await;
        log::info!("================================================");

        // 5. Poll until expected state is reached (with timeout)
        // Note: D stays in Leaving status because gracefully shutdown nodes are not pinged anymore (by design)
        //       E transitions Alive → Suspected → Down because it was abruptly shutdown
        let final_state = String::from("A: Alive - {\"hello\": [119, 111, 114, 108, 100]}, B: Alive - {}, C: Alive - {}, D: Leaving - {}, E: Down - {}");

        let max_wait = Duration::from_secs(30);
        let poll_interval = Duration::from_millis(500);
        let start = std::time::Instant::now();

        loop {
            let state_a = print_state(&handler_a);
            let state_b = print_state(&handler_b);
            let state_c = print_state(&handler_c);

            if state_a == final_state && state_b == final_state && state_c == final_state {
                log::info!("All nodes converged to expected state");
                break;
            }

            if start.elapsed() > max_wait {
                log::info!("================================================");
                panic!(
                    "Timeout waiting for state convergence.\nExpected: {:?}\nState A: {:?}\nState B: {:?}\nState C: {:?}",
                    final_state, state_a, state_b, state_c
                );
            }

            tokio::time::sleep(poll_interval).await;
        }
    }
}
