//! Predictive tool prefetch via invocation sequence tracking.
//!
//! Tracks which tools are called after which within a session, learns
//! transition probabilities, and predicts likely next tools to invoke.
//!
//! # Design
//!
//! - `TransitionTracker` owns all state (thread-safe via `DashMap` + `parking_lot`).
//! - `record_transition(session_id, tool)` updates per-session "last tool" and
//!   increments the transition counter for `last → current`.
//! - `predict_next(tool, min_confidence, min_count)` returns candidates whose
//!   observed frequency clears both thresholds, sorted by descending confidence.
//! - Save/load follows the same pattern as [`crate::ranking::SearchRanker`].
//!
//! # Thresholds
//!
//! Predictions are only emitted when a candidate satisfies **both**:
//! - `count ≥ min_count` — prevents noise from single observations
//! - `confidence ≥ min_confidence` — expressed as a fraction (0.0–1.0)

use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};

use dashmap::DashMap;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};

// ============================================================================
// Public types
// ============================================================================

/// A predicted next tool with its confidence score.
#[derive(Debug, Clone, PartialEq)]
pub struct Prediction {
    /// Fully-qualified tool identifier (`"server:tool"`)
    pub tool: String,
    /// Confidence in the range `[0.0, 1.0]`
    pub confidence: f64,
}

// ============================================================================
// TransitionTracker
// ============================================================================

/// Thread-safe tracker for tool invocation sequence learning.
///
/// Internally uses two `DashMap` layers:
/// - outer key: `"from_tool"` (the tool just invoked)
/// - inner key: `"to_tool"` (the tool invoked next)
/// - value: `AtomicU64` transition count
///
/// Per-session "last invoked tool" is tracked in a separate `DashMap`
/// with `Mutex<Option<String>>` values so sessions are independent.
pub struct TransitionTracker {
    /// `from_tool -> (to_tool -> count)`
    transitions: DashMap<String, DashMap<String, AtomicU64>>,
    /// `session_id -> last_invoked_tool`
    last_per_session: DashMap<String, Mutex<Option<String>>>,
}

impl TransitionTracker {
    /// Create a new, empty tracker.
    #[must_use]
    pub fn new() -> Self {
        Self {
            transitions: DashMap::new(),
            last_per_session: DashMap::new(),
        }
    }

    /// Record a tool invocation for a session.
    ///
    /// If the session has a previous tool, increments the `previous → tool`
    /// transition counter. Always updates the session's "last tool".
    ///
    /// # Arguments
    /// * `session_id` — opaque session identifier
    /// * `tool` — fully-qualified tool key, e.g. `"server:tool_name"`
    pub fn record_transition(&self, session_id: &str, tool: &str) {
        // Get-or-create the session slot, then swap the last-tool under the lock.
        let previous = {
            let entry = self
                .last_per_session
                .entry(session_id.to_string())
                .or_insert_with(|| Mutex::new(None));
            let mut guard = entry.value().lock();
            guard.replace(tool.to_string())
        };

        // Record the transition if there was a previous tool.
        if let Some(from) = previous {
            let inner = self.transitions.entry(from).or_default();
            inner
                .entry(tool.to_string())
                .or_insert_with(|| AtomicU64::new(0))
                .fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Predict the most likely next tools after `from_tool`.
    ///
    /// Returns candidates sorted by descending confidence where both
    /// `count ≥ min_count` and `confidence ≥ min_confidence`.
    ///
    /// Runs in **O(k)** where k is the number of distinct successors of
    /// `from_tool` — typically very small (< 20).
    ///
    /// # Arguments
    /// * `from_tool` — the tool just invoked
    /// * `min_confidence` — minimum fraction threshold (e.g. `0.30` for 30 %)
    /// * `min_count` — minimum absolute occurrence count (e.g. `3`)
    #[must_use]
    pub fn predict_next(
        &self,
        from_tool: &str,
        min_confidence: f64,
        min_count: u64,
    ) -> Vec<Prediction> {
        let Some(successors) = self.transitions.get(from_tool) else {
            return Vec::new();
        };

        let total: u64 = successors
            .iter()
            .map(|e| e.value().load(Ordering::Relaxed))
            .sum();

        if total == 0 {
            return Vec::new();
        }

        #[allow(clippy::cast_precision_loss)]
        let mut predictions: Vec<Prediction> = successors
            .iter()
            .filter_map(|entry| {
                let count = entry.value().load(Ordering::Relaxed);
                let confidence = count as f64 / total as f64;
                if count >= min_count && confidence >= min_confidence {
                    Some(Prediction {
                        tool: entry.key().clone(),
                        confidence,
                    })
                } else {
                    None
                }
            })
            .collect();

        // Deterministic descending sort by confidence, then tool name for ties.
        predictions.sort_by(|a, b| {
            b.confidence
                .partial_cmp(&a.confidence)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.tool.cmp(&b.tool))
        });

        predictions
    }

    /// Total number of recorded transitions (sum across all from-tools).
    #[must_use]
    pub fn total_transitions(&self) -> u64 {
        self.transitions
            .iter()
            .flat_map(|outer| {
                outer
                    .value()
                    .iter()
                    .map(|inner| inner.value().load(Ordering::Relaxed))
                    .collect::<Vec<_>>()
            })
            .sum()
    }

    /// Save transition data to a JSON file.
    ///
    /// # Errors
    ///
    /// Returns an error if serialization fails or the file cannot be written.
    pub fn save(&self, path: &Path) -> std::io::Result<()> {
        let entries: Vec<TransitionEntry> = self
            .transitions
            .iter()
            .flat_map(|outer| {
                let from = outer.key().clone();
                outer
                    .value()
                    .iter()
                    .map(move |inner| TransitionEntry {
                        from: from.clone(),
                        to: inner.key().clone(),
                        count: inner.value().load(Ordering::Relaxed),
                    })
                    .collect::<Vec<_>>()
            })
            .collect();

        let json = serde_json::to_string_pretty(&entries)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        std::fs::write(path, json)
    }

    /// Load transition data from a JSON file.
    ///
    /// Merges the loaded data with any existing in-memory state — useful for
    /// hot-reloading without losing in-flight session data.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read or JSON is invalid.
    pub fn load(&self, path: &Path) -> std::io::Result<()> {
        let content = std::fs::read_to_string(path)?;
        let entries: Vec<TransitionEntry> = serde_json::from_str(&content)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;

        for entry in entries {
            let inner = self.transitions.entry(entry.from).or_default();
            inner
                .entry(entry.to)
                .or_insert_with(|| AtomicU64::new(0))
                .fetch_add(entry.count, Ordering::Relaxed);
        }

        Ok(())
    }
}

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

// ============================================================================
// Serialization helpers (private)
// ============================================================================

#[derive(Debug, Serialize, Deserialize)]
struct TransitionEntry {
    from: String,
    to: String,
    count: u64,
}

// ============================================================================
// Tests
// ============================================================================

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

    // ── record_transition ────────────────────────────────────────────────

    #[test]
    fn record_transition_first_call_has_no_predecessor() {
        // GIVEN: a fresh tracker and a new session
        // WHEN: recording the first tool invocation
        // THEN: no transitions are recorded (nothing to chain from)
        let tracker = TransitionTracker::new();
        tracker.record_transition("session-1", "s1:tool_a");

        assert_eq!(tracker.total_transitions(), 0);
    }

    #[test]
    fn record_transition_second_call_creates_one_transition() {
        // GIVEN: session with one prior invocation
        // WHEN: a second tool is invoked
        // THEN: exactly one A→B transition is recorded
        let tracker = TransitionTracker::new();
        tracker.record_transition("s", "s1:tool_a");
        tracker.record_transition("s", "s1:tool_b");

        assert_eq!(tracker.total_transitions(), 1);
        let preds = tracker.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds.len(), 1);
        assert_eq!(preds[0].tool, "s1:tool_b");
        assert!((preds[0].confidence - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn record_transition_repeated_pair_increments_count() {
        // GIVEN: the same A→B pair observed three times
        // WHEN: predicting next after A
        // THEN: count is 3 and confidence is 1.0
        let tracker = TransitionTracker::new();
        for _ in 0..3 {
            tracker.record_transition("s", "s1:tool_a");
            tracker.record_transition("s", "s1:tool_b");
        }

        let preds = tracker.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds.len(), 1);
        assert_eq!(preds[0].tool, "s1:tool_b");
        // 3/3 = 1.0 — but the session ends on "s1:tool_b" so only 3 A→B are recorded
        assert!((preds[0].confidence - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn record_transition_multiple_successors_tracked_independently() {
        // GIVEN: A→B twice, A→C once (interleaved via session resets)
        // WHEN: predicting next after A
        // THEN: B confidence = 2/3, C confidence = 1/3
        let tracker = TransitionTracker::new();

        // Session 1: A→B
        tracker.record_transition("s1", "s1:tool_a");
        tracker.record_transition("s1", "s1:tool_b");

        // Session 2: A→B
        tracker.record_transition("s2", "s1:tool_a");
        tracker.record_transition("s2", "s1:tool_b");

        // Session 3: A→C
        tracker.record_transition("s3", "s1:tool_a");
        tracker.record_transition("s3", "s1:tool_c");

        let preds = tracker.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds.len(), 2);
        assert_eq!(preds[0].tool, "s1:tool_b");
        assert!((preds[0].confidence - 2.0 / 3.0).abs() < 0.001);
        assert_eq!(preds[1].tool, "s1:tool_c");
        assert!((preds[1].confidence - 1.0 / 3.0).abs() < 0.001);
    }

    #[test]
    fn record_transition_different_sessions_are_independent() {
        // GIVEN: two sessions starting with different tools
        // WHEN: both sessions call tool_b second
        // THEN: transitions are attributed to each session's own "last tool"
        let tracker = TransitionTracker::new();

        tracker.record_transition("s1", "s1:tool_a");
        tracker.record_transition("s2", "s1:tool_x");
        tracker.record_transition("s1", "s1:tool_b"); // A→B
        tracker.record_transition("s2", "s1:tool_b"); // X→B

        let preds_a = tracker.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds_a.len(), 1);
        assert_eq!(preds_a[0].tool, "s1:tool_b");

        let preds_x = tracker.predict_next("s1:tool_x", 0.0, 1);
        assert_eq!(preds_x.len(), 1);
        assert_eq!(preds_x[0].tool, "s1:tool_b");
    }

    // ── predict_next ─────────────────────────────────────────────────────

    #[test]
    fn predict_next_returns_empty_for_unknown_tool() {
        // GIVEN: tracker with no data
        // WHEN: predicting next for an unknown tool
        // THEN: empty result (no panic)
        let tracker = TransitionTracker::new();
        let preds = tracker.predict_next("nonexistent:tool", 0.30, 3);
        assert!(preds.is_empty());
    }

    #[test]
    fn predict_next_filters_by_min_count() {
        // GIVEN: A→B seen once (below min_count = 3)
        // WHEN: predicting with min_count = 3
        // THEN: no predictions returned
        let tracker = TransitionTracker::new();
        tracker.record_transition("s", "s1:tool_a");
        tracker.record_transition("s", "s1:tool_b");

        let preds = tracker.predict_next("s1:tool_a", 0.0, 3);
        assert!(preds.is_empty());
    }

    #[test]
    fn predict_next_filters_by_min_confidence() {
        // GIVEN: A→B once, A→C four times (total 5), B confidence = 0.20
        // WHEN: predicting with min_confidence = 0.30
        // THEN: only C appears (confidence 0.80)
        let tracker = TransitionTracker::new();

        tracker.record_transition("s1", "s1:tool_a");
        tracker.record_transition("s1", "s1:tool_b"); // A→B ×1

        for _ in 0..4 {
            tracker.record_transition("s2", "s1:tool_a");
            tracker.record_transition("s2", "s1:tool_c"); // A→C ×4
        }

        let preds = tracker.predict_next("s1:tool_a", 0.30, 1);
        assert_eq!(preds.len(), 1);
        assert_eq!(preds[0].tool, "s1:tool_c");
        assert!((preds[0].confidence - 0.80).abs() < 0.001);
    }

    #[test]
    fn predict_next_satisfies_both_thresholds() {
        // GIVEN: default thresholds 30% confidence + 3 min occurrences
        // A→B: 8 times out of 10 total → confidence 0.80, count 8 → PASSES
        // A→C: 2 times out of 10 total → confidence 0.20, count 2 → FAILS BOTH
        // WHEN: predicting with defaults
        // THEN: only B is returned
        let tracker = TransitionTracker::new();

        for _ in 0..8 {
            tracker.record_transition("sx", "s1:tool_a");
            tracker.record_transition("sx", "s1:tool_b");
        }
        for _ in 0..2 {
            tracker.record_transition("sy", "s1:tool_a");
            tracker.record_transition("sy", "s1:tool_c");
        }

        let preds = tracker.predict_next("s1:tool_a", 0.30, 3);
        assert_eq!(preds.len(), 1);
        assert_eq!(preds[0].tool, "s1:tool_b");
    }

    #[test]
    fn predict_next_is_sorted_by_descending_confidence() {
        // GIVEN: A→B 6×, A→C 3×, A→D 1× (all clear min_count=1, min_confidence=0)
        // WHEN: predicting
        // THEN: order is B (0.60), C (0.30), D (0.10)
        let tracker = TransitionTracker::new();

        for _ in 0..6 {
            tracker.record_transition("sx", "s1:tool_a");
            tracker.record_transition("sx", "s1:tool_b");
        }
        for _ in 0..3 {
            tracker.record_transition("sy", "s1:tool_a");
            tracker.record_transition("sy", "s1:tool_c");
        }
        tracker.record_transition("sz", "s1:tool_a");
        tracker.record_transition("sz", "s1:tool_d");

        let preds = tracker.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds.len(), 3);
        assert!(preds[0].confidence >= preds[1].confidence);
        assert!(preds[1].confidence >= preds[2].confidence);
        assert_eq!(preds[0].tool, "s1:tool_b");
    }

    #[test]
    fn predict_next_tie_broken_by_tool_name_alphabetically() {
        // GIVEN: A→B 5×, A→C 5× (equal confidence)
        // WHEN: predicting
        // THEN: B comes before C (alphabetical tiebreak)
        let tracker = TransitionTracker::new();

        for _ in 0..5 {
            tracker.record_transition("s1", "s1:tool_a");
            tracker.record_transition("s1", "s1:tool_b");
        }
        for _ in 0..5 {
            tracker.record_transition("s2", "s1:tool_a");
            tracker.record_transition("s2", "s1:tool_c");
        }

        let preds = tracker.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds.len(), 2);
        assert_eq!(preds[0].tool, "s1:tool_b");
        assert_eq!(preds[1].tool, "s1:tool_c");
    }

    // ── cold start / edge cases ──────────────────────────────────────────

    #[test]
    fn predict_next_cold_start_no_data_returns_empty() {
        // GIVEN: brand-new tracker
        // WHEN: predicting for any tool
        // THEN: always empty, no panic
        let tracker = TransitionTracker::new();
        assert!(tracker.predict_next("any:tool", 0.30, 3).is_empty());
    }

    #[test]
    fn record_transition_single_session_single_call_no_transitions() {
        // GIVEN: only one call in a session (no chain possible)
        // THEN: transitions = 0 and predictions are empty
        let tracker = TransitionTracker::new();
        tracker.record_transition("session-x", "s1:lonely_tool");

        assert_eq!(tracker.total_transitions(), 0);
        assert!(tracker.predict_next("s1:lonely_tool", 0.0, 1).is_empty());
    }

    // ── persistence ──────────────────────────────────────────────────────

    #[test]
    fn save_and_load_round_trips_transition_counts() {
        // GIVEN: tracker with A→B×5, A→C×2
        let tracker = TransitionTracker::new();
        for _ in 0..5 {
            tracker.record_transition("s1", "s1:tool_a");
            tracker.record_transition("s1", "s1:tool_b");
        }
        for _ in 0..2 {
            tracker.record_transition("s2", "s1:tool_a");
            tracker.record_transition("s2", "s1:tool_c");
        }

        let file = NamedTempFile::new().unwrap();
        tracker.save(file.path()).unwrap();

        // WHEN: loading into a fresh tracker
        let loaded = TransitionTracker::new();
        loaded.load(file.path()).unwrap();

        // THEN: predictions match original
        let preds = loaded.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds.len(), 2);
        assert_eq!(preds[0].tool, "s1:tool_b");
        assert!((preds[0].confidence - 5.0 / 7.0).abs() < 0.001);
        assert_eq!(preds[1].tool, "s1:tool_c");
        assert!((preds[1].confidence - 2.0 / 7.0).abs() < 0.001);
    }

    #[test]
    fn load_merges_with_existing_data() {
        // GIVEN: in-memory A→B×3 and file with A→B×2 (total should be 5)
        let tracker = TransitionTracker::new();
        for _ in 0..3 {
            tracker.record_transition("s", "s1:tool_a");
            tracker.record_transition("s", "s1:tool_b");
        }

        // Save the file-based data separately
        let file_tracker = TransitionTracker::new();
        for _ in 0..2 {
            file_tracker.record_transition("s", "s1:tool_a");
            file_tracker.record_transition("s", "s1:tool_b");
        }
        let file = NamedTempFile::new().unwrap();
        file_tracker.save(file.path()).unwrap();

        // Load into original tracker (merge)
        tracker.load(file.path()).unwrap();

        // WHEN: predicting after merge
        let preds = tracker.predict_next("s1:tool_a", 0.0, 1);
        assert_eq!(preds.len(), 1);
        // 3 in-memory + 2 from file = 5, confidence 5/5 = 1.0
        assert!((preds[0].confidence - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn save_empty_tracker_produces_valid_json() {
        // GIVEN: tracker with no data
        // WHEN: saving
        // THEN: writes valid (empty array) JSON
        let tracker = TransitionTracker::new();
        let file = NamedTempFile::new().unwrap();
        tracker.save(file.path()).unwrap();

        let content = std::fs::read_to_string(file.path()).unwrap();
        let parsed: serde_json::Value = serde_json::from_str(&content).unwrap();
        assert!(parsed.is_array());
        assert_eq!(parsed.as_array().unwrap().len(), 0);
    }

    #[test]
    fn load_empty_file_is_a_no_op() {
        // GIVEN: empty JSON array on disk
        let file = NamedTempFile::new().unwrap();
        std::fs::write(file.path(), "[]").unwrap();

        let tracker = TransitionTracker::new();
        tracker.load(file.path()).unwrap();

        assert_eq!(tracker.total_transitions(), 0);
    }

    #[test]
    fn total_transitions_counts_all_edges() {
        // GIVEN: three sessions with interleaved tool calls.
        //
        // Session "sx" (2 iterations of a→b):
        //   iter1: record(a) → last=a; record(b) → a→b×1, last=b
        //   iter2: record(a) → b→a×1, last=a; record(b) → a→b×2, last=b
        //   edges from sx: a→b×2, b→a×1  → 3 counted edges
        //
        // Session "sy":
        //   record(a) → last=a; record(c) → a→c×1, last=c
        //   edges from sy: a→c×1            → 1 counted edge
        //
        // Session "sz" (3 iterations of x→y):
        //   iter1: record(x) → last=x; record(y) → x→y×1, last=y
        //   iter2: record(x) → y→x×1, last=x; record(y) → x→y×2, last=y
        //   iter3: record(x) → y→x×2, last=x; record(y) → x→y×3, last=y
        //   edges from sz: x→y×3, y→x×2   → 5 counted edges
        //
        // Total: 3 + 1 + 5 = 9
        let tracker = TransitionTracker::new();
        for _ in 0..2 {
            tracker.record_transition("sx", "t:a");
            tracker.record_transition("sx", "t:b");
        }
        tracker.record_transition("sy", "t:a");
        tracker.record_transition("sy", "t:c");
        for _ in 0..3 {
            tracker.record_transition("sz", "t:x");
            tracker.record_transition("sz", "t:y");
        }

        assert_eq!(tracker.total_transitions(), 9);
    }
}
