use flaredb_proto::pdpb::{Region, Store};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

#[derive(Clone)]
pub struct Cluster {
    inner: Arc<Mutex<ClusterState>>,
}

struct ClusterState {
    stores: HashMap<u64, Store>,
    regions: Vec<Region>, // Sorted by start_key? Or just linear scan for MVP
    region_leaders: HashMap<u64, u64>, // region_id -> store_id
    next_store_id: u64,
    next_region_id: u64,
}

impl Cluster {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(ClusterState {
                stores: HashMap::new(),
                regions: Vec::new(),
                region_leaders: HashMap::new(),
                next_store_id: 1,
                next_region_id: 1,
            })),
        }
    }

    pub fn register_store(&self, addr: String, requested_id: Option<u64>) -> u64 {
        let mut state = self.inner.lock().unwrap();

        // Dedup check? For now, always new ID.
        // In real system, check if addr exists.
        for store in state.stores.values() {
            if store.addr == addr {
                return store.id;
            }
        }

        let id = requested_id
            .filter(|id| *id != 0 && !state.stores.contains_key(id))
            .unwrap_or_else(|| {
                while state.stores.contains_key(&state.next_store_id) {
                    state.next_store_id += 1;
                }
                state.next_store_id
            });
        state.next_store_id = state.next_store_id.max(id.saturating_add(1));

        state.stores.insert(id, Store { id, addr });

        // Bootstrap check: If we have 3 stores (Raft requires 3 for quorum), create Raft group
        if state.stores.len() == 3 && state.regions.is_empty() {
            let r1_id = state.next_region_id;
            state.next_region_id += 1;

            let r1 = Region {
                id: r1_id,
                start_key: vec![],
                end_key: vec![], // All keys
                peers: Vec::new(),
                leader_id: 0,
            };

            state.regions.push(r1);

            // Assign ALL 3 stores as peers (Leader = S1, but peers = [S1, S2, S3])
            // We need to update Region definition to include Peers?
            // Currently Region struct in pdpb doesn't have peers field (it's commented out).
            // We can just assign leader for routing.
            // Real Raft: Server needs to know peers.
            // PD should return Peers in GetRegionResponse.

            let ids: Vec<u64> = state.stores.keys().cloned().collect();
            let mut sorted_ids = ids;
            sorted_ids.sort();

            state.region_leaders.insert(r1_id, sorted_ids[0]);

            println!(
                "Bootstrapped Raft Cluster: Region {} -> Leader S{}, Peers {:?}",
                r1_id, sorted_ids[0], sorted_ids
            );
        } else if state.stores.len() == 1 && state.regions.is_empty() {
            // Single node: one region ["", "")
            let r1_id = state.next_region_id;
            state.next_region_id += 1;

            let r1 = Region {
                id: r1_id,
                start_key: vec![],
                end_key: vec![],
                peers: Vec::new(),
                leader_id: 0,
            };
            state.regions.push(r1);

            // Assign to this store
            state.region_leaders.insert(r1_id, id);
            println!("Bootstrapped Single Node: R1->S{}", id);
        }

        id
    }

    pub fn get_region_for_key(&self, key: &[u8]) -> Option<(Region, Store)> {
        let state = self.inner.lock().unwrap();

        for region in &state.regions {
            let start_ok = region.start_key.is_empty() || key >= &region.start_key[..];
            let end_ok = region.end_key.is_empty() || key < &region.end_key[..];

            if start_ok && end_ok {
                // Found region
                if let Some(leader_id) = state.region_leaders.get(&region.id) {
                    if let Some(store) = state.stores.get(leader_id) {
                        let mut r = region.clone();
                        r.leader_id = *leader_id;
                        r.peers = state.stores.keys().cloned().collect();
                        return Some((r, store.clone()));
                    }
                }
            }
        }

        None
    }

    pub fn list_regions(&self) -> (Vec<Region>, Vec<Store>) {
        let state = self.inner.lock().unwrap();
        let mut regions = Vec::new();
        for r in &state.regions {
            let mut clone = r.clone();
            clone.leader_id = *state.region_leaders.get(&r.id).unwrap_or(&0);
            clone.peers = state.stores.keys().cloned().collect();
            regions.push(clone);
        }
        let stores: Vec<Store> = state.stores.values().cloned().collect();
        (regions, stores)
    }
}