//! Idempotency key support for `gateway_invoke`
//!
//! Prevents duplicate side effects when LLMs retry tool calls due to timeouts.
//!
//! # How it works
//!
//! 1. Client supplies an optional `idempotency_key` in `gateway_invoke` arguments.
//! 2. For side-effecting tools without an explicit key, one is auto-generated from
//!    `SHA-256(tool_name || canonical_json(arguments))`.
//! 3. Before dispatch the key is looked up:
//!    - Not found → mark `InFlight`, execute, store `Completed`.
//!    - `InFlight` and not timed-out → return `Err(Error::DuplicateRequest)`.
//!    - `Completed` → return cached result immediately (no re-execution).
//! 4. A background task periodically evicts stale entries to bound memory usage.

use std::sync::Arc;
use std::time::{Duration, Instant};

use dashmap::DashMap;
use serde_json::Value;
use tracing::debug;

use crate::hashing::{canonical_json, sha256_hex_chunks};
use crate::{Error, Result};

// ── Public constants ──────────────────────────────────────────────────────────

/// TTL for completed results (24 hours).
pub const COMPLETED_TTL: Duration = Duration::from_secs(24 * 60 * 60);

/// Timeout for in-flight markers (5 minutes).
///
/// If a tool call does not complete within this window the in-flight marker
/// is treated as stale and a new execution is allowed.
pub const IN_FLIGHT_TIMEOUT: Duration = Duration::from_secs(5 * 60);

// ── State machine ─────────────────────────────────────────────────────────────

/// State of an idempotency entry.
#[derive(Debug, Clone)]
pub enum IdempotencyState {
    /// Tool call is currently executing.  Holds the instant it was registered.
    InFlight(Instant),
    /// Tool call completed successfully.  Holds the result and when it was stored.
    Completed(Value, Instant),
}

impl IdempotencyState {
    /// Return `true` when this entry is stale and should be evicted.
    #[must_use]
    pub fn is_expired(&self) -> bool {
        match self {
            Self::InFlight(started) => started.elapsed() > IN_FLIGHT_TIMEOUT,
            Self::Completed(_, stored) => stored.elapsed() > COMPLETED_TTL,
        }
    }

    /// Return `true` when this is a live in-flight entry (not yet timed out).
    #[must_use]
    pub fn is_in_flight(&self) -> bool {
        matches!(self, Self::InFlight(t) if t.elapsed() <= IN_FLIGHT_TIMEOUT)
    }
}

// ── IdempotencyCache ──────────────────────────────────────────────────────────

/// Thread-safe cache that tracks in-flight and completed idempotent requests.
///
/// All operations are O(1) amortised thanks to the underlying `DashMap`.
///
/// # Example
///
/// ```
/// use mcp_gateway::idempotency::{IdempotencyCache, CheckOutcome};
/// use serde_json::json;
///
/// let cache = IdempotencyCache::new();
/// let key = "my-idempotency-key";
///
/// // Mark as in-flight before dispatching
/// cache.mark_in_flight(key);
///
/// // After the call completes, store the result
/// cache.mark_completed(key, json!({"status": "ok"}));
///
/// // Subsequent calls with the same key return the cached result
/// let result = cache.check(key);
/// assert!(matches!(result, CheckOutcome::Completed(_)));
/// ```
#[derive(Debug, Default)]
pub struct IdempotencyCache {
    entries: DashMap<String, IdempotencyState>,
}

/// Outcome of checking the idempotency cache before executing a tool call.
#[derive(Debug)]
pub enum CheckOutcome {
    /// Key not found (or stale in-flight) — proceed with execution.
    Proceed,
    /// A live in-flight entry exists — reject with 409.
    InFlight,
    /// A completed entry exists — return cached result.
    Completed(Value),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CacheEntryStatus {
    Missing,
    LiveInFlight,
    StaleInFlight,
    LiveCompleted,
    ExpiredCompleted,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CheckPlan {
    Proceed,
    InFlight,
    Completed,
}

#[must_use]
pub(crate) fn decide_check_plan(status: CacheEntryStatus) -> (CheckPlan, bool) {
    match status {
        CacheEntryStatus::Missing => (CheckPlan::Proceed, false),
        CacheEntryStatus::LiveInFlight => (CheckPlan::InFlight, false),
        CacheEntryStatus::StaleInFlight | CacheEntryStatus::ExpiredCompleted => {
            (CheckPlan::Proceed, true)
        }
        CacheEntryStatus::LiveCompleted => (CheckPlan::Completed, false),
    }
}

impl IdempotencyCache {
    /// Create a new, empty cache.
    #[must_use]
    pub fn new() -> Self {
        Self {
            entries: DashMap::new(),
        }
    }

    /// Check the cache state for `key` and return what the caller should do.
    ///
    /// Stale in-flight entries (exceeded [`IN_FLIGHT_TIMEOUT`]) are evicted and
    /// treated as `Proceed` so a fresh execution can start.
    pub fn check(&self, key: &str) -> CheckOutcome {
        let Some(entry) = self.entries.get(key) else {
            let (plan, evict) = decide_check_plan(CacheEntryStatus::Missing);
            debug_assert!(!evict);
            return match plan {
                CheckPlan::Proceed => CheckOutcome::Proceed,
                CheckPlan::InFlight => CheckOutcome::InFlight,
                CheckPlan::Completed => unreachable!("missing entries cannot be completed"),
            };
        };

        let status = match entry.value() {
            IdempotencyState::InFlight(started) if started.elapsed() <= IN_FLIGHT_TIMEOUT => {
                CacheEntryStatus::LiveInFlight
            }
            IdempotencyState::InFlight(_) => CacheEntryStatus::StaleInFlight,
            IdempotencyState::Completed(_, stored) if stored.elapsed() <= COMPLETED_TTL => {
                CacheEntryStatus::LiveCompleted
            }
            IdempotencyState::Completed(_, _) => CacheEntryStatus::ExpiredCompleted,
        };

        let (decision, evict) = decide_check_plan(status);
        if evict {
            drop(entry);
            self.entries.remove(key);
            debug!(key, "Evicted stale idempotency entry");
            return CheckOutcome::Proceed;
        }

        match decision {
            CheckPlan::Proceed => CheckOutcome::Proceed,
            CheckPlan::InFlight => CheckOutcome::InFlight,
            CheckPlan::Completed => {
                let IdempotencyState::Completed(value, _) = entry.value() else {
                    unreachable!("live completed status must hold a completed value");
                };
                CheckOutcome::Completed(value.clone())
            }
        }
    }

    /// Register `key` as in-flight.  Overwrites any stale entry.
    pub fn mark_in_flight(&self, key: &str) {
        self.entries
            .insert(key.to_string(), IdempotencyState::InFlight(Instant::now()));
    }

    /// Transition `key` from in-flight to completed with `result`.
    pub fn mark_completed(&self, key: &str, result: Value) {
        self.entries.insert(
            key.to_string(),
            IdempotencyState::Completed(result, Instant::now()),
        );
    }

    /// Remove `key` entirely (used when a call fails and should be retryable).
    pub fn remove(&self, key: &str) {
        self.entries.remove(key);
    }

    /// Evict all stale entries.  Called by the background maintenance task.
    pub fn evict_expired(&self) {
        let stale: Vec<String> = self
            .entries
            .iter()
            .filter_map(|e| e.value().is_expired().then(|| e.key().clone()))
            .collect();

        let count = stale.len();
        for key in stale {
            self.entries.remove(&key);
        }
        if count > 0 {
            debug!(count, "Evicted stale idempotency entries");
        }
    }

    /// Current number of tracked entries.
    #[must_use]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Return `true` when the cache is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }
}

#[cfg(kani)]
mod verification {
    use super::*;

    fn any_entry_status() -> CacheEntryStatus {
        match kani::any::<u8>() % 5 {
            0 => CacheEntryStatus::Missing,
            1 => CacheEntryStatus::LiveInFlight,
            2 => CacheEntryStatus::StaleInFlight,
            3 => CacheEntryStatus::LiveCompleted,
            _ => CacheEntryStatus::ExpiredCompleted,
        }
    }

    #[kani::proof]
    fn idempotency_decision_contract() {
        let status = any_entry_status();
        let (plan, evict) = decide_check_plan(status);

        match status {
            CacheEntryStatus::Missing => {
                assert_eq!(plan, CheckPlan::Proceed);
                assert!(!evict);
            }
            CacheEntryStatus::LiveInFlight => {
                assert_eq!(plan, CheckPlan::InFlight);
                assert!(!evict);
            }
            CacheEntryStatus::StaleInFlight => {
                assert_eq!(plan, CheckPlan::Proceed);
                assert!(evict);
            }
            CacheEntryStatus::LiveCompleted => {
                assert_eq!(plan, CheckPlan::Completed);
                assert!(!evict);
            }
            CacheEntryStatus::ExpiredCompleted => {
                assert_eq!(plan, CheckPlan::Proceed);
                assert!(evict);
            }
        }
    }
}

// ── Key generation ────────────────────────────────────────────────────────────

/// Derive an idempotency key from `tool_name` and `arguments`.
///
/// The key is the hex-encoded SHA-256 digest of
/// `"{tool_name}\0{canonical_json(arguments)}"`.
/// Using a NUL separator prevents collisions between tool names that share a
/// common prefix and arguments.
///
/// The resulting key is stable: identical `(tool_name, arguments)` pairs
/// always produce the same key regardless of JSON key ordering.
#[must_use]
pub fn derive_key(tool_name: &str, arguments: &Value) -> String {
    let canonical = canonical_json(arguments);
    sha256_hex_chunks([tool_name.as_bytes(), &b"\0"[..], canonical.as_bytes()])
}

// ── Idempotency enforcement ───────────────────────────────────────────────────

/// Outcome of the idempotency guard.
#[derive(Debug)]
pub enum GuardOutcome {
    /// Proceed with execution; the key has been registered as in-flight.
    Proceed,
    /// Return the cached result — no execution needed.
    CachedResult(Value),
}

/// Check the idempotency cache and either return a cached result or register
/// the key as in-flight for execution.
///
/// # Errors
///
/// Returns [`Error::DuplicateRequest`] (HTTP 409 equivalent) when an identical
/// request is already in flight.
pub fn enforce(cache: &IdempotencyCache, key: &str) -> Result<GuardOutcome> {
    match cache.check(key) {
        CheckOutcome::Proceed => {
            cache.mark_in_flight(key);
            Ok(GuardOutcome::Proceed)
        }
        CheckOutcome::InFlight => Err(Error::json_rpc(
            409,
            format!("Duplicate request in progress for key: {key}"),
        )),
        CheckOutcome::Completed(value) => Ok(GuardOutcome::CachedResult(value)),
    }
}

/// Spawn a background tokio task that periodically evicts stale idempotency
/// entries from `cache`.
///
/// The task runs every `interval` and stops when the `Arc` reference count
/// drops to 1 (i.e., all other owners have dropped their handles).
pub fn spawn_cleanup_task(cache: Arc<IdempotencyCache>, interval: Duration) {
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(interval);
        loop {
            ticker.tick().await;
            // Stop if we are the sole Arc holder (server is shutting down).
            if Arc::strong_count(&cache) <= 1 {
                break;
            }
            cache.evict_expired();
        }
    });
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use std::thread;

    // ── derive_key ────────────────────────────────────────────────────────────

    #[test]
    fn derive_key_is_deterministic_for_same_inputs() {
        // GIVEN: identical tool name and arguments
        // WHEN: deriving the key twice
        // THEN: both keys are identical
        let k1 = derive_key("gmail_send_email", &json!({"to": "a@b.com", "body": "hi"}));
        let k2 = derive_key("gmail_send_email", &json!({"to": "a@b.com", "body": "hi"}));
        assert_eq!(k1, k2);
    }

    #[test]
    fn derive_key_is_64_hex_chars() {
        // GIVEN: any tool + arguments
        // WHEN: deriving the key
        // THEN: result is a 64-character hex string (SHA-256)
        let key = derive_key("my_tool", &json!({}));
        assert_eq!(key.len(), 64);
        assert!(key.chars().all(|c| c.is_ascii_hexdigit()));
    }

    #[test]
    fn derive_key_differs_for_different_tool_names() {
        // GIVEN: same arguments but different tool names
        // WHEN: deriving keys
        // THEN: keys are different
        let k1 = derive_key("tool_a", &json!({"x": 1}));
        let k2 = derive_key("tool_b", &json!({"x": 1}));
        assert_ne!(k1, k2);
    }

    #[test]
    fn derive_key_differs_for_different_arguments() {
        // GIVEN: same tool name but different arguments
        // WHEN: deriving keys
        // THEN: keys are different
        let k1 = derive_key("send", &json!({"to": "a@b.com"}));
        let k2 = derive_key("send", &json!({"to": "c@d.com"}));
        assert_ne!(k1, k2);
    }

    #[test]
    fn derive_key_prevents_prefix_collision() {
        // GIVEN: a tool whose name is a prefix of another (tool, tool_extended)
        // WHEN: deriving keys with the same suffix args
        // THEN: keys are different (NUL separator prevents collision)
        let k1 = derive_key("tool", &json!({"a": "extended"}));
        let k2 = derive_key("tool_extended", &json!({"a": ""}));
        assert_ne!(k1, k2);
    }

    // ── IdempotencyState ──────────────────────────────────────────────────────

    #[test]
    fn state_in_flight_is_not_expired_immediately() {
        // GIVEN: a freshly created InFlight state
        // WHEN: checking expiry
        // THEN: not expired
        let state = IdempotencyState::InFlight(Instant::now());
        assert!(!state.is_expired());
    }

    #[test]
    fn state_completed_is_not_expired_immediately() {
        // GIVEN: a freshly created Completed state
        // WHEN: checking expiry
        // THEN: not expired
        let state = IdempotencyState::Completed(json!({"ok": true}), Instant::now());
        assert!(!state.is_expired());
    }

    #[test]
    fn state_in_flight_is_live_immediately() {
        // GIVEN: a freshly created InFlight state
        // WHEN: checking is_in_flight
        // THEN: true
        let state = IdempotencyState::InFlight(Instant::now());
        assert!(state.is_in_flight());
    }

    #[test]
    fn state_completed_is_not_in_flight() {
        // GIVEN: a Completed state
        // WHEN: checking is_in_flight
        // THEN: false
        let state = IdempotencyState::Completed(json!(null), Instant::now());
        assert!(!state.is_in_flight());
    }

    // ── IdempotencyCache::check ───────────────────────────────────────────────

    #[test]
    fn check_returns_proceed_for_unknown_key() {
        // GIVEN: an empty cache
        // WHEN: checking an unknown key
        // THEN: Proceed
        let cache = IdempotencyCache::new();
        assert!(matches!(cache.check("unknown"), CheckOutcome::Proceed));
    }

    #[test]
    fn check_returns_in_flight_for_live_in_flight_key() {
        // GIVEN: cache with a live in-flight key
        // WHEN: checking the same key
        // THEN: InFlight
        let cache = IdempotencyCache::new();
        cache.mark_in_flight("key-1");
        assert!(matches!(cache.check("key-1"), CheckOutcome::InFlight));
    }

    #[test]
    fn check_returns_completed_result_for_completed_key() {
        // GIVEN: cache with a completed key
        // WHEN: checking the same key
        // THEN: Completed with the stored value
        let cache = IdempotencyCache::new();
        let result = json!({"issue_id": "LIN-42"});
        cache.mark_in_flight("key-2");
        cache.mark_completed("key-2", result.clone());
        match cache.check("key-2") {
            CheckOutcome::Completed(v) => assert_eq!(v, result),
            other => panic!("expected Completed, got {other:?}"),
        }
    }

    #[test]
    fn check_evicts_stale_in_flight_and_returns_proceed() {
        // GIVEN: an in-flight entry whose timestamp is older than IN_FLIGHT_TIMEOUT
        // WHEN: checking the key
        // THEN: Proceed (stale entry evicted)
        let cache = IdempotencyCache::new();
        // Insert an entry with a timestamp in the distant past
        cache.entries.insert(
            "stale".to_string(),
            IdempotencyState::InFlight(
                Instant::now()
                    .checked_sub(IN_FLIGHT_TIMEOUT)
                    .unwrap()
                    .checked_sub(Duration::from_secs(1))
                    .unwrap(),
            ),
        );
        assert!(matches!(cache.check("stale"), CheckOutcome::Proceed));
        assert_eq!(cache.len(), 0, "stale entry must be removed");
    }

    #[test]
    fn check_evicts_expired_completed_and_returns_proceed() {
        // GIVEN: a completed entry whose TTL has elapsed
        // WHEN: checking the key
        // THEN: Proceed (expired entry evicted)
        let cache = IdempotencyCache::new();
        cache.entries.insert(
            "old".to_string(),
            IdempotencyState::Completed(
                json!(null),
                Instant::now()
                    .checked_sub(COMPLETED_TTL)
                    .unwrap()
                    .checked_sub(Duration::from_secs(1))
                    .unwrap(),
            ),
        );
        assert!(matches!(cache.check("old"), CheckOutcome::Proceed));
        assert_eq!(cache.len(), 0, "expired entry must be removed");
    }

    // ── evict_expired ─────────────────────────────────────────────────────────

    #[test]
    fn evict_expired_removes_only_stale_entries() {
        // GIVEN: one fresh and one stale completed entry
        // WHEN: calling evict_expired
        // THEN: only the stale entry is removed
        let cache = IdempotencyCache::new();
        cache.mark_in_flight("fresh");
        cache.mark_completed("fresh", json!(1));
        cache.entries.insert(
            "stale".to_string(),
            IdempotencyState::Completed(
                json!(2),
                Instant::now()
                    .checked_sub(COMPLETED_TTL)
                    .unwrap()
                    .checked_sub(Duration::from_secs(1))
                    .unwrap(),
            ),
        );

        cache.evict_expired();

        assert_eq!(cache.len(), 1);
        assert!(matches!(cache.check("fresh"), CheckOutcome::Completed(_)));
    }

    // ── enforce ───────────────────────────────────────────────────────────────

    #[test]
    fn enforce_marks_in_flight_and_returns_proceed_for_new_key() {
        // GIVEN: an empty cache
        // WHEN: enforcing on a new key
        // THEN: Proceed, and the key is now in-flight
        let cache = IdempotencyCache::new();
        let outcome = enforce(&cache, "k1").expect("should not fail");
        assert!(matches!(outcome, GuardOutcome::Proceed));
        assert!(matches!(cache.check("k1"), CheckOutcome::InFlight));
    }

    #[test]
    fn enforce_returns_cached_result_for_completed_key() {
        // GIVEN: a completed key in cache
        // WHEN: enforcing on that key
        // THEN: CachedResult with the stored value
        let cache = IdempotencyCache::new();
        let expected = json!({"done": true});
        cache.mark_in_flight("k2");
        cache.mark_completed("k2", expected.clone());
        match enforce(&cache, "k2").expect("should not fail") {
            GuardOutcome::CachedResult(v) => assert_eq!(v, expected),
            GuardOutcome::Proceed => panic!("expected CachedResult"),
        }
    }

    #[test]
    fn enforce_returns_error_for_in_flight_key() {
        // GIVEN: a live in-flight key
        // WHEN: enforcing on the same key from a concurrent caller
        // THEN: Err with code 409
        let cache = IdempotencyCache::new();
        cache.mark_in_flight("k3");
        let err = enforce(&cache, "k3").expect_err("should return 409");
        match err {
            crate::Error::JsonRpc { code, .. } => assert_eq!(code, 409),
            _ => panic!("expected JsonRpc error"),
        }
    }

    // ── remove ────────────────────────────────────────────────────────────────

    #[test]
    fn remove_clears_key_making_it_retryable() {
        // GIVEN: an in-flight key
        // WHEN: calling remove (e.g. on tool failure)
        // THEN: key is gone and check returns Proceed
        let cache = IdempotencyCache::new();
        cache.mark_in_flight("fail-key");
        cache.remove("fail-key");
        assert!(matches!(cache.check("fail-key"), CheckOutcome::Proceed));
        assert_eq!(cache.len(), 0);
    }

    // ── concurrent access ─────────────────────────────────────────────────────

    #[test]
    fn concurrent_mark_completed_is_safe() {
        // GIVEN: cache shared across threads
        // WHEN: 10 threads each mark different keys completed
        // THEN: all entries are present without data races
        let cache = Arc::new(IdempotencyCache::new());
        let handles: Vec<_> = (0..10)
            .map(|i| {
                let c = Arc::clone(&cache);
                thread::spawn(move || {
                    let key = format!("key-{i}");
                    c.mark_in_flight(&key);
                    c.mark_completed(&key, json!(i));
                })
            })
            .collect();

        for h in handles {
            h.join().expect("thread panicked");
        }

        assert_eq!(cache.len(), 10);
    }

    // ── is_empty / len ────────────────────────────────────────────────────────

    #[test]
    fn new_cache_is_empty() {
        let cache = IdempotencyCache::new();
        assert!(cache.is_empty());
        assert_eq!(cache.len(), 0);
    }

    #[test]
    fn len_increases_on_insert() {
        let cache = IdempotencyCache::new();
        cache.mark_in_flight("a");
        cache.mark_in_flight("b");
        assert_eq!(cache.len(), 2);
        assert!(!cache.is_empty());
    }

    // ── cleanup task (tokio) ──────────────────────────────────────────────────

    #[tokio::test]
    async fn spawn_cleanup_task_evicts_expired_entries() {
        // GIVEN: a cache with one stale completed entry
        // WHEN: the cleanup task runs
        // THEN: the entry is evicted
        let cache = Arc::new(IdempotencyCache::new());
        cache.entries.insert(
            "stale".to_string(),
            IdempotencyState::Completed(
                json!(null),
                Instant::now()
                    .checked_sub(COMPLETED_TTL)
                    .unwrap()
                    .checked_sub(Duration::from_secs(1))
                    .unwrap(),
            ),
        );

        spawn_cleanup_task(Arc::clone(&cache), Duration::from_millis(10));

        // Wait a bit for the task to run
        tokio::time::sleep(Duration::from_millis(50)).await;

        assert_eq!(cache.len(), 0, "stale entry should have been evicted");
    }
}
