use crate::config::{Config, ConsistencyMode, NamespaceManager}; // Updated imports
use flaredb_raft::{FlareNetworkFactory, FlareRaftNode};
use flaredb_storage::rocks_engine::RocksEngine;
use flaredb_types::RegionMeta;
use std::collections::HashMap;
use std::sync::Arc;

pub struct Store {
    store_id: u64,
    engine: Arc<RocksEngine>,
    raft_nodes: tokio::sync::RwLock<HashMap<u64, Arc<FlareRaftNode>>>,
    regions: tokio::sync::RwLock<Vec<RegionMeta>>,
    config: Arc<Config>,                          // Use Config
    pub namespace_manager: Arc<NamespaceManager>, // Make public
    peer_addrs: Arc<HashMap<u64, String>>,
}

impl Store {
    const META_KEY: &'static [u8] = b"_region_meta";

    pub fn new(
        store_id: u64,
        engine: Arc<RocksEngine>,
        config: Arc<Config>,                      // Use Config
        namespace_manager: Arc<NamespaceManager>, // New argument
        peer_addrs: Arc<HashMap<u64, String>>,
    ) -> Self {
        Self {
            store_id,
            engine,
            raft_nodes: tokio::sync::RwLock::new(HashMap::new()),
            regions: tokio::sync::RwLock::new(Vec::new()),
            config,
            namespace_manager, // Initialize new field
            peer_addrs,
        }
    }

    fn validate_regions(regions: &[(RegionMeta, Vec<u64>)]) -> Result<(), String> {
        let mut metas: Vec<&RegionMeta> = regions.iter().map(|(m, _)| m).collect();
        metas.sort_by(|a, b| a.start_key.cmp(&b.start_key));
        for win in metas.windows(2) {
            let a = &win[0];
            let b = &win[1];
            let a_end = if a.end_key.is_empty() {
                None
            } else {
                Some(a.end_key.as_slice())
            };
            if let Some(ae) = a_end {
                if ae > b.start_key.as_slice() {
                    return Err(format!(
                        "region overlap: {} [{:?},{:?}) vs {} [{:?},{:?})",
                        a.id, a.start_key, a.end_key, b.id, b.start_key, b.end_key
                    ));
                }
            } else {
                return Err(format!(
                    "region {} has infinite end but not last (next id {})",
                    a.id, b.id
                ));
            }
        }
        Ok(())
    }

    pub async fn bootstrap_regions(
        &self,
        regions_with_voters: Vec<(RegionMeta, Vec<u64>)>,
    ) -> Result<(), String> {
        Self::validate_regions(&regions_with_voters)?;
        let mut raft_nodes = self.raft_nodes.write().await;
        let mut regions = self.regions.write().await;

        for (meta, voters) in regions_with_voters {
            if raft_nodes.contains_key(&meta.id) {
                continue;
            }

            let network_factory = FlareNetworkFactory::new(self.store_id, meta.id);
            // Register peer addresses with the network factory
            for (&peer_id, addr) in self.peer_addrs.iter() {
                network_factory.add_peer(peer_id, addr.clone()).await;
            }
            // Use persistent storage backed by RocksDB
            let raft_node = FlareRaftNode::new_persistent(
                self.store_id,
                meta.id,
                network_factory,
                voters,
                &self.peer_addrs,
                self.engine.clone(),
            )
            .await
            .map_err(|e| format!("Failed to create raft node: {}", e))?;

            raft_nodes.insert(meta.id, Arc::new(raft_node));
            regions.push(meta);
        }

        regions.sort_by(|a, b| a.start_key.cmp(&b.start_key));
        self.persist_regions(&regions).map_err(|e| e.to_string())?;
        Ok(())
    }

    pub async fn get_raft_node(&self, region_id: u64) -> Option<Arc<FlareRaftNode>> {
        self.raft_nodes.read().await.get(&region_id).cloned()
    }

    pub async fn any_raft_node(&self) -> Option<Arc<FlareRaftNode>> {
        self.raft_nodes.read().await.values().next().cloned()
    }

    pub async fn list_region_ids(&self) -> Vec<u64> {
        self.regions.read().await.iter().map(|r| r.id).collect()
    }

    pub async fn route_key(&self, key: &[u8]) -> Option<u64> {
        let regions = self.regions.read().await;
        for r in regions.iter() {
            let start_ok = key >= r.start_key.as_slice();
            let end_ok = r.end_key.is_empty() || key < r.end_key.as_slice();
            if start_ok && end_ok {
                return Some(r.id);
            }
        }
        None
    }

    pub async fn refresh_regions(&self, metas: Vec<(RegionMeta, Vec<u64>)>) -> Result<(), String> {
        Self::validate_regions(&metas)?;
        self.bootstrap_regions(metas.clone()).await?;
        let mut regions = self.regions.write().await;
        regions.clear();
        for (meta, _) in metas {
            regions.push(meta);
        }
        regions.sort_by(|a, b| a.start_key.cmp(&b.start_key));
        self.persist_regions(&regions).map_err(|e| e.to_string())?;
        Ok(())
    }

    pub fn config(&self) -> Arc<Config> {
        self.config.clone()
    }

    pub fn store_id(&self) -> u64 {
        self.store_id
    }

    pub fn get_namespace_mode(&self, namespace: &str) -> ConsistencyMode {
        self.namespace_manager.get_namespace(namespace).mode
    }

    #[allow(dead_code)]
    pub fn engine(&self) -> Arc<RocksEngine> {
        self.engine.clone()
    }

    pub async fn split_region(
        &self,
        region_id: u64,
        split_key: Vec<u8>,
        new_region_id: u64,
        voters: Vec<u64>,
    ) -> Result<(), String> {
        let mut regions = self.regions.write().await;
        let pos = regions
            .iter()
            .position(|r| r.id == region_id)
            .ok_or_else(|| {
                format!(
                    "region {} not found for split (split_key {:?})",
                    region_id, split_key
                )
            })?;
        let target = regions[pos].clone();
        if split_key.is_empty()
            || (!target.end_key.is_empty() && split_key >= target.end_key)
            || split_key <= target.start_key
        {
            return Err("invalid split key".to_string());
        }

        let left = RegionMeta {
            id: target.id,
            start_key: target.start_key.clone(),
            end_key: split_key.clone(),
        };
        let right = RegionMeta {
            id: new_region_id,
            start_key: split_key,
            end_key: target.end_key.clone(),
        };

        let mut new_layout: Vec<(RegionMeta, Vec<u64>)> = regions
            .iter()
            .enumerate()
            .filter_map(|(i, r)| {
                if i == pos {
                    None
                } else {
                    Some((r.clone(), voters.clone()))
                }
            })
            .collect();
        new_layout.push((left.clone(), voters.clone()));
        new_layout.push((right.clone(), voters.clone()));
        Self::validate_regions(&new_layout)?;

        regions.remove(pos);
        regions.push(left);
        regions.push(right);
        regions.sort_by(|a, b| a.start_key.cmp(&b.start_key));
        drop(regions);

        // Create OpenRaft node for new region
        let mut raft_nodes = self.raft_nodes.write().await;
        if let std::collections::hash_map::Entry::Vacant(e) = raft_nodes.entry(new_region_id) {
            let network_factory = FlareNetworkFactory::new(self.store_id, new_region_id);
            // Register peer addresses with the network factory
            for (&peer_id, addr) in self.peer_addrs.iter() {
                network_factory.add_peer(peer_id, addr.clone()).await;
            }
            // Use persistent storage backed by RocksDB
            let raft_node = FlareRaftNode::new_persistent(
                self.store_id,
                new_region_id,
                network_factory,
                voters,
                &self.peer_addrs,
                self.engine.clone(),
            )
            .await
            .map_err(|e| format!("Failed to create raft node: {}", e))?;
            e.insert(Arc::new(raft_node));
        }
        drop(raft_nodes);

        let updated_regions = self.regions.read().await;
        self.persist_regions(&updated_regions)
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    fn persist_regions(&self, regions: &[RegionMeta]) -> Result<(), rocksdb::Error> {
        let mut buf = Vec::new();
        let len = regions.len() as u32;
        buf.extend_from_slice(&len.to_be_bytes());
        for r in regions {
            buf.extend_from_slice(&r.id.to_be_bytes());
            buf.extend_from_slice(&(r.start_key.len() as u32).to_be_bytes());
            buf.extend_from_slice(&(r.end_key.len() as u32).to_be_bytes());
            buf.extend_from_slice(&r.start_key);
            buf.extend_from_slice(&r.end_key);
        }
        self.engine.db().put(Self::META_KEY, buf)
    }

    #[allow(dead_code)]
    pub async fn allocate_region_id(&self) -> u64 {
        let regions = self.regions.read().await;
        regions.iter().map(|r| r.id).max().unwrap_or(0) + 1
    }
}
