//! FlareRaftNode - High-level wrapper for OpenRaft
//!
//! This module provides a convenient wrapper around OpenRaft's Raft<FlareTypeConfig>
//! that mirrors the interface of the old Peer struct, making migration easier.

use crate::network::FlareNetworkFactory;
use crate::persistent_storage::PersistentFlareStore;
use crate::storage::FlareStore;
use crate::types::{
    FlareNode, FlareNodeId, FlareRaft, FlareRequest, FlareResponse,
};
use flaredb_storage::rocks_engine::RocksEngine;
use openraft::error::{ClientWriteError, InitializeError, RaftError};
use openraft::{Config, Raft};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::sync::Arc;
use tracing::{info, warn};

/// Type alias for ClientWriteError with our types
pub type FlareClientWriteError = RaftError<FlareNodeId, ClientWriteError<FlareNodeId, FlareNode>>;

/// Storage backend for FlareRaftNode
pub enum FlareStorage {
    /// In-memory storage (for testing or non-persistent deployments)
    InMemory(Arc<FlareStore>),
    /// Persistent storage backed by RocksDB
    Persistent(Arc<PersistentFlareStore>),
}

impl FlareStorage {
    /// Read a value (eventual consistency)
    pub async fn read_kv(&self, namespace_id: u32, key: &[u8]) -> Option<(Vec<u8>, u64)> {
        match self {
            FlareStorage::InMemory(s) => s.read_kv(namespace_id, key).await,
            FlareStorage::Persistent(s) => s.read_kv(namespace_id, key).await,
        }
    }

    /// Read a CAS value (strong consistency)
    pub async fn read_cas(&self, namespace_id: u32, key: &[u8]) -> Option<(Vec<u8>, u64, u64)> {
        match self {
            FlareStorage::InMemory(s) => s.read_cas(namespace_id, key).await,
            FlareStorage::Persistent(s) => s.read_cas(namespace_id, key).await,
        }
    }

    /// Scan eventual-consistency KV data from the replicated state machine.
    pub async fn scan_kv(
        &self,
        namespace_id: u32,
        start_key: &[u8],
        end_key: &[u8],
        limit: usize,
    ) -> Vec<(Vec<u8>, Vec<u8>, u64)> {
        let sm = match self {
            FlareStorage::InMemory(s) => s.get_state_machine().await,
            FlareStorage::Persistent(s) => s.get_state_machine().await,
        };

        sm.kv_data
            .iter()
            .filter(|((ns_id, key), _)| {
                *ns_id == namespace_id
                    && key.as_slice() >= start_key
                    && (end_key.is_empty() || key.as_slice() < end_key)
            })
            .take(limit)
            .map(|((_ns_id, key), (value, ts))| (key.clone(), value.clone(), *ts))
            .collect()
    }

    /// Scan strong-consistency CAS data from the replicated state machine.
    pub async fn scan_cas(
        &self,
        namespace_id: u32,
        start_key: &[u8],
        end_key: &[u8],
        limit: usize,
    ) -> Vec<(Vec<u8>, Vec<u8>, u64, u64)> {
        let sm = match self {
            FlareStorage::InMemory(s) => s.get_state_machine().await,
            FlareStorage::Persistent(s) => s.get_state_machine().await,
        };

        sm.cas_data
            .iter()
            .filter(|((ns_id, key), _)| {
                *ns_id == namespace_id
                    && key.as_slice() >= start_key
                    && (end_key.is_empty() || key.as_slice() < end_key)
            })
            .take(limit)
            .map(|((_ns_id, key), (value, version, ts))| {
                (key.clone(), value.clone(), *version, *ts)
            })
            .collect()
    }
}

/// Wrapper around OpenRaft for easier migration from raft-rs
pub struct FlareRaftNode {
    /// The underlying OpenRaft instance
    pub raft: FlareRaft,

    /// Node ID (same as store_id)
    store_id: FlareNodeId,

    /// Region ID this node belongs to
    region_id: u64,

    /// Network factory for peer communication
    network_factory: FlareNetworkFactory,

    /// Storage instance for direct reads
    storage: FlareStorage,
}

impl FlareRaftNode {
    /// Create OpenRaft config with sensible defaults
    fn create_raft_config(region_id: u64) -> Result<Arc<Config>, Box<dyn std::error::Error + Send + Sync>> {
        let raft_config = Config {
            cluster_name: format!("flare-region-{}", region_id),
            // VM-backed cluster tests can stall for >1s while other services build images,
            // upload artifacts, or flush RocksDB state. Keep leader heartbeats conservative so
            // brief scheduler pauses do not trigger avoidable leader churn.
            heartbeat_interval: 2500,
            election_timeout_min: 10000,
            election_timeout_max: 20000,
            install_snapshot_timeout: 60000,
            replication_lag_threshold: 1000,
            snapshot_policy: openraft::SnapshotPolicy::LogsSinceLast(1000),
            max_in_snapshot_log_to_keep: 100,
            ..Default::default()
        };
        Ok(Arc::new(raft_config.validate()?))
    }

    /// Create a new Raft node with in-memory storage (for testing)
    pub async fn new(
        store_id: FlareNodeId,
        region_id: u64,
        network_factory: FlareNetworkFactory,
        voters: Vec<FlareNodeId>,
        peer_addrs: &HashMap<u64, String>,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let raft_config = Self::create_raft_config(region_id)?;

        // Create in-memory storage
        let storage = FlareStore::new_arc();

        // Create the Raft instance (storage-v2: storage implements both RaftLogStorage and RaftStateMachine)
        let raft = Raft::new(
            store_id,
            raft_config,
            network_factory.clone(),
            storage.clone(),
            storage.clone(),
        )
        .await?;

        let node = Self {
            raft,
            store_id,
            region_id,
            network_factory,
            storage: FlareStorage::InMemory(storage),
        };

        // Initialize cluster if voters are provided and cluster is not already initialized
        if !voters.is_empty() {
            node.try_initialize(voters, peer_addrs).await;
        }

        Ok(node)
    }

    /// Create a new Raft node with persistent RocksDB storage
    ///
    /// This constructor persists Raft logs, votes, and state to RocksDB,
    /// ensuring data survives node restarts.
    pub async fn new_persistent(
        store_id: FlareNodeId,
        region_id: u64,
        network_factory: FlareNetworkFactory,
        voters: Vec<FlareNodeId>,
        peer_addrs: &HashMap<u64, String>,
        engine: Arc<RocksEngine>,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let raft_config = Self::create_raft_config(region_id)?;

        // Create persistent storage backed by RocksDB
        let storage = PersistentFlareStore::new_arc(engine);

        // Create the Raft instance (storage-v2: storage implements both RaftLogStorage and RaftStateMachine)
        let raft = Raft::new(
            store_id,
            raft_config,
            network_factory.clone(),
            storage.clone(),
            storage.clone(),
        )
        .await?;

        let node = Self {
            raft,
            store_id,
            region_id,
            network_factory,
            storage: FlareStorage::Persistent(storage),
        };

        // Initialize cluster if voters are provided and cluster is not already initialized
        if !voters.is_empty() {
            node.try_initialize(voters, peer_addrs).await;
        }

        info!(
            store_id,
            region_id,
            "Created persistent Raft node with RocksDB storage"
        );

        Ok(node)
    }

    /// Try to initialize the cluster with the given voters
    async fn try_initialize(&self, voters: Vec<FlareNodeId>, peer_addrs: &HashMap<u64, String>) {
        // Build membership from voters with their addresses
        let membership: BTreeMap<FlareNodeId, FlareNode> = voters
            .into_iter()
            .map(|id| {
                let addr = peer_addrs.get(&id).cloned().unwrap_or_default();
                (
                    id,
                    FlareNode {
                        addr,
                        region_id: self.region_id,
                    },
                )
            })
            .collect();

        match self.raft.initialize(membership).await {
            Ok(_) => {
                info!(
                    store_id = self.store_id,
                    region_id = self.region_id,
                    "cluster initialized"
                );
            }
            Err(RaftError::APIError(InitializeError::NotAllowed(_))) => {
                // Already initialized, this is fine
            }
            Err(e) => {
                warn!(
                    error = %e,
                    store_id = self.store_id,
                    region_id = self.region_id,
                    "failed to initialize cluster"
                );
            }
        }
    }

    /// Get the store ID
    pub fn store_id(&self) -> FlareNodeId {
        self.store_id
    }

    /// Get the region ID
    pub fn region_id(&self) -> u64 {
        self.region_id
    }

    /// Check if this node is the current leader
    pub async fn is_leader(&self) -> bool {
        self.raft
            .current_leader()
            .await
            .map(|id| id == self.store_id)
            .unwrap_or(false)
    }

    /// Get the current leader ID
    pub async fn leader_id(&self) -> Option<FlareNodeId> {
        self.raft.current_leader().await
    }

    /// Get current membership
    pub async fn get_membership(&self) -> BTreeSet<FlareNodeId> {
        self.raft
            .metrics()
            .borrow()
            .membership_config
            .membership()
            .voter_ids()
            .collect()
    }

    /// Propose a KV write operation (replaces peer.propose_kv)
    /// Uses Last-Write-Wins for eventual consistency
    pub async fn write_kv(
        &self,
        namespace_id: u32,
        key: Vec<u8>,
        value: Vec<u8>,
        ts: u64,
    ) -> Result<FlareResponse, FlareClientWriteError> {
        let req = FlareRequest::KvWrite {
            namespace_id,
            key,
            value,
            ts,
        };

        let result = self.raft.client_write(req).await?;
        Ok(result.data)
    }

    /// Propose a CAS (Compare-And-Swap) write operation
    /// CAS is executed atomically within the Raft state machine,
    /// ensuring linearizable semantics across all replicas.
    pub async fn cas_write(
        &self,
        namespace_id: u32,
        key: Vec<u8>,
        value: Vec<u8>,
        expected_version: u64,
        ts: u64,
    ) -> Result<FlareResponse, FlareClientWriteError> {
        let req = FlareRequest::CasWrite {
            namespace_id,
            key,
            value,
            expected_version,
            ts,
        };

        let result = self.raft.client_write(req).await?;
        Ok(result.data)
    }

    /// Propose a KV delete operation (eventual consistency)
    /// Uses Last-Write-Wins semantics, similar to write_kv
    pub async fn delete_kv(
        &self,
        namespace_id: u32,
        key: Vec<u8>,
        ts: u64,
    ) -> Result<FlareResponse, FlareClientWriteError> {
        let req = FlareRequest::KvDelete {
            namespace_id,
            key,
            ts,
        };

        let result = self.raft.client_write(req).await?;
        Ok(result.data)
    }

    /// Propose a CAS delete operation (strong consistency)
    /// Deletes the key if the version matches expected_version.
    /// If expected_version is 0, deletes if key exists (no version check).
    pub async fn cas_delete(
        &self,
        namespace_id: u32,
        key: Vec<u8>,
        expected_version: u64,
        ts: u64,
    ) -> Result<FlareResponse, FlareClientWriteError> {
        let req = FlareRequest::CasDelete {
            namespace_id,
            key,
            expected_version,
            ts,
        };

        let result = self.raft.client_write(req).await?;
        Ok(result.data)
    }

    /// Read a value directly from the state machine (local read, eventual consistency)
    pub async fn read_kv(&self, namespace_id: u32, key: &[u8]) -> Option<(Vec<u8>, u64)> {
        self.storage.read_kv(namespace_id, key).await
    }

    /// Read a CAS value directly from the state machine (local read)
    /// Returns (value, version, timestamp) if found
    pub async fn read_cas(&self, namespace_id: u32, key: &[u8]) -> Option<(Vec<u8>, u64, u64)> {
        self.storage.read_cas(namespace_id, key).await
    }

    /// Scan eventual-consistency KV data from the replicated state machine.
    pub async fn scan_kv(
        &self,
        namespace_id: u32,
        start_key: &[u8],
        end_key: &[u8],
        limit: usize,
    ) -> Vec<(Vec<u8>, Vec<u8>, u64)> {
        self.storage
            .scan_kv(namespace_id, start_key, end_key, limit)
            .await
    }

    /// Scan strong-consistency CAS data from the replicated state machine.
    pub async fn scan_cas(
        &self,
        namespace_id: u32,
        start_key: &[u8],
        end_key: &[u8],
        limit: usize,
    ) -> Vec<(Vec<u8>, Vec<u8>, u64, u64)> {
        self.storage
            .scan_cas(namespace_id, start_key, end_key, limit)
            .await
    }

    /// Perform a linearizable read (ensures we read from committed state)
    ///
    /// This method first calls `ensure_linearizable()` which confirms this node
    /// has the most recent committed data, then reads from local state machine.
    /// This provides strong consistency guarantees.
    pub async fn linearizable_read_kv(
        &self,
        namespace_id: u32,
        key: &[u8],
    ) -> Result<Option<(Vec<u8>, u64)>, String> {
        // Ensure we have the latest committed state
        self.raft
            .ensure_linearizable()
            .await
            .map_err(|e| format!("Linearizable read failed: {}", e))?;

        // Now safe to read from local state machine
        Ok(self.storage.read_kv(namespace_id, key).await)
    }

    /// Perform a linearizable CAS read (ensures we read from committed state)
    /// Returns (value, version, timestamp) if found
    pub async fn linearizable_read_cas(
        &self,
        namespace_id: u32,
        key: &[u8],
    ) -> Result<Option<(Vec<u8>, u64, u64)>, String> {
        // Ensure we have the latest committed state
        self.raft
            .ensure_linearizable()
            .await
            .map_err(|e| format!("Linearizable read failed: {}", e))?;

        // Now safe to read from local state machine
        Ok(self.storage.read_cas(namespace_id, key).await)
    }

    /// Propose a split operation
    pub async fn propose_split(
        &self,
        region_id: u64,
        split_key: Vec<u8>,
        new_region_id: u64,
    ) -> Result<FlareResponse, FlareClientWriteError> {
        let req = FlareRequest::Split {
            region_id,
            split_key,
            new_region_id,
        };

        let result = self.raft.client_write(req).await?;
        Ok(result.data)
    }

    /// Trigger a leader election (replaces peer.campaign)
    pub async fn trigger_election(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        self.raft.trigger().elect().await?;
        Ok(())
    }

    /// Add a learner node (first step in membership change)
    pub async fn add_learner(
        &self,
        node_id: FlareNodeId,
        node: FlareNode,
    ) -> Result<(), FlareClientWriteError> {
        self.raft.add_learner(node_id, node, true).await?;
        Ok(())
    }

    /// Change the cluster membership
    pub async fn change_membership(
        &self,
        members: BTreeSet<FlareNodeId>,
    ) -> Result<(), FlareClientWriteError> {
        self.raft.change_membership(members, false).await?;
        Ok(())
    }

    /// Add a voter to the cluster (convenience method)
    pub async fn add_voter(&self, node_id: FlareNodeId, addr: String) -> Result<(), String> {
        let node = FlareNode {
            addr: addr.clone(),
            region_id: self.region_id,
        };

        // First add as learner
        self.add_learner(node_id, node)
            .await
            .map_err(|e| e.to_string())?;

        // Then promote to voter
        let mut members = self.get_membership().await;
        members.insert(node_id);

        self.change_membership(members)
            .await
            .map_err(|e| e.to_string())?;

        // Register with network factory
        self.network_factory.add_peer(node_id, addr).await;

        Ok(())
    }

    /// Remove a voter from the cluster
    pub async fn remove_voter(&self, node_id: FlareNodeId) -> Result<(), String> {
        let mut members = self.get_membership().await;
        members.remove(&node_id);

        self.change_membership(members)
            .await
            .map_err(|e| e.to_string())?;

        // Remove from network factory
        self.network_factory.remove_peer(node_id).await;

        Ok(())
    }

    /// Get the network factory
    pub fn network_factory(&self) -> &FlareNetworkFactory {
        &self.network_factory
    }

    /// Get the storage instance
    pub fn storage(&self) -> &FlareStorage {
        &self.storage
    }

    /// Shutdown the Raft node gracefully
    pub async fn shutdown(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        self.raft.shutdown().await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    async fn create_test_node() -> FlareRaftNode {
        let network_factory = FlareNetworkFactory::new(1, 1);
        let mut peer_addrs = HashMap::new();
        peer_addrs.insert(1, "127.0.0.1:50051".to_string());

        FlareRaftNode::new(1, 1, network_factory, vec![1], &peer_addrs)
            .await
            .unwrap()
    }

    #[tokio::test]
    async fn test_create_node() {
        let node = create_test_node().await;
        assert_eq!(node.store_id(), 1);
        assert_eq!(node.region_id(), 1);
    }

    #[tokio::test]
    async fn test_leader_election() {
        let node = create_test_node().await;

        // In a single-node cluster, the node should become leader after triggering election
        node.trigger_election().await.unwrap();

        // Give some time for election
        tokio::time::sleep(Duration::from_millis(500)).await;

        // Should be leader in a single-node cluster
        assert!(node.is_leader().await);
    }
}
