use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};

use tokio::sync::{Mutex, Notify, RwLock, mpsc};
use tokio::task::JoinHandle;

use crate::types::{ClientId, GroupName, MessageId, PublishedMessage, TopicName, TopicOffset};
use crate::wildcard::{SubscriptionId, SubscriptionTrie};

/// Tracks an individual message that has been delivered but not yet acked.
#[derive(Debug)]
pub struct InflightEntry {
    pub offset: TopicOffset,
    pub message_id: MessageId,
    pub delivered_to: ClientId,
    pub delivered_at: Instant,
    pub delivery_count: u32,
    /// Wall-clock timestamp (millis since epoch) of the first delivery attempt.
    /// Set on first delivery (`delivery_count == 1`), preserved across redeliveries.
    /// Used by the DLQ envelope to report total time-in-flight.
    pub first_delivered_at: u64,
    /// Number of consecutive DLQ write failures for this message.
    /// Used to escalate logging from warn to error after repeated failures.
    pub dlq_failures: u32,
    /// When this entry becomes eligible for redelivery. Each redelivery uses
    /// exponential backoff: `min(base * 2^(count-1), max_backoff) * jitter`.
    pub next_redeliver_at: Instant,
    /// Original message timestamp (millis since epoch). Stored here so the
    /// inflight sweep can check TTL expiry without re-reading from the engine.
    pub msg_timestamp: u64,
    /// Original message TTL in milliseconds. 0 = no expiry.
    pub msg_ttl_ms: u32,
}

impl InflightEntry {
    /// Check whether the message tracked by this inflight entry has expired.
    pub fn is_expired(&self, now_ms: u64) -> bool {
        self.msg_ttl_ms > 0 && self.msg_timestamp + (self.msg_ttl_ms as u64) <= now_ms
    }
}

impl InflightEntry {
    /// Calculate the redelivery delay for the *next* attempt using exponential
    /// backoff with ±25% jitter.
    ///
    /// `delivery_count` is the count *after* this delivery (i.e., the attempt
    /// number that just happened). The exponent is `count - 1` so the first
    /// delivery uses `base_timeout` as-is.
    pub fn redelivery_delay(
        delivery_count: u32,
        base_timeout: Duration,
        max_backoff: Duration,
    ) -> Duration {
        let exponent = (delivery_count.saturating_sub(1)).min(10);
        let backoff = base_timeout.saturating_mul(1u32 << exponent);
        let capped = backoff.min(max_backoff);
        // Jitter: multiply by a random factor in [0.75, 1.25)
        let jitter = 0.75 + fastrand::f64() * 0.5;
        capped.mul_f64(jitter)
    }
}

/// Per-subscriber delivery state: cursor position, inflight messages, and the
/// channel used to push messages to the subscriber's connection handler.
pub struct SubscriberDeliveryState {
    pub next_offset: TopicOffset,
    pub inflight: HashMap<TopicOffset, InflightEntry>,
    pub tx: mpsc::Sender<PublishedMessage>,
    pub last_activity: Instant,
    /// Number of consecutive `try_send` failures (channel full).
    pub consecutive_send_failures: u64,
    /// Total number of messages dropped due to slow consumer policy.
    pub total_drops: u64,
    /// Whether this subscriber has been marked as slow.
    pub is_slow: bool,
}

impl SubscriberDeliveryState {
    pub fn new(offset: TopicOffset, tx: mpsc::Sender<PublishedMessage>) -> Self {
        Self {
            next_offset: offset,
            inflight: HashMap::new(),
            tx,
            last_activity: Instant::now(),
            consecutive_send_failures: 0,
            total_drops: 0,
            is_slow: false,
        }
    }
}

/// Shared state for a consumer group on a single topic.
pub struct ConsumerGroupState {
    pub name: GroupName,
    pub members: HashMap<ClientId, mpsc::Sender<PublishedMessage>>,
    pub next_offset: TopicOffset,
    pub inflight: HashMap<TopicOffset, InflightEntry>,
    pub round_robin_counter: AtomicU64,
    pub last_activity: Instant,
    /// Set when the group becomes empty; cleared when a new member joins.
    /// Used for grace period cleanup of empty groups.
    pub emptied_at: Option<Instant>,
    /// Number of consecutive `try_send` failures (channel full).
    pub consecutive_send_failures: u64,
    /// Total number of messages dropped due to slow consumer policy.
    pub total_drops: u64,
    /// Whether this consumer group has been marked as slow.
    pub is_slow: bool,
}

impl ConsumerGroupState {
    pub fn new(name: GroupName, offset: TopicOffset) -> Self {
        Self {
            name,
            members: HashMap::new(),
            next_offset: offset,
            inflight: HashMap::new(),
            round_robin_counter: AtomicU64::new(0),
            last_activity: Instant::now(),
            emptied_at: None,
            consecutive_send_failures: 0,
            total_drops: 0,
            is_slow: false,
        }
    }

    /// Pick the next member via round-robin. Returns `None` if no members.
    pub fn pick_member(&self) -> Option<(&ClientId, &mpsc::Sender<PublishedMessage>)> {
        if self.members.is_empty() {
            return None;
        }
        let members: Vec<_> = self.members.iter().collect();
        let idx = self.round_robin_counter.fetch_add(1, Ordering::Relaxed) as usize % members.len();
        Some(members[idx])
    }
}

/// Delivery state for a single topic — individual subscribers and consumer groups.
pub struct TopicDeliveryState {
    pub individual: HashMap<ClientId, SubscriberDeliveryState>,
    pub groups: HashMap<GroupName, ConsumerGroupState>,
    pub notify: Arc<Notify>,
    /// JoinHandles for individual subscriber delivery tasks.
    pub individual_handles: HashMap<ClientId, JoinHandle<()>>,
    /// JoinHandles for consumer group delivery tasks.
    pub group_handles: HashMap<GroupName, JoinHandle<()>>,
    /// Cached head offset for this topic, updated on publish notification.
    /// Used by the ack path to compute consumer lag without an engine lock.
    pub cached_head_offset: TopicOffset,
    /// Wall-clock timestamp (millis since epoch) of the last subscribe or fetch
    /// operation on this topic. Initialized to 0 (no activity). Used by the
    /// inactivity GC sweep to determine whether a topic is eligible for reaping.
    pub last_subscribe_millis: u64,
}

impl Default for TopicDeliveryState {
    fn default() -> Self {
        Self::new()
    }
}

impl TopicDeliveryState {
    pub fn new() -> Self {
        Self {
            individual: HashMap::new(),
            groups: HashMap::new(),
            notify: Arc::new(Notify::new()),
            individual_handles: HashMap::new(),
            group_handles: HashMap::new(),
            cached_head_offset: TopicOffset::ZERO,
            last_subscribe_millis: 0,
        }
    }

    /// Scan all individual and group cursors to find the minimum live offset.
    /// Returns `None` when there are no subscribers or groups.
    pub fn minimum_live_offset(&self) -> Option<TopicOffset> {
        let individual_min = self.individual.values().map(|s| s.next_offset).min();

        let inflight_individual_min = self
            .individual
            .values()
            .flat_map(|s| s.inflight.keys().copied())
            .min();

        let group_min = self.groups.values().map(|g| g.next_offset).min();

        let inflight_group_min = self
            .groups
            .values()
            .flat_map(|g| g.inflight.keys().copied())
            .min();

        [
            individual_min,
            inflight_individual_min,
            group_min,
            inflight_group_min,
        ]
        .into_iter()
        .flatten()
        .min()
    }

    /// Remove inflight entries past their individual `next_redeliver_at` deadline
    /// from all subscribers and groups. This prevents expired entries from
    /// blocking compaction forever.
    pub fn evict_expired_inflight(&mut self) {
        let now = Instant::now();
        for sub in self.individual.values_mut() {
            sub.inflight
                .retain(|_, entry| now < entry.next_redeliver_at);
        }
        for group in self.groups.values_mut() {
            group
                .inflight
                .retain(|_, entry| now < entry.next_redeliver_at);
        }
    }

    /// Reset all inflight entries' `next_redeliver_at` to `now + timeout`.
    /// Called during drain to ensure long-backoff entries are flushed promptly.
    pub fn reset_inflight_deadlines(&mut self, timeout: Duration) {
        let deadline = Instant::now() + timeout;
        for sub in self.individual.values_mut() {
            for entry in sub.inflight.values_mut() {
                if entry.next_redeliver_at > deadline {
                    entry.next_redeliver_at = deadline;
                }
            }
        }
        for group in self.groups.values_mut() {
            for entry in group.inflight.values_mut() {
                if entry.next_redeliver_at > deadline {
                    entry.next_redeliver_at = deadline;
                }
            }
        }
    }

    /// Abort all delivery task handles (used during shutdown).
    pub fn abort_all_handles(&mut self) {
        for (_, handle) in self.individual_handles.drain() {
            handle.abort();
        }
        for (_, handle) in self.group_handles.drain() {
            handle.abort();
        }
    }
}

/// Info needed to auto-subscribe a wildcard subscriber to a newly created topic.
/// Stored in the reverse index keyed by `SubscriptionId`.
/// Used for **individual** (non-group) wildcard subscriptions only.
pub struct WildcardSubInfo {
    pub client: ClientId,
    /// The pattern string as a `TopicName` — needed for metric labels and logging.
    pub pattern_topic: TopicName,
    /// Shared delivery channel sender — all per-topic forwarders fan into this.
    pub tx: mpsc::Sender<PublishedMessage>,
}

/// Info needed to auto-subscribe wildcard consumer group members to newly
/// created topics. One entry per `(pattern, group)` combination.
///
/// When a new topic matches the pattern, all current members are added to a
/// new consumer group for that topic.
pub struct WildcardGroupSubInfo {
    /// The pattern string as a `TopicName` — for metric labels and logging.
    pub pattern_topic: TopicName,
    /// The consumer group name shared across all matched topics.
    pub group: GroupName,
    /// Current members of this wildcard group subscription.
    /// Each member has their own delivery channel sender.
    pub members: HashMap<ClientId, mpsc::Sender<PublishedMessage>>,
    /// Concrete topics this wildcard group currently matches.
    pub matched_topics: Vec<TopicName>,
}

/// Top-level state for a single namespace.
pub struct NamespaceState {
    pub topics: RwLock<HashMap<TopicName, Arc<Mutex<TopicDeliveryState>>>>,
    pub message_count: AtomicU64,
    pub topic_count: AtomicUsize,
    /// Number of active connections to this namespace.
    pub connection_count: AtomicUsize,
    /// Suppression flags for 80% threshold warnings (fire once, reset when
    /// usage drops back below threshold).
    pub messages_warning_fired: AtomicBool,
    pub disk_bytes_warning_fired: AtomicBool,
    pub connections_warning_fired: AtomicBool,
    /// Trie for matching concrete topic names against active wildcard
    /// subscription patterns. Read-locked on new topic creation (rare),
    /// write-locked on subscribe/unsubscribe.
    pub subscription_trie: RwLock<SubscriptionTrie>,
    /// Reverse index: SubscriptionId → individual wildcard sub info.
    /// Used during new-topic-creation to find the client and delivery channel
    /// for each matching individual wildcard subscription.
    pub wildcard_subs: RwLock<HashMap<SubscriptionId, WildcardSubInfo>>,
    /// Reverse index: SubscriptionId → wildcard group sub info.
    /// Used during new-topic-creation to find group members and create
    /// consumer groups for newly matched topics. One entry per (pattern, group).
    pub wildcard_group_subs: RwLock<HashMap<SubscriptionId, WildcardGroupSubInfo>>,
}

impl Default for NamespaceState {
    fn default() -> Self {
        Self::new()
    }
}

impl NamespaceState {
    pub fn new() -> Self {
        Self {
            topics: RwLock::new(HashMap::new()),
            message_count: AtomicU64::new(0),
            topic_count: AtomicUsize::new(0),
            connection_count: AtomicUsize::new(0),
            messages_warning_fired: AtomicBool::new(false),
            disk_bytes_warning_fired: AtomicBool::new(false),
            connections_warning_fired: AtomicBool::new(false),
            subscription_trie: RwLock::new(SubscriptionTrie::new()),
            wildcard_subs: RwLock::new(HashMap::new()),
            wildcard_group_subs: RwLock::new(HashMap::new()),
        }
    }
}

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

    #[test]
    fn minimum_live_offset_mixed_cursors() {
        let mut state = TopicDeliveryState::new();

        // Individual subscriber at offset 5
        let (tx, _rx) = mpsc::channel(1);
        state.individual.insert(
            ClientId(1),
            SubscriberDeliveryState::new(TopicOffset(5), tx),
        );

        // Individual subscriber at offset 3
        let (tx2, _rx2) = mpsc::channel(1);
        state.individual.insert(
            ClientId(2),
            SubscriberDeliveryState::new(TopicOffset(3), tx2),
        );

        // Consumer group at offset 7
        let mut group = ConsumerGroupState::new(GroupName("g1".into()), TopicOffset(7));
        let (gtx, _grx) = mpsc::channel(1);
        group.members.insert(ClientId(3), gtx);
        state.groups.insert(GroupName("g1".into()), group);

        assert_eq!(state.minimum_live_offset(), Some(TopicOffset(3)));
    }

    #[test]
    fn minimum_live_offset_includes_inflight() {
        let mut state = TopicDeliveryState::new();

        // Individual subscriber at offset 10 with inflight at offset 2
        let (tx, _rx) = mpsc::channel(1);
        let mut sub = SubscriberDeliveryState::new(TopicOffset(10), tx);
        sub.inflight.insert(
            TopicOffset(2),
            InflightEntry {
                offset: TopicOffset(2),
                message_id: MessageId(1),
                delivered_to: ClientId(1),
                delivered_at: Instant::now(),
                delivery_count: 1,
                first_delivered_at: 0,
                dlq_failures: 0,
                next_redeliver_at: Instant::now() + Duration::from_secs(300),
                msg_timestamp: 0,
                msg_ttl_ms: 0,
            },
        );
        state.individual.insert(ClientId(1), sub);

        assert_eq!(state.minimum_live_offset(), Some(TopicOffset(2)));
    }

    #[test]
    fn minimum_live_offset_empty_returns_none() {
        let state = TopicDeliveryState::new();
        assert_eq!(state.minimum_live_offset(), None);
    }

    #[test]
    fn pick_member_round_robins() {
        let mut group = ConsumerGroupState::new(GroupName("g".into()), TopicOffset::ZERO);
        let (tx1, _rx1) = mpsc::channel(1);
        let (tx2, _rx2) = mpsc::channel(1);
        let (tx3, _rx3) = mpsc::channel(1);
        group.members.insert(ClientId(1), tx1);
        group.members.insert(ClientId(2), tx2);
        group.members.insert(ClientId(3), tx3);

        // Collect several picks and verify we see all members
        let mut seen = HashMap::new();
        for _ in 0..30 {
            let (id, _) = group.pick_member().unwrap();
            *seen.entry(*id).or_insert(0u32) += 1;
        }

        // Each member should be picked exactly 10 times
        assert_eq!(seen.len(), 3);
        for &count in seen.values() {
            assert_eq!(count, 10);
        }
    }

    #[test]
    fn pick_member_empty_returns_none() {
        let group = ConsumerGroupState::new(GroupName("g".into()), TopicOffset::ZERO);
        assert!(group.pick_member().is_none());
    }

    #[test]
    fn evict_expired_inflight_removes_old_entries() {
        let mut state = TopicDeliveryState::new();

        let (tx, _rx) = mpsc::channel(1);
        let mut sub = SubscriberDeliveryState::new(TopicOffset(10), tx);
        // Insert an inflight entry with next_redeliver_at in the past
        sub.inflight.insert(
            TopicOffset(2),
            InflightEntry {
                offset: TopicOffset(2),
                message_id: MessageId(1),
                delivered_to: ClientId(1),
                delivered_at: Instant::now(),
                delivery_count: 1,
                first_delivered_at: 0,
                dlq_failures: 0,
                // Already expired
                next_redeliver_at: Instant::now() - Duration::from_secs(1),
                msg_timestamp: 0,
                msg_ttl_ms: 0,
            },
        );
        state.individual.insert(ClientId(1), sub);

        state.evict_expired_inflight();
        assert!(state.individual[&ClientId(1)].inflight.is_empty());
    }

    #[test]
    fn evict_expired_inflight_keeps_fresh_entries() {
        let mut state = TopicDeliveryState::new();

        let (tx, _rx) = mpsc::channel(1);
        let mut sub = SubscriberDeliveryState::new(TopicOffset(10), tx);
        sub.inflight.insert(
            TopicOffset(2),
            InflightEntry {
                offset: TopicOffset(2),
                message_id: MessageId(1),
                delivered_to: ClientId(1),
                delivered_at: Instant::now(),
                delivery_count: 1,
                first_delivered_at: 0,
                dlq_failures: 0,
                next_redeliver_at: Instant::now() + Duration::from_secs(3600),
                msg_timestamp: 0,
                msg_ttl_ms: 0,
            },
        );
        state.individual.insert(ClientId(1), sub);

        state.evict_expired_inflight();
        assert_eq!(state.individual[&ClientId(1)].inflight.len(), 1);
    }

    #[test]
    fn redelivery_delay_first_attempt_equals_base() {
        let base = Duration::from_secs(300);
        let max = Duration::from_secs(3600);
        // delivery_count=1 → exponent=0 → base * 1 * jitter
        let delay = InflightEntry::redelivery_delay(1, base, max);
        // With jitter [0.75, 1.25), delay should be in [225, 375)
        assert!(
            delay >= Duration::from_millis(224_000),
            "delay too small: {delay:?}"
        );
        assert!(
            delay < Duration::from_millis(376_000),
            "delay too large: {delay:?}"
        );
    }

    #[test]
    fn redelivery_delay_doubles_each_attempt() {
        let base = Duration::from_secs(10);
        let max = Duration::from_secs(3600);
        // Collect many samples to verify the mean roughly doubles
        let samples = 1000;
        let mean_1: f64 = (0..samples)
            .map(|_| InflightEntry::redelivery_delay(1, base, max).as_secs_f64())
            .sum::<f64>()
            / samples as f64;
        let mean_2: f64 = (0..samples)
            .map(|_| InflightEntry::redelivery_delay(2, base, max).as_secs_f64())
            .sum::<f64>()
            / samples as f64;
        let ratio = mean_2 / mean_1;
        assert!(
            ratio > 1.7 && ratio < 2.3,
            "ratio should be ~2.0, got {ratio}"
        );
    }

    #[test]
    fn redelivery_delay_caps_at_max_backoff() {
        let base = Duration::from_secs(300);
        let max = Duration::from_secs(3600);
        // delivery_count=20 → exponent capped at 10 → 300 * 1024 = 307200s → capped to 3600s
        for _ in 0..100 {
            let delay = InflightEntry::redelivery_delay(20, base, max);
            // max * 1.25 jitter = 4500
            assert!(
                delay <= Duration::from_secs(4500),
                "delay exceeded max with jitter: {delay:?}"
            );
            // max * 0.75 jitter = 2700
            assert!(
                delay >= Duration::from_millis(2_700_000),
                "delay below min jitter: {delay:?}"
            );
        }
    }

    #[test]
    fn redelivery_delay_jitter_within_bounds() {
        let base = Duration::from_secs(100);
        let max = Duration::from_secs(10_000);
        // delivery_count=1 → base * jitter. Check many samples stay in [75, 125).
        for _ in 0..1000 {
            let delay = InflightEntry::redelivery_delay(1, base, max);
            assert!(delay >= Duration::from_secs(75));
            assert!(delay < Duration::from_secs(126));
        }
    }
}
