pub mod cache;

use std::{
    collections::{HashMap, VecDeque},
    time::{Instant, SystemTime, UNIX_EPOCH},
};

use crate::{
    providers::{EnvSnapshot, GitSnapshot},
    types::{
        AgentSummary, CompletedToolCount, Line3Metrics, PendingTask, TaskItem, TodoInProgressItem,
        TodoSummary, ToolSummary,
    },
};
use cache::{CacheEntry, SessionCache, CACHE_TTL_MS};

const MAX_RECENT_TOOLS_CAP: usize = 10;

/// Maximum number of CTX% samples retained for the ledger sparkline. The
/// widget draws 12 (6 braille cells × 2 samples each); the surplus lets
/// the adaptive 1-minute window expand when samples arrive in bursts.
pub const MAX_CTX_HISTORY: usize = 30;

#[derive(Debug, Clone, Default)]
pub struct SessionState {
    pub last_transcript_offset: u64,
    pub last_transcript_path: Option<String>,
    pub last_transcript_poll: Option<Instant>,
    pub active_tools: Vec<ToolSummary>,
    pub recent_tools: Vec<ToolSummary>,
    pub active_agents: Vec<AgentSummary>,
    pub completed_agents: Vec<AgentSummary>,
    pub completed_tool_counts: HashMap<String, (u32, u64)>,
    pub todo: Option<TodoSummary>,
    // Agent linking: Agent tool_use → agent_progress ID linking
    pub pending_tasks: Vec<PendingTask>,
    pub task_agent_links: HashMap<String, String>, // tool_use_id → agentId
    // Todo tracking: TaskCreate/TaskUpdate stateful accumulation
    pub task_items: HashMap<String, TaskItem>,
    pub task_counter: u32,
    pub cached_env: Option<(String, EnvSnapshot)>,
    pub cached_git: Option<(String, GitSnapshot)>,
    pub cached_line3: Option<Line3Metrics>,
    // Token speed tracking (output only)
    pub last_output_tokens: Option<u64>,
    pub last_output_token_time_ms: Option<u64>,
    pub output_speed_toks_per_sec: Option<f64>,
    /// Sparkline source: rolling window of `(pct, epoch_ms)` samples.
    /// Timestamps drive the ledger's adaptive 1-minute window + velocity-
    /// based aurora coloring.
    pub ctx_history: VecDeque<(u8, u64)>,
}

impl SessionState {
    pub fn reset_transcript_if_path_changed(&mut self, transcript_path: &str) {
        if self.last_transcript_path.as_deref() != Some(transcript_path) {
            self.last_transcript_path = Some(transcript_path.to_string());
            self.last_transcript_offset = 0;
            self.last_transcript_poll = None;
            self.active_tools.clear();
            self.recent_tools.clear();
            self.active_agents.clear();
            self.completed_agents.clear();
            self.completed_tool_counts.clear();
            self.todo = None;
            self.pending_tasks.clear();
            self.task_agent_links.clear();
            self.task_items.clear();
            self.task_counter = 0;
            self.cached_line3 = None;
            self.last_output_tokens = None;
            self.last_output_token_time_ms = None;
            self.output_speed_toks_per_sec = None;
            self.ctx_history.clear();
        }
    }

    /// Push a fresh CTX% sample stamped with `ts_ms`, discarding the oldest
    /// when at capacity.
    pub fn push_ctx_sample(&mut self, pct: u8, ts_ms: u64) {
        if self.ctx_history.len() == MAX_CTX_HISTORY {
            self.ctx_history.pop_front();
        }
        self.ctx_history.push_back((pct.min(100), ts_ms));
    }

    /// Compute output token speed from successive payload snapshots.
    /// Returns the current speed estimate, or keeps the last known value if the
    /// delta window exceeds 2 seconds (likely idle/paused).
    /// When `current_tokens` is `None`, returns the last known speed without
    /// corrupting the time anchor (bug #2 fix).
    pub fn update_output_speed(&mut self, current_tokens: Option<u64>) -> Option<f64> {
        let curr = match current_tokens {
            Some(c) => c,
            None => return self.output_speed_toks_per_sec, // early-return: don't touch state
        };
        let now_ms = cache::now_epoch_ms();
        let result = match (self.last_output_tokens, self.last_output_token_time_ms) {
            (Some(prev), Some(prev_time)) => {
                let delta_ms = now_ms.saturating_sub(prev_time);
                let delta_toks = curr.saturating_sub(prev);
                if delta_ms > 0 && delta_ms <= 2000 && delta_toks > 0 {
                    Some(delta_toks as f64 / delta_ms as f64 * 1000.0)
                } else {
                    self.output_speed_toks_per_sec // keep last known
                }
            }
            _ => None,
        };
        self.last_output_tokens = Some(curr);
        self.last_output_token_time_ms = Some(now_ms);
        self.output_speed_toks_per_sec = result;
        result
    }

    pub fn cached_env_for(&self, cwd: &str) -> Option<EnvSnapshot> {
        self.cached_env.as_ref().and_then(|(path, snapshot)| {
            if path == cwd {
                Some(snapshot.clone())
            } else {
                None
            }
        })
    }

    pub fn set_cached_env(&mut self, cwd: String, snapshot: EnvSnapshot) {
        self.cached_env = Some((cwd, snapshot));
    }

    pub fn cached_git_for(&self, cwd: &str) -> Option<GitSnapshot> {
        self.cached_git.as_ref().and_then(|(path, snapshot)| {
            if path == cwd {
                Some(snapshot.clone())
            } else {
                None
            }
        })
    }

    pub fn set_cached_git(&mut self, cwd: String, snapshot: GitSnapshot) {
        self.cached_git = Some((cwd, snapshot));
    }

    pub fn upsert_tool(&mut self, id: String, name: String, target: Option<String>) {
        if let Some(position) = self.active_tools.iter().position(|tool| tool.id == id) {
            self.active_tools.remove(position);
        }
        let tool = ToolSummary { id, name, target };
        self.active_tools.push(tool.clone());

        // recent_tools: display list (dedup, push, cap)
        self.recent_tools.retain(|t| t.id != tool.id);
        self.recent_tools.push(tool);
        if self.recent_tools.len() > MAX_RECENT_TOOLS_CAP {
            self.recent_tools
                .drain(..self.recent_tools.len() - MAX_RECENT_TOOLS_CAP);
        }
    }

    pub fn remove_tool(&mut self, id: &str, event_ts: Option<u64>) {
        if let Some(tool) = self.active_tools.iter().find(|t| t.id == id) {
            self.record_tool_completion(&tool.name.clone(), event_ts);
        }
        self.active_tools.retain(|tool| tool.id != id);
        // recent_tools: intentionally NOT touched — tool stays visible until displaced
    }

    pub fn record_tool_completion(&mut self, name: &str, event_ts: Option<u64>) {
        let ts = event_ts.unwrap_or_else(cache::now_epoch_ms);
        let entry = self
            .completed_tool_counts
            .entry(name.to_string())
            .or_insert((0, 0));
        entry.0 += 1;
        entry.1 = ts;
    }

    /// Hybrid scoring: count + recency bonus (decays over 2 minutes).
    /// Recently used tools float up even with low counts.
    pub fn scored_completed_tools(&self, max: usize) -> Vec<CompletedToolCount> {
        const RECENCY_WEIGHT: f64 = 3.0;
        const DECAY_WINDOW_MS: u64 = 120_000; // 2 minutes

        let now = cache::now_epoch_ms();
        let mut scored: Vec<(CompletedToolCount, f64)> = self
            .completed_tool_counts
            .iter()
            .map(|(name, (count, last_at))| {
                let age_ms = now.saturating_sub(*last_at);
                let recency = if age_ms < DECAY_WINDOW_MS {
                    RECENCY_WEIGHT * (1.0 - age_ms as f64 / DECAY_WINDOW_MS as f64)
                } else {
                    0.0
                };
                let score = *count as f64 + recency;
                (
                    CompletedToolCount {
                        name: name.clone(),
                        count: *count,
                        last_completed_at: Some(*last_at),
                    },
                    score,
                )
            })
            .collect();
        scored.sort_by(|a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then(a.0.name.cmp(&b.0.name))
        });
        scored.truncate(max);
        scored.into_iter().map(|(tool, _)| tool).collect()
    }

    pub fn upsert_agent(
        &mut self,
        id: String,
        description: String,
        agent_type: Option<String>,
        started_at: Option<u64>,
        model: Option<String>,
        message_id: Option<String>,
    ) {
        let (started_at, existing_model, existing_message_id) =
            if let Some(position) = self.active_agents.iter().position(|agent| agent.id == id) {
                let old = self.active_agents.remove(position);
                (old.started_at, old.model, old.message_id)
            } else {
                let ts = started_at.or_else(|| {
                    Some(
                        SystemTime::now()
                            .duration_since(UNIX_EPOCH)
                            .unwrap_or_default()
                            .as_millis() as u64,
                    )
                });
                (ts, None, None)
            };
        self.active_agents.push(AgentSummary {
            id,
            description,
            agent_type,
            started_at,
            model: model.or(existing_model),
            completed_at: None,
            message_id: message_id.or(existing_message_id),
        });
    }

    pub fn remove_agent(&mut self, id: &str) {
        if let Some(pos) = self.active_agents.iter().position(|a| a.id == id) {
            let mut agent = self.active_agents.remove(pos);
            agent.completed_at = Some(
                SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_millis() as u64,
            );
            self.completed_agents.push(agent);
            // Prune to max 10 completed agents
            if self.completed_agents.len() > 10 {
                let drain_count = self.completed_agents.len() - 10;
                self.completed_agents.drain(0..drain_count);
            }
        }
    }

    /// Drop an agent from the active list without recording it as
    /// completed. Used when an agent_progress event re-keys an existing
    /// active agent under a new id (e.g. progress assigns a runtime
    /// `agent_id` distinct from the original `tool_use_id`).
    pub fn discard_active_agent(&mut self, id: &str) {
        if let Some(pos) = self.active_agents.iter().position(|a| a.id == id) {
            self.active_agents.remove(pos);
        }
    }

    pub fn set_todo(&mut self, todo: Option<TodoSummary>) {
        self.todo = todo;
    }

    // ── Agent linking methods ────────────────────────────────────────

    pub fn push_pending_task(
        &mut self,
        tool_use_id: String,
        description: String,
        agent_type: Option<String>,
        model: Option<String>,
        event_ts: Option<u64>,
        message_id: Option<String>,
    ) {
        self.pending_tasks.push(PendingTask {
            tool_use_id,
            description,
            agent_type,
            model,
            event_ts,
            message_id,
        });
    }

    /// FIFO pop from pending queue, store bidirectional link.
    pub fn link_agent_to_pending_task(&mut self, agent_id: &str) -> Option<PendingTask> {
        if self.pending_tasks.is_empty() {
            return None;
        }
        let pending = self.pending_tasks.remove(0);
        self.task_agent_links
            .insert(pending.tool_use_id.clone(), agent_id.to_string());
        Some(pending)
    }

    /// Lookup linked agentId for a tool_use_id.
    pub fn resolve_task_agent(&self, tool_use_id: &str) -> Option<String> {
        self.task_agent_links.get(tool_use_id).cloned()
    }

    /// Remove an unlinked pending task by its tool_use_id.
    pub fn drain_pending_task(&mut self, tool_use_id: &str) -> Option<PendingTask> {
        let pos = self
            .pending_tasks
            .iter()
            .position(|p| p.tool_use_id == tool_use_id)?;
        Some(self.pending_tasks.remove(pos))
    }

    /// Check if an agent was linked from an Agent tool_use.
    pub fn is_task_linked_agent(&self, agent_id: &str) -> bool {
        self.task_agent_links.values().any(|id| id == agent_id)
    }

    // ── Todo tracking methods ────────────────────────────────────────

    pub fn create_task_item(&mut self, subject: String, active_form: Option<String>) {
        self.task_counter += 1;
        let id = self.task_counter.to_string();
        self.task_items.insert(
            id,
            TaskItem {
                subject,
                status: "pending".to_string(),
                active_form,
                started_at: None,
            },
        );
        self.rebuild_todo_from_tasks();
    }

    pub fn update_task_item(&mut self, task_id: &str, status: &str) {
        if status == "deleted" {
            self.task_items.remove(task_id);
        } else if let Some(item) = self.task_items.get_mut(task_id) {
            if status == "in_progress" && item.status != "in_progress" {
                item.started_at = Some(cache::now_epoch_ms());
            }
            item.status = status.to_string();
        }
        self.rebuild_todo_from_tasks();
    }

    fn rebuild_todo_from_tasks(&mut self) {
        if self.task_items.is_empty() {
            self.todo = None;
            return;
        }
        let total = self.task_items.len();
        let completed = self
            .task_items
            .values()
            .filter(|item| item.status == "completed")
            .count();
        let pending = total.saturating_sub(completed);

        // Collect in-progress items, sorted by started_at (earliest first)
        let mut in_progress_items: Vec<TodoInProgressItem> = self
            .task_items
            .values()
            .filter(|item| item.status == "in_progress")
            .map(|item| TodoInProgressItem {
                text: item
                    .active_form
                    .clone()
                    .unwrap_or_else(|| item.subject.clone()),
                started_at: item.started_at,
            })
            .collect();
        in_progress_items.sort_by_key(|item| item.started_at.unwrap_or(u64::MAX));

        let all_done = pending == 0 && completed > 0;

        self.todo = Some(TodoSummary {
            text: format!("{completed}/{total} done, {pending} pending"),
            pending,
            completed,
            total,
            in_progress_items,
            all_done,
            is_task_api: true,
        });
    }

    pub fn capped_tools(&self, max_lines: usize) -> Vec<ToolSummary> {
        if max_lines == 0 {
            return Vec::new();
        }
        let start = self.recent_tools.len().saturating_sub(max_lines);
        self.recent_tools[start..].to_vec()
    }

    /// Active agents first, then most recent completed, sliced to max_total.
    pub fn agents_for_display(&self, max_total: usize) -> Vec<AgentSummary> {
        if max_total == 0 {
            return Vec::new();
        }
        let mut result: Vec<AgentSummary> = Vec::new();

        // Active agents (running) — most recent N
        let active_start = self.active_agents.len().saturating_sub(max_total);
        result.extend_from_slice(&self.active_agents[active_start..]);

        // Fill remaining slots with completed agents (most recent first)
        let remaining = max_total.saturating_sub(result.len());
        if remaining > 0 {
            let mut completed: Vec<&AgentSummary> = self.completed_agents.iter().collect();
            completed.sort_by_key(|agent| std::cmp::Reverse(agent.completed_at));
            for agent in completed.into_iter().take(remaining) {
                result.push(agent.clone());
            }
        }

        result
    }

    /// Load state from a disk cache. Only restores env/git if within TTL.
    pub fn load_from_cache(&mut self, cache: SessionCache) {
        let now = cache::now_epoch_ms();

        // Transcript state
        self.last_transcript_offset = cache.transcript_offset;
        self.last_transcript_path = cache.transcript_path;
        self.active_tools = cache.active_tools;
        self.recent_tools = cache.recent_tools;
        self.active_agents = cache.active_agents;
        self.completed_agents = cache.completed_agents;
        self.completed_tool_counts = cache.completed_tool_counts;
        self.todo = cache.todo;
        self.pending_tasks = cache.pending_tasks;
        self.task_agent_links = cache.task_agent_links;
        self.task_items = cache.task_items;
        self.task_counter = cache.task_counter;
        self.cached_line3 = cache.line3;
        self.last_output_tokens = cache.last_output_tokens;
        self.last_output_token_time_ms = cache.last_output_token_time_ms;
        self.output_speed_toks_per_sec = cache.output_speed_toks_per_sec;
        // Defensive trim: a future cap reduction must not over-restore here.
        let mut history = cache.ctx_history;
        while history.len() > MAX_CTX_HISTORY {
            history.pop_front();
        }
        self.ctx_history = history;

        // Env/Git only if within TTL
        if let Some(entry) = cache.env {
            if now.saturating_sub(entry.cached_at_ms) < CACHE_TTL_MS {
                self.cached_env = Some((entry.path, entry.snapshot));
            }
        }
        if let Some(entry) = cache.git {
            if now.saturating_sub(entry.cached_at_ms) < CACHE_TTL_MS {
                self.cached_git = Some((entry.path, entry.snapshot));
            }
        }
    }

    /// Export current state to a cache struct for disk persistence.
    pub fn to_cache(&self) -> SessionCache {
        let now = cache::now_epoch_ms();
        SessionCache {
            transcript_offset: self.last_transcript_offset,
            transcript_path: self.last_transcript_path.clone(),
            active_tools: self.active_tools.clone(),
            recent_tools: self.recent_tools.clone(),
            active_agents: self.active_agents.clone(),
            completed_agents: self.completed_agents.clone(),
            completed_tool_counts: self.completed_tool_counts.clone(),
            todo: self.todo.clone(),
            pending_tasks: self.pending_tasks.clone(),
            task_agent_links: self.task_agent_links.clone(),
            task_items: self.task_items.clone(),
            task_counter: self.task_counter,
            line3: self.cached_line3.clone(),
            last_output_tokens: self.last_output_tokens,
            last_output_token_time_ms: self.last_output_token_time_ms,
            output_speed_toks_per_sec: self.output_speed_toks_per_sec,
            ctx_history: self.ctx_history.clone(),
            env: self.cached_env.as_ref().map(|(path, snapshot)| CacheEntry {
                path: path.clone(),
                snapshot: snapshot.clone(),
                cached_at_ms: now,
            }),
            git: self.cached_git.as_ref().map(|(path, snapshot)| CacheEntry {
                path: path.clone(),
                snapshot: snapshot.clone(),
                cached_at_ms: now,
            }),
        }
    }
}

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

    #[test]
    fn update_output_speed_first_call_returns_none() {
        let mut state = SessionState::default();
        let result = state.update_output_speed(Some(100));
        assert!(result.is_none(), "first call has no previous data");
        assert_eq!(state.last_output_tokens, Some(100));
    }

    #[test]
    fn update_output_speed_computes_rate_within_window() {
        let now = cache::now_epoch_ms();
        let mut state = SessionState {
            last_output_tokens: Some(100),
            last_output_token_time_ms: Some(now.saturating_sub(500)),
            output_speed_toks_per_sec: None,
            ..SessionState::default()
        };

        let result = state.update_output_speed(Some(150));
        assert!(result.is_some(), "should compute speed");
        let speed = result.unwrap();
        assert!(
            speed > 50.0 && speed < 200.0,
            "speed {speed} should be ~100 tok/s"
        );
    }

    #[test]
    fn update_output_speed_keeps_last_beyond_2s() {
        let now = cache::now_epoch_ms();
        let mut state = SessionState {
            last_output_tokens: Some(100),
            last_output_token_time_ms: Some(now.saturating_sub(3000)),
            output_speed_toks_per_sec: Some(42.0),
            ..SessionState::default()
        };

        let result = state.update_output_speed(Some(200));
        assert_eq!(result, Some(42.0), "should keep last known speed");
    }

    #[test]
    fn update_speed_resets_on_transcript_change() {
        let mut state = SessionState {
            last_output_tokens: Some(100),
            last_output_token_time_ms: Some(1000),
            output_speed_toks_per_sec: Some(50.0),
            last_transcript_path: Some("/old/path".to_string()),
            ..SessionState::default()
        };

        state.reset_transcript_if_path_changed("/new/path");

        assert!(state.last_output_tokens.is_none());
        assert!(state.last_output_token_time_ms.is_none());
        assert!(state.output_speed_toks_per_sec.is_none());
    }

    #[test]
    fn update_output_speed_none_tokens_preserves_state() {
        let mut state = SessionState {
            last_output_tokens: Some(100),
            last_output_token_time_ms: Some(1000),
            output_speed_toks_per_sec: Some(50.0),
            ..SessionState::default()
        };

        let result = state.update_output_speed(None);
        // Bug #2 fix: None tokens should NOT corrupt state
        assert_eq!(result, Some(50.0), "should return last known speed");
        assert_eq!(
            state.last_output_tokens,
            Some(100),
            "should not overwrite tokens"
        );
        assert_eq!(
            state.last_output_token_time_ms,
            Some(1000),
            "should not overwrite time"
        );
    }

    #[test]
    fn scored_completed_tools_recent_low_count_beats_stale_moderate_count() {
        let mut state = SessionState::default();
        let now = cache::now_epoch_ms();

        // Moderate count tool, last used 3 minutes ago (beyond 2-min decay window)
        // Score: 3 + 0 = 3.0
        state
            .completed_tool_counts
            .insert("Grep".to_string(), (3, now.saturating_sub(180_000)));

        // Low count tool, used just now
        // Score: 1 + 3.0 = 4.0
        state
            .completed_tool_counts
            .insert("Skill".to_string(), (1, now));

        let top = state.scored_completed_tools(2);
        assert_eq!(top[0].name, "Skill", "recent tool should rank first");
        assert_eq!(top[1].name, "Grep", "stale tool should rank second");
    }

    #[test]
    fn scored_completed_tools_count_dominates_after_decay() {
        let mut state = SessionState::default();
        let now = cache::now_epoch_ms();

        // Both tools last used 3 minutes ago (beyond 2-min decay window)
        state
            .completed_tool_counts
            .insert("Read".to_string(), (15, now.saturating_sub(180_000)));
        state
            .completed_tool_counts
            .insert("Skill".to_string(), (1, now.saturating_sub(180_000)));

        let top = state.scored_completed_tools(2);
        assert_eq!(
            top[0].name, "Read",
            "higher count should win when both are stale"
        );
    }

    #[test]
    fn scored_completed_tools_tiebreak_alphabetical() {
        let mut state = SessionState::default();
        let now = cache::now_epoch_ms();

        // Same count, same recency
        state
            .completed_tool_counts
            .insert("Edit".to_string(), (5, now));
        state
            .completed_tool_counts
            .insert("Bash".to_string(), (5, now));

        let top = state.scored_completed_tools(2);
        assert_eq!(
            top[0].name, "Bash",
            "alphabetical tiebreak: Bash before Edit"
        );
        assert_eq!(top[1].name, "Edit");
    }
}
