mirage 0.0.1

mod protocol;

use clap::Subcommand;
use openssh::{Session, SessionBuilder, Stdio};
use std::path::{Path, PathBuf};
use tokio::io::BufReader;
use tracing::{debug, info};

use crate::{Host, Port, Workspace};

#[derive(Subcommand)]
pub enum LspServerCommand {
    /// Rust Language Server
    RustAnalyzer {
        /// Additional arguments to forward to rust-analyzer
        #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
        args: Vec<String>,
    },
    /// C/C++ Language Server
    Clangd {
        /// Additional arguments to forward to clangd
        #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
        args: Vec<String>,
    },
    /// Python Language Server
    Pyright {
        /// Additional arguments to forward to pyright
        #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
        args: Vec<String>,
    },
}

#[derive(Debug, Clone)]
pub enum LspServer {
    RustAnalyzer { args: Vec<String> },
    Clangd { args: Vec<String> },
    Pyright { args: Vec<String> },
}

impl LspServer {
    pub fn name(&self) -> &str {
        match self {
            LspServer::RustAnalyzer { .. } => "rust-analyzer",
            LspServer::Clangd { .. } => "clangd",
            LspServer::Pyright { .. } => "pyright",
        }
    }

    pub fn args(&self) -> &[String] {
        match self {
            LspServer::RustAnalyzer { args } => args,
            LspServer::Clangd { args } => args,
            LspServer::Pyright { args } => args,
        }
    }
}

impl From<&LspServerCommand> for LspServer {
    fn from(cmd: &LspServerCommand) -> Self {
        match cmd {
            LspServerCommand::RustAnalyzer { args } => {
                LspServer::RustAnalyzer { args: args.clone() }
            }
            LspServerCommand::Clangd { args } => LspServer::Clangd { args: args.clone() },
            LspServerCommand::Pyright { args } => LspServer::Pyright { args: args.clone() },
        }
    }
}

/// Copy a local folder to a remote system and return the remote path
async fn sync_workspace(
    session: &Session,
    local_path: &Path,
    exclude_patterns: &[String],
) -> Result<String, Box<dyn std::error::Error>> {
    // Verify local path exists
    if !local_path.exists() {
        return Err(format!("Local path does not exist: {}", local_path.display()).into());
    }

    // Determine remote destination path (use basename of local path)
    let folder_name = local_path
        .file_name()
        .ok_or("Invalid path")?
        .to_string_lossy();
    let remote_path = format!("/tmp/{}", folder_name);

    info!(
        local = %local_path.display(),
        remote = %remote_path,
        "Copying folder to remote system"
    );

    // Create remote directory
    let output = session
        .command("mkdir")
        .arg("-p")
        .arg(&remote_path)
        .output()
        .await?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("Failed to create remote directory: {}", stderr).into());
    }

    debug!("Created remote directory: {}", remote_path);

    // Build rsync command with exclude patterns
    let mut rsync_cmd = "rsync -avz".to_string();

    for pattern in exclude_patterns {
        rsync_cmd.push_str(&format!(" --exclude='{}'", pattern));
        debug!("Excluding pattern: {}", pattern);
    }

    rsync_cmd.push_str(&format!(" {}/ {}/", local_path.display(), remote_path));

    info!("Running: {}", rsync_cmd);

    // Use rsync over SSH for efficient copying
    let output = session
        .command("bash")
        .arg("-c")
        .arg(&rsync_cmd)
        .output()
        .await?;

    if output.status.success() {
        let stdout = String::from_utf8_lossy(&output.stdout);
        debug!("Rsync output: {}", stdout);
        info!("Successfully copied folder to {}", remote_path);
    } else {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("Failed to copy folder: {}", stderr).into());
    }

    // Verify the copy
    let output = session
        .command("ls")
        .arg("-la")
        .arg(&remote_path)
        .output()
        .await?;

    if output.status.success() {
        let stdout = String::from_utf8_lossy(&output.stdout);
        debug!("Remote directory contents:\n{}", stdout);
    }

    Ok(remote_path)
}

/// Run LSP server in the remote directory and proxy LSP messages bidirectionally
async fn run_lsp_server(
    session: &Session,
    remote_path: &str,
    server_name: &str,
    server_args: &[String],
) -> Result<(), Box<dyn std::error::Error>> {
    // Check if LSP server exists on remote machine
    info!("Starting {} on remote machine", server_name);

    let output = session
        .command("bash")
        .arg("-c")
        .arg(format!("cd {} && {} --version", remote_path, server_name))
        .output()
        .await?;

    if output.status.success() {
        let stdout = String::from_utf8_lossy(&output.stdout);
        info!("{} version: {}", server_name, stdout.trim());
    } else {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("{} not found: {}", server_name, stderr).into());
    }

    // Build command with arguments
    let mut cmd = format!("cd {} && {}", remote_path, server_name);
    for arg in server_args {
        cmd.push_str(&format!(" {}", arg));
    }

    info!("Running: {}", cmd);
    info!("Starting LSP proxy (bidirectional JSON-RPC over stdin/stdout)");

    // Start LSP server with piped stdin/stdout
    let mut child = session
        .command("bash")
        .arg("-c")
        .arg(&cmd)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()
        .await?;

    // Get stdin/stdout handles
    let remote_stdin = child.stdin().take().ok_or("Failed to get remote stdin")?;
    let remote_stdout = child.stdout().take().ok_or("Failed to get remote stdout")?;

    // Spawn task to forward local stdin → remote stdin
    let stdin_to_remote = tokio::spawn(async move {
        let local_stdin = tokio::io::stdin();
        let mut local_stdin_reader = BufReader::new(local_stdin);
        let mut remote_stdin = remote_stdin;

        loop {
            match protocol::read_lsp_message(&mut local_stdin_reader).await {
                Ok(Some(message)) => {
                    debug!("→ Forwarding message to remote: {} bytes", message.len());
                    if let Err(e) = protocol::write_lsp_message(&mut remote_stdin, &message).await {
                        debug!("Failed to write to remote stdin: {}", e);
                        break;
                    }
                }
                Ok(None) => {
                    info!("Local stdin closed");
                    break;
                }
                Err(e) => {
                    debug!("Failed to read from local stdin: {}", e);
                    break;
                }
            }
        }
    });

    // Spawn task to forward remote stdout → local stdout
    let remote_to_stdout = tokio::spawn(async move {
        let mut remote_stdout_reader = BufReader::new(remote_stdout);
        let mut local_stdout = tokio::io::stdout();

        loop {
            match protocol::read_lsp_message(&mut remote_stdout_reader).await {
                Ok(Some(message)) => {
                    debug!("← Forwarding message from remote: {} bytes", message.len());
                    if let Err(e) = protocol::write_lsp_message(&mut local_stdout, &message).await {
                        debug!("Failed to write to local stdout: {}", e);
                        break;
                    }
                }
                Ok(None) => {
                    info!("Remote stdout closed");
                    break;
                }
                Err(e) => {
                    debug!("Failed to read from remote stdout: {}", e);
                    break;
                }
            }
        }
    });

    // Wait for either task to complete (or the child process to exit)
    tokio::select! {
        _ = stdin_to_remote => {
            info!("stdin→remote task completed");
        }
        _ = remote_to_stdout => {
            info!("remote→stdout task completed");
        }
        status = child.wait() => {
            match status {
                Ok(s) => info!("{} exited with status: {:?}", server_name, s),
                Err(e) => info!("Failed to wait for child: {}", e),
            }
        }
    }

    Ok(())
}

/// Start remote LSP: connect, sync workspace, and run LSP server
pub async fn start(
    host: Host,
    port: Port,
    workspace: Workspace,
    exclude_patterns: &[String],
    server: &LspServer,
) -> Result<(), Box<dyn std::error::Error>> {
    let address: String = host.into();
    let port = port.into();
    let local_path: PathBuf = workspace.into();

    info!(
        address = %address,
        port = port,
        path = %local_path.display(),
        server = %server.name(),
        "Connecting to remote host"
    );

    // Establish SSH session
    let session = SessionBuilder::default()
        .port(port)
        .connect(address)
        .await?;

    info!("SSH connection established");

    // Sync workspace to remote
    let remote_path = sync_workspace(&session, &local_path, exclude_patterns).await?;

    // Start LSP server
    run_lsp_server(&session, &remote_path, server.name(), server.args()).await?;

    session.close().await?;
    info!("SSH session closed");

    Ok(())
}