use arboard::Clipboard;
#[cfg(target_os = "linux")]
use arboard::SetExtLinux as _;
use clap::{Parser, ValueEnum};
use std::fs;
use std::io::{self, Read as _, Write as _};
use std::{env, error::Error, process};

/// dacopy - Copy text to the clipboard from stdin or files.
#[derive(Parser, Debug)]
#[command(author, version, about)]
struct Args {
    /// Files to copy. If none provided, reads from stdin.
    files: Vec<String>,

    /// Also print the copied text to stdout (tee mode).
    #[arg(short, long)]
    tee: bool,

    /// Do not trim leading/trailing whitespace from the input.
    #[arg(long = "no-trim", alias = "notrim", visible_alias = "notrim")]
    no_trim: bool,

    /// Store the copied text to the clipboard using the specified mode.
    #[arg(short, long, conflicts_with = "osc52", default_value = "auto")]
    mode: StoreToClipboardModeArg,

    /// Force using OSC-52 escape sequence instead of the system clipboard.
    /// Useful in SSH sessions or terminals that support OSC-52
    /// (e.g. kitty, iTerm2, tmux with set-clipboard on).
    ///
    /// Alias for '--mode osc52'.
    #[arg(short, long, conflicts_with = "mode")]
    osc52: bool,

    /// Verbose output.
    #[arg(short, long)]
    verbose: bool,
}

// An argument that can be passed into the program to signal that it should daemonize itself. This
// can be anything as long as it is unlikely to be passed in by the user by mistake.
const DAEMONIZE_ARG: &str = "__internal_daemonize";

#[derive(Debug, Copy, Clone, Default, ValueEnum)]
pub enum StoreToClipboardModeArg {
    #[default]
    #[value(name = "auto")]
    Auto,
    #[value(name = "osc52")]
    ForceOsc52,
    #[value(name = "system")]
    ForceSystemClipboard,
}

#[derive(Debug, Copy, Clone)]
pub enum StoreToClipboardMode {
    Osc52,
    SystemClipboard,
}

fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
    // Incantation: https://github.com/1Password/arboard/blob/0a73f82d454eff18f3423c8fd9b89477a17281aa/examples/daemonize.rs#L13-L14
    #[cfg(target_os = "linux")]
    if env::args().any(|x| x == DAEMONIZE_ARG) {
        let mut text = String::new();
        io::stdin().read_to_string(&mut text)?;
        Clipboard::new()?.set().wait().text(text)?;
        return Ok(());
    }

    let args = Args::parse();

    let content = if args.files.is_empty() {
        let mut buf = String::new();
        if let Err(e) = io::stdin().read_to_string(&mut buf) {
            eprintln!("dacopy: error reading stdin: {e}");
            process::exit(1);
        }
        buf
    } else {
        let mut buf = String::new();
        for path in &args.files {
            if path == DAEMONIZE_ARG {
                continue;
            }

            match fs::read_to_string(path) {
                Ok(text) => buf.push_str(&text),
                Err(e) => {
                    eprintln!("dacopy: cannot read '{path}': {e}");
                    process::exit(1);
                }
            }
        }
        buf
    };

    let text = if args.no_trim {
        content
    } else {
        content.trim().to_string()
    };

    let arg_selected_copy_mode: StoreToClipboardModeArg = if args.osc52 {
        StoreToClipboardModeArg::ForceOsc52
    } else {
        args.mode
    };

    let copy_mode: StoreToClipboardMode = match arg_selected_copy_mode {
        StoreToClipboardModeArg::ForceOsc52 => StoreToClipboardMode::Osc52,
        StoreToClipboardModeArg::ForceSystemClipboard => StoreToClipboardMode::SystemClipboard,
        StoreToClipboardModeArg::Auto => {
            // If connection is over SSH, use OSC 52 escape sequences. Otherwise, use the system clipboard.
            if std::env::var("SSH_CLIENT").is_ok() {
                StoreToClipboardMode::Osc52
            } else {
                StoreToClipboardMode::SystemClipboard
            }
        }
    };

    if args.verbose {
        eprintln!(
            "Copy mode requests by arguments: {arg_selected_copy_mode:?}. Copy mode: {copy_mode:?}."
        );
    }

    eprintln!(
        "Copying {} characters ({}).",
        text.chars().count(),
        format_size(text.len())
    );

    if args.tee {
        print!("{text}");
        if !text.ends_with("\n") {
            println!(); // Flush.
        }
    }

    match copy_mode {
        StoreToClipboardMode::Osc52 => {
            copy_osc52(&text)?;
        }
        StoreToClipboardMode::SystemClipboard => {
            copy_with_system_clipboard(&text)?;
        }
    };

    Ok(())
}

/// Write an OSC-52 clipboard escape sequence to the terminal.
///
/// OSC-52 is supported by kitty, iTerm2, foot, WezTerm, and tmux (with
/// `set-clipboard on`).  It works over SSH without any special forwarding
/// because the sequence travels through the terminal emulator itself.
///
/// The sequence is written to `/dev/tty` when available so that it reaches
/// the terminal even when stdout/stderr have been redirected (e.g. when the
/// user pipes output elsewhere while also using `--tee`).
fn copy_osc52(text: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
    let encoded = base64_encode(text.as_bytes());
    // OSC 52 ; c ; <base64> ST
    // "c" selects the primary clipboard; use "p" for primary selection on X11.
    let sequence = format!("\x1b]52;c;{encoded}\x07");

    // Prefer /dev/tty so the sequence reaches the terminal even when stdout
    // is redirected.
    if let Ok(mut tty) = fs::OpenOptions::new().write(true).open("/dev/tty") {
        tty.write_all(sequence.as_bytes())?;
    } else {
        // Fallback: write to stderr, which is usually still a tty.
        eprint!("{sequence}");
    }

    Ok(())
}

/// Copies text to the system clipboard using the appropriate method for the platform.
fn copy_with_system_clipboard(text: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
    if cfg!(target_os = "linux") {
        let mut child = process::Command::new(env::current_exe()?)
            .arg(DAEMONIZE_ARG)
            .stdin(process::Stdio::piped())
            .stdout(process::Stdio::null())
            .stderr(process::Stdio::null())
            .current_dir("/")
            .spawn()?;

        // Pass the text to the daemonized process via its stdin.
        child.stdin.take().unwrap().write_all(text.as_bytes())?;
    } else {
        Clipboard::new()?.set_text(text)?;
    }

    Ok(())
}

/// Minimal base64 encoder.
/// Implemented here to avoid pulling in an extra dependency.
fn base64_encode(input: &[u8]) -> String {
    const TABLE: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
    for chunk in input.chunks(3) {
        let b0 = chunk[0] as u32;
        let b1 = chunk.get(1).copied().unwrap_or(0) as u32;
        let b2 = chunk.get(2).copied().unwrap_or(0) as u32;
        let n = (b0 << 16) | (b1 << 8) | b2;
        out.push(TABLE[((n >> 18) & 0x3f) as usize] as char);
        out.push(TABLE[((n >> 12) & 0x3f) as usize] as char);
        out.push(if chunk.len() > 1 {
            TABLE[((n >> 6) & 0x3f) as usize] as char
        } else {
            '='
        });
        out.push(if chunk.len() > 2 {
            TABLE[(n & 0x3f) as usize] as char
        } else {
            '='
        });
    }
    out
}

fn format_size(size_bytes: usize) -> String {
    if size_bytes < 1024 {
        return format!("{} B", size_bytes);
    }
    let size_kb = size_bytes as f64 / 1024.0;
    if size_kb < 1024.0 {
        return format!("{:.2} KiB", size_kb);
    }
    let size_mb = size_kb / 1024.0;
    if size_mb < 1024.0 {
        return format!("{:.2} MiB", size_mb);
    }
    let size_gb = size_mb / 1024.0;
    format!("{:.2} GiB", size_gb)
}

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

    #[test]
    fn test_format_size() {
        assert_eq!(format_size(281655), "275.05 KiB");
    }

    #[test]
    fn test_base64_encode() {
        assert_eq!(base64_encode(b""), "");
        assert_eq!(base64_encode(b"f"), "Zg==");
        assert_eq!(base64_encode(b"fo"), "Zm8=");
        assert_eq!(base64_encode(b"foo"), "Zm9v");
        assert_eq!(base64_encode(b"hello world"), "aGVsbG8gd29ybGQ=");
    }
}
