dissolve-python 0.3.0

//! Concurrent LSP client implementation using message passing
//!
//! This module provides high-performance, concurrent access to LSP servers
//! like pyright-langserver using sync threads and channels - no async needed!

use anyhow::{anyhow, Result};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read, Write};
use std::process::{Command, Stdio};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use tracing::{info, warn};
use url::Url;

/// LSP request with response channel
struct LspRequest {
    id: u64,
    method: String,
    params: Value,
    response_tx: mpsc::Sender<Result<Value>>,
}

/// LSP notification (no response expected)  
struct LspNotification {
    method: String,
    params: Value,
}

enum LspMessage {
    Request(LspRequest),
    Notification(LspNotification),
    Shutdown,
}

/// Sync concurrent Pyright LSP client - uses message passing for true parallelism
/// Multiple threads can make requests simultaneously without blocking each other!
pub struct SyncConcurrentPyrightClient {
    message_tx: mpsc::Sender<LspMessage>,
    request_id: Arc<AtomicU64>,
    workspace_files: Arc<Mutex<HashMap<String, i32>>>, // file_path -> version
    _workspace_root: Option<String>,
    _process_handle: Arc<thread::JoinHandle<()>>,
}

impl SyncConcurrentPyrightClient {
    /// Create and initialize a new sync concurrent pyright client
    pub fn new(workspace_root: Option<&str>) -> Result<Self> {
        info!("Starting sync concurrent pyright-langserver...");

        // Start pyright-langserver process
        let mut process = Command::new("pyright-langserver")
            .arg("--stdio")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .map_err(|e| anyhow!("Failed to start pyright-langserver: {}", e))?;

        let mut stdin = process
            .stdin
            .take()
            .ok_or_else(|| anyhow!("Failed to get stdin"))?;
        let stdout = process
            .stdout
            .take()
            .ok_or_else(|| anyhow!("Failed to get stdout"))?;

        let (message_tx, message_rx) = mpsc::channel();
        let request_id = Arc::new(AtomicU64::new(1));

        // Shared pending requests map for coordinating responses
        let pending_requests_arc = Arc::new(Mutex::new(
            HashMap::<u64, mpsc::Sender<Result<Value>>>::new(),
        ));
        let pending_clone = pending_requests_arc.clone();

        // Thread to read server responses and forward to pending requests
        let _response_reader = thread::spawn(move || {
            let mut stdout_reader = BufReader::new(stdout);
            while let Ok(response) = Self::read_lsp_message(&mut stdout_reader) {
                if let Some(id) = response.get("id").and_then(|v| v.as_u64()) {
                    let mut pending = pending_clone.lock().unwrap();
                    if let Some(tx) = pending.remove(&id) {
                        let result = if let Some(error) = response.get("error") {
                            Err(anyhow!("LSP error: {}", error))
                        } else {
                            Ok(response.get("result").cloned().unwrap_or(Value::Null))
                        };
                        let _ = tx.send(result);
                    }
                }
            }
        });

        // Main message processing thread
        let process_handle = thread::spawn(move || {
            // Process client messages and send to LSP server
            while let Ok(msg) = message_rx.recv() {
                match msg {
                    LspMessage::Request(req) => {
                        {
                            let mut pending = pending_requests_arc.lock().unwrap();
                            pending.insert(req.id, req.response_tx);
                        }

                        let lsp_request = json!({
                            "jsonrpc": "2.0",
                            "id": req.id,
                            "method": req.method,
                            "params": req.params
                        });

                        let request_str = format!(
                            "Content-Length: {}\r\n\r\n{}",
                            lsp_request.to_string().len(),
                            lsp_request
                        );

                        if let Err(e) = stdin.write_all(request_str.as_bytes()) {
                            warn!("Failed to write LSP request: {}", e);
                            let mut pending = pending_requests_arc.lock().unwrap();
                            if let Some(tx) = pending.remove(&req.id) {
                                let _ = tx.send(Err(anyhow!("Failed to send request: {}", e)));
                            }
                        } else if let Err(e) = stdin.flush() {
                            warn!("Failed to flush LSP request: {}", e);
                            let mut pending = pending_requests_arc.lock().unwrap();
                            if let Some(tx) = pending.remove(&req.id) {
                                let _ = tx.send(Err(anyhow!("Failed to flush request: {}", e)));
                            }
                        }
                    }
                    LspMessage::Notification(notif) => {
                        let lsp_notif = json!({
                            "jsonrpc": "2.0",
                            "method": notif.method,
                            "params": notif.params
                        });

                        let notif_str = format!(
                            "Content-Length: {}\r\n\r\n{}",
                            lsp_notif.to_string().len(),
                            lsp_notif
                        );

                        if let Err(e) = stdin.write_all(notif_str.as_bytes()) {
                            warn!("Failed to write LSP notification: {}", e);
                        } else {
                            let _ = stdin.flush();
                        }
                    }
                    LspMessage::Shutdown => break,
                }
            }
        });

        // Initialize LSP server
        let client = Self {
            message_tx,
            request_id,
            workspace_files: Arc::new(Mutex::new(HashMap::new())),
            _workspace_root: workspace_root.map(|s| s.to_string()),
            _process_handle: Arc::new(process_handle),
        };

        client.initialize(workspace_root)?;
        Ok(client)
    }

    fn read_lsp_message(reader: &mut BufReader<std::process::ChildStdout>) -> Result<Value> {
        // Read headers
        let mut headers = Vec::new();
        loop {
            let mut line = String::new();
            reader.read_line(&mut line)?;
            let line = line.trim();
            if line.is_empty() {
                break; // End of headers
            }
            headers.push(line.to_string());
        }

        // Parse Content-Length header
        let content_length = headers
            .iter()
            .find(|h| h.starts_with("Content-Length:"))
            .and_then(|h| h.split(':').nth(1))
            .and_then(|v| v.trim().parse::<usize>().ok())
            .ok_or_else(|| anyhow!("Missing or invalid Content-Length header"))?;

        // Read content
        let mut content = vec![0u8; content_length];
        reader.read_exact(&mut content)?;

        // Parse JSON
        let content_str = String::from_utf8(content)?;
        serde_json::from_str(&content_str).map_err(|e| anyhow!("Failed to parse JSON: {}", e))
    }

    fn initialize(&self, workspace_root: Option<&str>) -> Result<()> {
        let root_uri = workspace_root
            .and_then(|root| Url::from_file_path(root).ok())
            .map(|url| url.to_string());

        let params = json!({
            "processId": std::process::id(),
            "rootUri": root_uri,
            "capabilities": {
                "textDocument": {
                    "hover": {
                        "contentFormat": ["plaintext", "markdown"]
                    }
                }
            }
        });

        let (response_tx, response_rx) = mpsc::channel();
        let req_id = self.request_id.fetch_add(1, Ordering::SeqCst);

        let request = LspRequest {
            id: req_id,
            method: "initialize".to_string(),
            params,
            response_tx,
        };

        self.message_tx
            .send(LspMessage::Request(request))
            .map_err(|e| anyhow!("Failed to send initialize request: {}", e))?;

        // Wait for response with timeout
        match response_rx.recv_timeout(std::time::Duration::from_secs(30)) {
            Ok(Ok(_)) => {
                // Send initialized notification
                let notif = LspNotification {
                    method: "initialized".to_string(),
                    params: json!({}),
                };

                self.message_tx
                    .send(LspMessage::Notification(notif))
                    .map_err(|e| anyhow!("Failed to send initialized notification: {}", e))?;

                info!("Sync concurrent Pyright LSP client initialized successfully");
                Ok(())
            }
            Ok(Err(e)) => Err(anyhow!("Initialize failed: {}", e)),
            Err(_) => Err(anyhow!("Initialize timed out")),
        }
    }

    /// Query type - truly concurrent across multiple threads!
    /// This is the key method that enables parallel test execution
    pub fn query_type_concurrent(
        &self,
        file_path: &str,
        content: &str,
        line: u32,
        column: u32,
    ) -> Result<Option<String>> {
        // Convert to absolute path if relative
        let abs_path = if std::path::Path::new(file_path).is_relative() {
            std::env::current_dir()?.join(file_path)
        } else {
            std::path::PathBuf::from(file_path)
        };

        // Open document first if needed
        self.open_document(&abs_path.to_string_lossy(), content)?;

        let uri = format!("file://{}", abs_path.display());
        let params = json!({
            "textDocument": {
                "uri": uri
            },
            "position": {
                "line": line - 1, // Convert to 0-based line numbering for LSP
                "character": column
            }
        });

        let (response_tx, response_rx) = mpsc::channel();
        let req_id = self.request_id.fetch_add(1, Ordering::SeqCst);

        let request = LspRequest {
            id: req_id,
            method: "textDocument/hover".to_string(),
            params,
            response_tx,
        };

        self.message_tx
            .send(LspMessage::Request(request))
            .map_err(|e| anyhow!("Failed to send hover request: {}", e))?;

        // Wait for response with timeout - return proper errors, not None!
        match response_rx.recv_timeout(std::time::Duration::from_secs(5)) {
            Ok(Ok(response)) => {
                // Parse hover response - match pyright's actual response format
                if let Some(hover) = response.as_object() {
                    if let Some(contents) = hover.get("contents") {
                        let type_info = match contents {
                            Value::String(s) => s.clone(),
                            Value::Object(obj) => {
                                if let Some(Value::String(s)) = obj.get("value") {
                                    s.clone()
                                } else {
                                    return Ok(None);
                                }
                            }
                            _ => return Ok(None),
                        };

                        // Parse pyright's hover format like the original client
                        tracing::debug!("Pyright hover response: {}", type_info);

                        // Check for module format first
                        if type_info.starts_with("(module) ") {
                            let module_start = "(module) ".len();
                            let module_end = type_info[module_start..]
                                .find('\n')
                                .map(|pos| module_start + pos)
                                .unwrap_or(type_info.len());
                            let module_name = type_info[module_start..module_end].trim();
                            return Ok(Some(module_name.to_string()));
                        }

                        // Check for class format
                        if type_info.starts_with("(class) ") {
                            let class_start = "(class) ".len();
                            let class_end = type_info[class_start..]
                                .find('\n')
                                .map(|pos| class_start + pos)
                                .unwrap_or(type_info.len());
                            let class_name = type_info[class_start..class_end].trim();
                            return Ok(Some(class_name.to_string()));
                        }

                        // Look for colon format for variables
                        if let Some(colon_pos) = type_info.find(':') {
                            let type_part = type_info[colon_pos + 1..].trim();

                            // Check if pyright returned "Unknown" - treat as no type info
                            if type_part == "Unknown" {
                                tracing::warn!(
                                    "Pyright returned 'Unknown' type at {}:{}:{}",
                                    file_path,
                                    line,
                                    column
                                );
                                return Ok(None);
                            }

                            return Ok(Some(type_part.to_string()));
                        }
                    }
                }
                Ok(None)
            }
            Ok(Err(e)) => Err(anyhow!("Hover request failed: {}", e)),
            Err(_) => Err(anyhow!(
                "Hover request timed out after 5 seconds for {}:{}",
                line,
                column
            )),
        }
    }

    fn open_document(&self, file_path: &str, content: &str) -> Result<()> {
        let mut files = self.workspace_files.lock().unwrap();
        if !files.contains_key(file_path) {
            files.insert(file_path.to_string(), 1);
            drop(files);

            let params = json!({
                "textDocument": {
                    "uri": format!("file://{}", file_path),
                    "languageId": "python",
                    "version": 1,
                    "text": content
                }
            });

            let notif = LspNotification {
                method: "textDocument/didOpen".to_string(),
                params,
            };

            self.message_tx
                .send(LspMessage::Notification(notif))
                .map_err(|e| anyhow!("Failed to send didOpen notification: {}", e))?;

            // Wait for pyright to be ready by polling for analysis completion
            self.wait_for_analysis_ready(file_path)?;
        }
        Ok(())
    }

    /// Wait for pyright to complete analysis by polling for diagnostics or other readiness indicators
    fn wait_for_analysis_ready(&self, file_path: &str) -> Result<()> {
        const MAX_WAIT_MS: u64 = 2000; // Maximum 2 seconds
        const POLL_INTERVAL_MS: u64 = 50; // Check every 50ms

        let start = std::time::Instant::now();

        // Strategy: Try a simple hover request on line 1 and see if we get any response
        // Once pyright starts responding to hover requests, it's likely ready
        while start.elapsed().as_millis() < MAX_WAIT_MS as u128 {
            // Send a simple hover request to line 1, column 1 to test readiness
            let params = json!({
                "textDocument": {
                    "uri": format!("file://{}", file_path)
                },
                "position": {
                    "line": 0,
                    "character": 0
                }
            });

            let (response_tx, response_rx) = mpsc::channel();
            let req_id = self.request_id.fetch_add(1, Ordering::SeqCst);

            let request = LspRequest {
                id: req_id,
                method: "textDocument/hover".to_string(),
                params,
                response_tx,
            };

            if self.message_tx.send(LspMessage::Request(request)).is_ok() {
                // Wait for response with a short timeout
                match response_rx.recv_timeout(std::time::Duration::from_millis(200)) {
                    Ok(Ok(_)) => {
                        // Got a response (even if it's null) - pyright is ready
                        tracing::debug!("Pyright analysis ready for {}", file_path);
                        return Ok(());
                    }
                    Ok(Err(_)) => {
                        // Got an error response - pyright is responding, so it's ready
                        tracing::debug!(
                            "Pyright analysis ready for {} (error response)",
                            file_path
                        );
                        return Ok(());
                    }
                    Err(_) => {
                        // Timeout - pyright might still be analyzing, continue polling
                    }
                }
            }

            std::thread::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS));
        }

        // If we've waited too long, log a warning but continue
        tracing::warn!(
            "Timeout waiting for pyright analysis of {}, continuing anyway",
            file_path
        );
        Ok(())
    }

    /// Shutdown the client gracefully  
    pub fn shutdown(&self) -> Result<()> {
        tracing::debug!("Shutting down sync concurrent pyright client");

        // Send shutdown message using the message channel
        if let Err(e) = self.message_tx.send(LspMessage::Shutdown) {
            tracing::warn!("Failed to send shutdown message: {}", e);
        }

        // Give the process a moment to shutdown gracefully
        std::thread::sleep(std::time::Duration::from_millis(100));

        Ok(())
    }
}