smart-tree 8.0.1

//! Real-time state synchronization between MCP and browser dashboard
//!
//! Provides WebSocket endpoint `/ws/state` for:
//! - Pushing MCP activity updates to browser at 60fps
//! - Receiving user hints/nudges from browser

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State,
    },
    response::Response,
};
use chrono::{DateTime, Utc};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use std::collections::{HashSet, VecDeque};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};

use super::SharedState;

// ============================================================================
// MCP Activity Tracking Types
// ============================================================================

/// Currently executing MCP tool
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActiveTool {
    /// Tool name (e.g., "search", "read_file")
    pub name: String,
    /// When execution started (serialized as epoch millis)
    #[serde(with = "instant_serde")]
    pub started_at: Instant,
    /// Tool parameters (simplified for display)
    pub parameters: serde_json::Value,
    /// Estimated progress 0.0 to 1.0 (if available)
    pub progress: Option<f32>,
}

/// File access event for Wave Compass visualization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileAccessEvent {
    /// File path accessed
    pub path: PathBuf,
    /// Type of access
    pub access_type: AccessType,
    /// When access occurred (serialized as epoch millis)
    #[serde(with = "instant_serde")]
    pub timestamp: Instant,
    /// Which MCP tool accessed it
    pub tool_name: String,
}

/// Type of file access
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AccessType {
    Read,
    Write,
    Analyze,
    Search,
}

/// Completed tool execution record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolExecution {
    /// Tool name
    pub name: String,
    /// When execution completed
    pub completed_at: DateTime<Utc>,
    /// How long it took
    pub duration_ms: u64,
    /// Whether it succeeded
    pub success: bool,
    /// Brief result summary
    pub summary: String,
}

/// Real-time MCP activity state for dashboard visualization
#[derive(Debug, Default)]
pub struct McpActivityState {
    /// Currently executing tool (if any)
    pub active_tool: Option<ActiveTool>,
    /// Recent file access events (ring buffer, last 100)
    pub file_access_log: VecDeque<FileAccessEvent>,
    /// Human-readable current operation
    pub current_operation: String,
    /// Files touched this session
    pub files_touched: HashSet<PathBuf>,
    /// Directories explored this session
    pub directories_explored: HashSet<PathBuf>,
    /// Recent tool executions (last 20)
    pub tool_history: VecDeque<ToolExecution>,
    /// Last state update timestamp
    pub last_update: Option<Instant>,
}

impl McpActivityState {
    /// Maximum file access events to keep
    const MAX_FILE_EVENTS: usize = 100;
    /// Maximum tool history entries
    const MAX_TOOL_HISTORY: usize = 20;

    /// Record start of a tool execution
    pub fn start_tool(&mut self, name: &str, parameters: serde_json::Value) {
        self.active_tool = Some(ActiveTool {
            name: name.to_string(),
            started_at: Instant::now(),
            parameters,
            progress: None,
        });
        self.current_operation = format!("Executing {}...", name);
        self.last_update = Some(Instant::now());
    }

    /// Update tool progress
    pub fn update_progress(&mut self, progress: f32) {
        if let Some(ref mut tool) = self.active_tool {
            tool.progress = Some(progress.clamp(0.0, 1.0));
            self.last_update = Some(Instant::now());
        }
    }

    /// Record file access event
    pub fn record_file_access(&mut self, path: PathBuf, access_type: AccessType, tool_name: &str) {
        let event = FileAccessEvent {
            path: path.clone(),
            access_type,
            timestamp: Instant::now(),
            tool_name: tool_name.to_string(),
        };

        self.file_access_log.push_back(event);
        if self.file_access_log.len() > Self::MAX_FILE_EVENTS {
            self.file_access_log.pop_front();
        }

        self.files_touched.insert(path.clone());
        if let Some(parent) = path.parent() {
            self.directories_explored.insert(parent.to_path_buf());
        }

        self.last_update = Some(Instant::now());
    }

    /// Complete tool execution
    pub fn complete_tool(&mut self, success: bool, summary: &str) {
        if let Some(tool) = self.active_tool.take() {
            let duration = tool.started_at.elapsed();
            let execution = ToolExecution {
                name: tool.name,
                completed_at: Utc::now(),
                duration_ms: duration.as_millis() as u64,
                success,
                summary: summary.to_string(),
            };

            self.tool_history.push_back(execution);
            if self.tool_history.len() > Self::MAX_TOOL_HISTORY {
                self.tool_history.pop_front();
            }
        }

        self.current_operation = if success {
            "Ready".to_string()
        } else {
            "Error occurred".to_string()
        };
        self.last_update = Some(Instant::now());
    }

    /// Get recent file events with age in milliseconds
    pub fn recent_file_events(&self, max_age_ms: u64) -> Vec<FileEventDto> {
        let now = Instant::now();
        self.file_access_log
            .iter()
            .filter_map(|event| {
                let age = now.duration_since(event.timestamp).as_millis() as u64;
                if age <= max_age_ms {
                    Some(FileEventDto {
                        path: event.path.to_string_lossy().to_string(),
                        access_type: event.access_type,
                        age_ms: age,
                        tool_name: event.tool_name.clone(),
                    })
                } else {
                    None
                }
            })
            .collect()
    }
}

// ============================================================================
// User Hints Types
// ============================================================================

/// User hint from browser to AI
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserHint {
    /// Type of hint
    pub hint_type: HintType,
    /// Content/description
    pub content: String,
    /// When hint was sent
    pub timestamp: DateTime<Utc>,
    /// Whether hint has been consumed by MCP
    #[serde(default)]
    pub consumed: bool,
}

/// Type of user hint
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum HintType {
    /// User clicked on Wave Compass region
    Click { target: String },
    /// User typed text hint
    Text { message: String },
    /// User voice input (transcribed)
    Voice { transcript: String, salience: f32 },
}

/// Queue of user hints waiting to be consumed by MCP
#[derive(Debug, Default)]
pub struct UserHintsQueue {
    pub hints: VecDeque<UserHint>,
    pub max_size: usize,
}

impl UserHintsQueue {
    const DEFAULT_MAX_SIZE: usize = 50;

    pub fn new() -> Self {
        Self {
            hints: VecDeque::new(),
            max_size: Self::DEFAULT_MAX_SIZE,
        }
    }

    /// Add a new hint
    pub fn push(&mut self, hint: UserHint) {
        self.hints.push_back(hint);
        while self.hints.len() > self.max_size {
            self.hints.pop_front();
        }
    }

    /// Get and consume the next unconsumed hint
    pub fn consume_next(&mut self) -> Option<UserHint> {
        for hint in &mut self.hints {
            if !hint.consumed {
                hint.consumed = true;
                return Some(hint.clone());
            }
        }
        None
    }

    /// Peek at unconsumed hints without consuming
    pub fn peek_unconsumed(&self) -> Vec<&UserHint> {
        self.hints.iter().filter(|h| !h.consumed).collect()
    }

    /// Count unconsumed hints
    pub fn unconsumed_count(&self) -> usize {
        self.hints.iter().filter(|h| !h.consumed).count()
    }

    /// Clear consumed hints older than threshold
    pub fn gc(&mut self, max_age: Duration) {
        let now = Utc::now();
        self.hints.retain(|h| {
            !h.consumed || (now - h.timestamp).num_milliseconds() < max_age.as_millis() as i64
        });
    }
}

// ============================================================================
// WebSocket Protocol DTOs
// ============================================================================

/// State update sent to browser (60fps)
#[derive(Debug, Serialize)]
pub struct StateUpdateDto {
    #[serde(rename = "type")]
    pub msg_type: &'static str,
    pub timestamp: i64,
    pub mcp: McpStateDto,
    pub file_log: Vec<FileEventDto>,
    pub wave_compass: WaveCompassDto,
    pub hints_pending: usize,
}

#[derive(Debug, Serialize)]
pub struct McpStateDto {
    pub active_tool: Option<String>,
    pub current_operation: String,
    pub progress: Option<f32>,
    pub tools_executed: usize,
}

#[derive(Debug, Serialize)]
pub struct FileEventDto {
    pub path: String,
    pub access_type: AccessType,
    pub age_ms: u64,
    pub tool_name: String,
}

#[derive(Debug, Serialize)]
pub struct WaveCompassDto {
    pub hot_regions: Vec<HotRegion>,
    pub trail: Vec<[f32; 2]>,
}

#[derive(Debug, Serialize)]
pub struct HotRegion {
    pub x: f32,
    pub y: f32,
    pub intensity: f32,
    pub label: String,
}

/// Hint message from browser
#[derive(Debug, Deserialize)]
pub struct HintMessageDto {
    #[serde(rename = "type")]
    pub msg_type: String,
    pub hint_type: String,
    #[serde(default)]
    pub target: Option<String>,
    #[serde(default)]
    pub content: Option<String>,
    #[serde(default)]
    pub transcript: Option<String>,
    #[serde(default)]
    pub salience: Option<f32>,
}

// ============================================================================
// WebSocket Handler
// ============================================================================

/// Handle WebSocket upgrade for state sync
pub async fn state_handler(
    ws: WebSocketUpgrade,
    State(state): State<SharedState>,
) -> Response {
    ws.on_upgrade(|socket| handle_state_socket(socket, state))
}

/// Handle the WebSocket connection for real-time state sync
async fn handle_state_socket(socket: WebSocket, state: SharedState) {
    let (mut sender, mut receiver) = socket.split();

    // Increment connection count
    {
        let mut dashboard = state.write().await;
        dashboard.connections += 1;
    }

    // Clone handles for tasks
    let mcp_activity = {
        let dashboard = state.read().await;
        Arc::clone(&dashboard.mcp_activity)
    };
    let user_hints = {
        let dashboard = state.read().await;
        Arc::clone(&dashboard.user_hints)
    };

    // Task: Send state updates at 60fps
    let mcp_activity_send = Arc::clone(&mcp_activity);
    let user_hints_send = Arc::clone(&user_hints);
    let send_task = tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_millis(16)); // ~60fps
        let mut last_update: Option<Instant> = None;

        loop {
            interval.tick().await;

            let activity = mcp_activity_send.read().await;

            // Skip if no updates since last send
            if activity.last_update == last_update {
                continue;
            }
            last_update = activity.last_update;

            // Build state update
            let update = build_state_update(&activity, &*user_hints_send.read().await);
            drop(activity);

            let json = match serde_json::to_string(&update) {
                Ok(j) => j,
                Err(_) => continue,
            };

            if sender.send(Message::Text(json)).await.is_err() {
                break; // Client disconnected
            }
        }
    });

    // Task: Receive hints from browser
    let recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = receiver.next().await {
            if let Message::Text(text) = msg {
                if let Ok(hint_msg) = serde_json::from_str::<HintMessageDto>(&text) {
                    if hint_msg.msg_type == "hint" {
                        let hint = parse_hint_message(hint_msg);
                        user_hints.write().await.push(hint);
                    }
                }
            }
        }
    });

    // Wait for either task to complete
    tokio::select! {
        _ = send_task => {},
        _ = recv_task => {},
    }

    // Decrement connection count
    {
        let mut dashboard = state.write().await;
        dashboard.connections = dashboard.connections.saturating_sub(1);
    }
}

/// Build state update DTO from current activity state
fn build_state_update(activity: &McpActivityState, hints: &UserHintsQueue) -> StateUpdateDto {
    let file_log = activity.recent_file_events(10_000); // Last 10 seconds

    // Build wave compass data from file events
    let wave_compass = build_wave_compass(&file_log, &activity.directories_explored);

    StateUpdateDto {
        msg_type: "state_update",
        timestamp: Utc::now().timestamp_millis(),
        mcp: McpStateDto {
            active_tool: activity.active_tool.as_ref().map(|t| t.name.clone()),
            current_operation: activity.current_operation.clone(),
            progress: activity.active_tool.as_ref().and_then(|t| t.progress),
            tools_executed: activity.tool_history.len(),
        },
        file_log,
        wave_compass,
        hints_pending: hints.unconsumed_count(),
    }
}

/// Build Wave Compass visualization data
fn build_wave_compass(file_log: &[FileEventDto], _directories: &HashSet<PathBuf>) -> WaveCompassDto {
    use std::collections::HashMap;

    // Aggregate intensity by directory
    let mut dir_intensity: HashMap<String, f32> = HashMap::new();
    for event in file_log {
        let dir = PathBuf::from(&event.path)
            .parent()
            .map(|p| p.to_string_lossy().to_string())
            .unwrap_or_default();

        // Decay intensity by age (newer = brighter)
        let intensity = 1.0 - (event.age_ms as f32 / 10_000.0).min(1.0);
        *dir_intensity.entry(dir).or_default() += intensity * 0.3;
    }

    // Convert to hot regions with pseudo-random but stable coordinates
    let hot_regions: Vec<HotRegion> = dir_intensity
        .into_iter()
        .map(|(dir, intensity)| {
            let (x, y) = path_to_coords(&dir);
            HotRegion {
                x,
                y,
                intensity: intensity.min(1.0),
                label: dir.split('/').next_back().unwrap_or(&dir).to_string(),
            }
        })
        .filter(|r| r.intensity > 0.05)
        .collect();

    // Build trail from recent file accesses
    let trail: Vec<[f32; 2]> = file_log
        .iter()
        .rev()
        .take(20)
        .map(|e| {
            let (x, y) = path_to_coords(&e.path);
            [x, y]
        })
        .collect();

    WaveCompassDto { hot_regions, trail }
}

/// Convert file path to Wave Compass coordinates using directory clustering
fn path_to_coords(path: &str) -> (f32, f32) {
    // Simple but stable hash-based positioning with directory clustering
    let parts: Vec<&str> = path.split('/').filter(|s| !s.is_empty()).collect();

    if parts.is_empty() {
        return (0.5, 0.5);
    }

    // First directory determines quadrant
    let quadrant = match parts.first().copied() {
        Some("src") => (0.0, 0.0),   // Upper-left
        Some("tests") => (0.5, 0.0), // Upper-right
        Some("docs") => (0.0, 0.5),  // Lower-left
        Some("scripts") => (0.5, 0.5), // Lower-right
        Some("examples") => (0.25, 0.25),
        _ => (0.25, 0.75),
    };

    // Sub-path determines position within quadrant
    let sub_hash = simple_hash(&parts[1..].join("/"));
    let x = quadrant.0 + (sub_hash % 100) as f32 / 200.0;
    let y = quadrant.1 + ((sub_hash / 100) % 100) as f32 / 200.0;

    (x.clamp(0.02, 0.98), y.clamp(0.02, 0.98))
}

/// Simple string hash for stable positioning
fn simple_hash(s: &str) -> u32 {
    s.bytes().fold(0u32, |acc, b| acc.wrapping_mul(31).wrapping_add(b as u32))
}

/// Parse hint message from browser
fn parse_hint_message(msg: HintMessageDto) -> UserHint {
    let content = msg.content.unwrap_or_default();

    let hint_type = match msg.hint_type.as_str() {
        "click" => HintType::Click {
            target: msg.target.unwrap_or_default(),
        },
        "text" => HintType::Text {
            message: content.clone(),
        },
        "voice" => HintType::Voice {
            transcript: msg.transcript.unwrap_or_default(),
            salience: msg.salience.unwrap_or(0.5),
        },
        _ => HintType::Text {
            message: content.clone(),
        },
    };

    UserHint {
        hint_type,
        content,
        timestamp: Utc::now(),
        consumed: false,
    }
}

// ============================================================================
// Instant serde helper (serialize as duration since UNIX_EPOCH-ish)
// ============================================================================

mod instant_serde {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use std::time::Instant;

    // We can't serialize Instant to absolute time, so we store "age" in ms
    // When deserializing, we approximate by subtracting from now

    pub fn serialize<S>(instant: &Instant, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        // Serialize as milliseconds ago
        let age_ms = instant.elapsed().as_millis() as u64;
        age_ms.serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Instant, D::Error>
    where
        D: Deserializer<'de>,
    {
        let age_ms = u64::deserialize(deserializer)?;
        Ok(Instant::now() - std::time::Duration::from_millis(age_ms))
    }
}