use std::sync::Arc;

use pd_proto_tso::Timestamp;

use crate::election::ElectionState;
use crate::error::ServiceError;
use crate::keypath::{election_path, expected_primary_path};
use crate::oracle::{OracleConfig, TimestampOracle};
use crate::persistence::TimestampStore;

/// Group-allocator runtime configuration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AllocatorConfig {
    /// Lease TTL for one elected group primary in milliseconds.
    pub lease_ms: i64,
    /// Timestamp-oracle runtime configuration.
    pub oracle: OracleConfig,
}

impl Default for AllocatorConfig {
    fn default() -> Self {
        Self {
            lease_ms: 5_000,
            oracle: OracleConfig::default(),
        }
    }
}

/// Allocator serving one keyspace-group timeline.
pub struct GroupAllocator {
    /// Keyspace-group id served by this allocator instance.
    pub group_id: u32,
    cluster_id: u64,
    member_value: String,
    store: Arc<dyn TimestampStore>,
    election: ElectionState,
    oracle: TimestampOracle,
    config: AllocatorConfig,
}

impl GroupAllocator {
    /// Creates one allocator instance for a single keyspace-group timeline.
    #[must_use]
    pub fn new(
        cluster_id: u64,
        group_id: u32,
        member_value: String,
        store: Arc<dyn TimestampStore>,
        config: AllocatorConfig,
    ) -> Self {
        let leader_key = election_path(cluster_id, group_id);
        let expected_primary_key = expected_primary_path(cluster_id, group_id);
        Self {
            group_id,
            cluster_id,
            member_value,
            store,
            election: ElectionState::new(leader_key, expected_primary_key),
            oracle: TimestampOracle::new(config.oracle),
            config,
        }
    }

    /// Returns current allocator runtime configuration.
    #[must_use]
    pub fn config(&self) -> AllocatorConfig {
        self.config
    }

    /// Returns allocator member identity used in election records.
    #[must_use]
    pub fn member_value(&self) -> &str {
        &self.member_value
    }

    /// Returns election leader key path for this allocator.
    #[must_use]
    pub fn leader_key(&self) -> &str {
        self.election.leader_key()
    }

    /// Returns expected-primary key path for this allocator.
    #[must_use]
    pub fn expected_primary_key(&self) -> &str {
        self.election.expected_primary_key()
    }

    /// Returns whether a primary identity is currently elected.
    #[must_use]
    pub fn is_primary_elected(&self) -> bool {
        self.election.is_primary_elected()
    }

    /// Returns whether allocator may serve requests at `now_ms`.
    #[must_use]
    pub fn is_serving(&self, now_ms: i64) -> bool {
        self.election.is_serving(&self.member_value, now_ms)
    }

    /// Returns allocator in-memory timestamp state.
    #[must_use]
    pub fn current_timestamp(&self) -> Timestamp {
        self.oracle.current_timestamp()
    }

    /// Returns local lease expiration timestamp in milliseconds.
    #[must_use]
    pub fn lease_expire_at_ms(&self) -> i64 {
        self.election.lease_expire_at_ms()
    }

    /// Returns persisted save-window upper bound in milliseconds.
    #[must_use]
    pub fn save_window_end_ms(&self) -> i64 {
        self.oracle.save_window_end_ms()
    }

    /// Applies dynamic allocator/oracle configuration updates.
    pub fn apply_config(&mut self, config: AllocatorConfig) {
        self.config = config;
        self.oracle.apply_config(config.oracle);
    }

    /// Campaigns local allocator as primary and persists leader fence when successful.
    pub fn campaign_primary(&mut self, now_ms: i64) -> Result<bool, ServiceError> {
        let campaigned = self
            .election
            .campaign(&self.member_value, now_ms, self.config.lease_ms);
        if campaigned {
            self.store
                .set_leader_value(self.election.leader_key(), &self.member_value)
                .map_err(map_store_err)?;
        }
        Ok(campaigned)
    }

    /// Campaigns local allocator state after an external election success.
    pub fn campaign_primary_external(&mut self, now_ms: i64) -> Result<bool, ServiceError> {
        Ok(self
            .election
            .campaign(&self.member_value, now_ms, self.config.lease_ms))
    }

    /// Renews in-memory lease state after successful keepalive.
    pub fn keep_alive(&mut self, now_ms: i64) {
        self.election.keep_alive(now_ms);
    }

    /// Updates expected-primary gate used by election campaign checks.
    pub fn set_expected_primary(&mut self, expected: Option<String>) {
        self.election.set_expected_primary(expected);
    }

    /// Steps down local serving state and resets in-memory oracle state.
    pub fn step_down(&mut self) {
        self.election.step_down();
        self.oracle.reset();
    }

    /// Initializes allocator from persisted window and persists a fresh save window fence.
    pub fn initialize(&mut self, now_ms: i64) -> Result<(), ServiceError> {
        let loaded = self
            .store
            .load_timestamp(self.cluster_id, self.group_id)
            .map_err(map_store_err)?;
        let store = Arc::clone(&self.store);
        let cluster_id = self.cluster_id;
        let group_id = self.group_id;
        let leader_key = self.election.leader_key();
        let expected = self.election.primary_value().unwrap_or_default();

        self.oracle.sync_timestamp(now_ms, loaded, |save_ms| {
            save_window_ms(
                store.as_ref(),
                cluster_id,
                group_id,
                save_ms,
                leader_key,
                expected,
            )
        })
    }

    /// Runs periodic timestamp-window maintenance for serving allocator.
    pub fn update_tso(&mut self, now_ms: i64) -> Result<(), ServiceError> {
        if !self.is_serving(now_ms) {
            return Err(ServiceError::NotLeader);
        }

        let store = Arc::clone(&self.store);
        let cluster_id = self.cluster_id;
        let group_id = self.group_id;
        let leader_key = self.election.leader_key();
        let expected = self.election.primary_value().unwrap_or_default();
        self.oracle.update_interval(now_ms, |save_ms| {
            save_window_ms(
                store.as_ref(),
                cluster_id,
                group_id,
                save_ms,
                leader_key,
                expected,
            )
        })
    }

    /// Allocates one timestamp range of `count` logical slots for this group timeline.
    pub fn generate_tso(&mut self, now_ms: i64, count: u32) -> Result<Timestamp, ServiceError> {
        if !self.is_serving(now_ms) {
            return Err(ServiceError::NotLeader);
        }
        let store = Arc::clone(&self.store);
        let cluster_id = self.cluster_id;
        let group_id = self.group_id;
        let leader_key = self.election.leader_key();
        let expected = self.election.primary_value().unwrap_or_default();

        self.oracle.generate_tso(now_ms, count, true, |save_ms| {
            save_window_ms(
                store.as_ref(),
                cluster_id,
                group_id,
                save_ms,
                leader_key,
                expected,
            )
        })
    }

    /// Rebases allocator timeline to an explicit TSO value.
    pub fn set_tso(
        &mut self,
        now_ms: i64,
        tso: u64,
        ignore_smaller: bool,
        skip_upper_bound_check: bool,
    ) -> Result<(), ServiceError> {
        if !self.is_serving(now_ms) {
            return Err(ServiceError::NotLeader);
        }

        let store = Arc::clone(&self.store);
        let cluster_id = self.cluster_id;
        let group_id = self.group_id;
        let leader_key = self.election.leader_key();
        let expected = self.election.primary_value().unwrap_or_default();

        self.oracle
            .set_tso(tso, ignore_smaller, skip_upper_bound_check, |save_ms| {
                save_window_ms(
                    store.as_ref(),
                    cluster_id,
                    group_id,
                    save_ms,
                    leader_key,
                    expected,
                )
            })
    }
}

fn save_window_ms(
    store: &dyn TimestampStore,
    cluster_id: u64,
    group_id: u32,
    save_ms: i64,
    leader_key: &str,
    expected_leader_value: &str,
) -> Result<(), ServiceError> {
    let save_nanos = save_ms.saturating_mul(1_000_000);
    store
        .save_timestamp_if_leader(
            cluster_id,
            group_id,
            save_nanos,
            leader_key,
            expected_leader_value,
        )
        .map_err(map_store_err)
}

fn map_store_err(err: impl std::error::Error) -> ServiceError {
    ServiceError::Storage(err.to_string())
}

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

    use super::{AllocatorConfig, GroupAllocator};
    use crate::persistence::InMemoryEtcdStore;

    #[test]
    fn allocator_requires_campaign_for_serving() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut allocator = GroupAllocator::new(
            1,
            0,
            "node-a".to_string(),
            store,
            AllocatorConfig::default(),
        );

        let err = allocator
            .generate_tso(1_000, 1)
            .expect_err("not serving must fail");
        assert_eq!(err.to_string(), "requested pd is not leader of cluster");
    }

    #[test]
    fn allocator_campaign_init_and_generate() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut allocator = GroupAllocator::new(
            1,
            7,
            "node-a".to_string(),
            store,
            AllocatorConfig::default(),
        );

        assert!(allocator.campaign_primary(10_000).expect("campaign"));
        allocator.initialize(10_000).expect("initialize");
        let ts = allocator.generate_tso(10_000, 3).expect("generate");

        assert_eq!(ts.logical, 3);
        assert_eq!(ts.physical, 10_000);
    }
}