use anyhow::{bail, Result};
use blahaj::Sharks;
use clap::{Parser, ValueEnum};
use std::io::Read;
use std::path::{Path, PathBuf};
use zeroize::Zeroize;

use keyquorum_core::share_format::{self, ShareEncoding, ShareFormatOptions};

#[derive(Parser)]
#[command(
    name = "keyquorum-split",
    about = "Split a secret into Shamir shares",
    long_about = "Split a secret into Shamir shares.\n\n\
        Reads a secret from stdin and generates N shares, any K of which can reconstruct \
        the original. By default, each share is output as a PEM envelope \
        (KEYQUORUM-SHARE-V1) with metadata headers and a CRC32 integrity check, and \
        a blake3 verification checksum is embedded in the secret before splitting.\n\n\
        Use --bare for the V1 binary payload without the PEM envelope, or --no-metadata \
        to keep the envelope but omit the headers.",
    version
)]
struct Cli {
    /// Total number of shares to generate (2-255)
    #[arg(short = 'n', long)]
    shares: u8,
    /// Minimum shares needed to reconstruct (2-N)
    #[arg(short = 'k', long)]
    threshold: u8,
    /// Output mode
    #[arg(short, long, value_enum, default_value_t = OutputMode::Stdout)]
    output: OutputMode,
    /// Output directory for file-per-share mode
    #[arg(short, long)]
    dir: Option<PathBuf>,
    /// Lockdown mode: rejects stdout output. May gain new restrictions between versions.
    #[arg(long)]
    lockdown: bool,
    /// Do not embed a blake3 verification checksum in the secret before splitting.
    ///
    /// WARNING: without a checksum, retry mode (on_failure="retry") cannot verify
    /// candidate secrets and will execute the configured action (cryptsetup, command,
    /// etc.) with each incorrect key attempt. This may cause repeated failed
    /// invocations of downstream tools.
    #[arg(long)]
    no_checksum: bool,
    /// Skip per-share CRC32 integrity check
    #[arg(long)]
    no_integrity: bool,
    /// Omit metadata headers (share number, threshold, scheme) from PEM envelope
    #[arg(long)]
    no_metadata: bool,
    /// Output V1 binary payload only, no PEM envelope
    #[arg(long)]
    bare: bool,
    /// Payload encoding
    #[arg(long, value_enum, default_value_t = Encoding::Base64)]
    encoding: Encoding,
    /// Disable strict hardening: allow operation if memory protections
    /// (mlock, madvise) fail. Not recommended for production.
    #[arg(long)]
    no_strict_hardening: bool,
    /// File containing age recipients (one age1... public key per line).
    /// Each share is encrypted to the corresponding recipient before writing.
    /// Requires --output age --dir <path>.
    // Future: --encrypt-to for inline comma-separated keys
    #[arg(long, value_name = "FILE")]
    recipients: Option<PathBuf>,
    /// Output age-encrypted shares as ASCII armor (PEM-like text).
    /// Useful for pasting into text-based channels (Signal, email, etc.).
    /// Files get .age.txt extension instead of .age.
    #[arg(long, alias = "armour")]
    armor: bool,
    /// Designate the last N generated shares as DURESS (canary) shares and
    /// print their indices plus a ready-to-paste [session.duress] config
    /// block (indices are deterministic 1..=N; this just saves you the
    /// bookkeeping and surfaces the security warning below).
    ///
    /// SECURITY: a duress share is a real Shamir share of the same secret.
    /// If you hand each participant a normal share AND a duress share, every
    /// person then holds 2 of N shares, so an attacker needs only ceil(K/2)
    /// people to collect K shares — roughly HALVING the people-collusion
    /// threshold. See the README "Duress shares" section before using this.
    #[arg(long, value_name = "N")]
    duress: Option<u8>,
    // Future: --raw-shares for truly raw sharks bytes (no KQ prefix),
    // for interoperability with other Shamir tools.
}

#[derive(Clone, ValueEnum)]
enum OutputMode {
    /// All shares to stdout (labels to stderr)
    Stdout,
    /// One file per share in the output directory
    Files,
    /// Encrypt each share to an age recipient before writing
    Age,
    /// Key ceremony: show one share at a time on the terminal, wait for
    /// confirmation, clear the screen between shares
    Interactive,
}

#[derive(Clone, ValueEnum)]
enum Encoding {
    Base64,
    Base32,
}

fn main() -> Result<()> {
    keyquorum_core::memory::warn_if_not_linux();
    keyquorum_core::memory::harden_process()?;

    let cli = Cli::parse();

    if cli.threshold < 2 {
        bail!("threshold must be at least 2");
    }
    if cli.threshold > cli.shares {
        bail!(
            "threshold ({}) must be <= total shares ({})",
            cli.threshold,
            cli.shares
        );
    }

    if matches!(cli.output, OutputMode::Files | OutputMode::Age) && cli.dir.is_none() {
        bail!("--dir is required when using --output files or --output age");
    }

    if matches!(cli.output, OutputMode::Age) && cli.recipients.is_none() {
        bail!("--recipients is required when using --output age");
    }

    if cli.recipients.is_some() && !matches!(cli.output, OutputMode::Age) {
        bail!("--recipients requires --output age");
    }

    if cli.armor && !matches!(cli.output, OutputMode::Age) {
        bail!("--armor requires --output age");
    }

    if let Some(d) = cli.duress {
        if d == 0 {
            bail!("--duress must be at least 1 (omit the flag for no duress shares)");
        }
        // Leave at least `threshold` non-duress shares so an honest quorum can
        // unlock using normal shares alone. duress > shares - threshold would
        // force duress shares into every legitimate reconstruction.
        if d > cli.shares - cli.threshold {
            bail!(
                "--duress {} leaves only {} normal share(s), but threshold is {}; \
                 use at most {} duress shares (shares - threshold)",
                d,
                cli.shares - d,
                cli.threshold,
                cli.shares - cli.threshold
            );
        }
    }

    if cli.lockdown && matches!(cli.output, OutputMode::Stdout) {
        bail!("lockdown mode rejects --output stdout: shares must not appear in terminal output. Use --output files --dir <path>");
    }

    if cli.lockdown && matches!(cli.output, OutputMode::Interactive) {
        bail!("lockdown mode rejects --output interactive: shares must not appear in terminal output. Use --output files or --output age");
    }

    // Read secret from stdin
    let mut secret = Vec::new();
    std::io::stdin().read_to_end(&mut secret)?;

    // Strip trailing newline if present (common when piping from echo)
    if secret.last() == Some(&b'\n') {
        secret.pop();
    }

    if secret.is_empty() {
        bail!("no secret provided on stdin");
    }

    // Embed blake3 verification checksum (default behavior).
    // Done before protect_secret so the final buffer (with checksum) is what gets mlocked.
    if !cli.no_checksum {
        // Reserve upfront to avoid reallocation after mlock
        secret.reserve(32);
        let hash = blake3::hash(&secret);
        secret.extend_from_slice(hash.as_bytes());
        eprintln!("Embedded blake3 verification checksum (32 bytes)");
        eprintln!("Use verification = \"embedded-blake3\" in daemon config.");
    } else {
        eprintln!("WARNING: no checksum embedded. If on_failure=\"retry\" is used,");
        eprintln!("the daemon will execute the configured action with each incorrect");
        eprintln!("candidate secret, causing repeated failed invocations of downstream tools.");
        eprintln!("Use verification = \"none\" in daemon config.");
    }

    // Apply memory protections to the final secret buffer (including checksum if appended)
    let strict_hardening = !cli.no_strict_hardening || cli.lockdown;
    if !strict_hardening {
        eprintln!("WARNING: strict_hardening disabled — memory protections will not be enforced");
    }
    let failures = keyquorum_core::memory::protect_secret(&secret);
    if !failures.is_empty() {
        for (name, err) in &failures {
            eprintln!("warning: memory protection {} failed: {}", name, err);
        }
        if strict_hardening {
            bail!("strict_hardening: failed to apply memory protections to secret");
        }
    }

    // Generate shares
    let sharks = Sharks(cli.threshold);
    let dealer = sharks.dealer(&secret);
    let shares: Vec<blahaj::Share> = dealer.take(cli.shares as usize).collect();

    // Build format options from CLI flags
    let encoding = match cli.encoding {
        Encoding::Base64 => ShareEncoding::Base64,
        Encoding::Base32 => ShareEncoding::Base32,
    };

    // Output shares
    match cli.output {
        OutputMode::Stdout => output_stdout(&shares, &cli, encoding),
        OutputMode::Files => output_files(
            &shares,
            &cli,
            encoding,
            cli.dir.as_ref().expect("validated above"),
        )?,
        OutputMode::Age => {
            let recipients =
                parse_recipients_file(cli.recipients.as_ref().expect("validated above"))?;
            if recipients.len() > shares.len() {
                bail!(
                    "too many recipients: {} recipients for {} shares",
                    recipients.len(),
                    shares.len()
                );
            }
            if recipients.len() < shares.len() {
                if cli.lockdown {
                    bail!(
                        "lockdown mode requires exactly {} recipients (one per share), got {}",
                        shares.len(),
                        recipients.len()
                    );
                }
                eprintln!(
                    "WARNING: {} recipients for {} shares — {} shares will be written unencrypted",
                    recipients.len(),
                    shares.len(),
                    shares.len() - recipients.len()
                );
            }
            output_age(
                &shares,
                &cli,
                encoding,
                cli.dir.as_ref().expect("validated above"),
                &recipients,
            )?
        }
        OutputMode::Interactive => output_interactive(&shares, &cli, encoding)?,
    }

    // Report duress shares. Indices are deterministic (blahaj assigns
    // x-coordinates 1..=N in order), but emitting the exact config block and
    // the security warning here removes a transcription footgun.
    if let Some(d) = cli.duress {
        let d = d as usize;
        let duress_indices: Vec<u8> = shares[shares.len() - d..]
            .iter()
            .map(share_x_coord)
            .collect();
        let normal = cli.shares as usize - d;
        let index_list = duress_indices
            .iter()
            .map(|i| i.to_string())
            .collect::<Vec<_>>()
            .join(", ");

        eprintln!();
        eprintln!("======================================================================");
        eprintln!("DURESS SHARES");
        eprintln!("======================================================================");
        eprintln!(
            "The last {} share(s) — share #{}..#{} — are duress shares.",
            d,
            normal + 1,
            cli.shares
        );
        eprintln!("Their indices are: {}", index_list);
        eprintln!();
        eprintln!("Add to the daemon config:");
        eprintln!();
        eprintln!("    [session.duress]");
        eprintln!("    indices = [{}]", index_list);
        eprintln!("    mode = \"alert\"   # or \"poison\"");
        eprintln!("    alert_program = \"/usr/local/bin/notify-security\"");
        eprintln!();
        eprintln!("!! SECURITY: each duress share is a REAL Shamir share of this secret.");
        eprintln!("!! If you give each participant a normal share AND a duress share, every");
        eprintln!(
            "!! person then holds 2 of {} shares, so an attacker needs only ceil(K/2)",
            cli.shares
        );
        eprintln!(
            "!! = {} people (not {}) to collect a quorum of {} shares.",
            cli.threshold.div_ceil(2),
            cli.threshold,
            cli.threshold
        );
        eprintln!("!! This roughly HALVES your people-collusion threshold. Choose -n/-k with");
        eprintln!("!! that in mind. For poison mode, every participant needs their own");
        eprintln!("!! distinct duress share, or coercion of a non-holder goes undetected.");
        eprintln!("======================================================================");
    }

    // Zeroize and unlock secret (in that order: Vec::zeroize() would clear
    // the Vec first, making the munlock a no-op on an empty slice)
    keyquorum_core::memory::wipe_and_unlock(&mut secret);

    Ok(())
}

/// The x-coordinate (index) of a share: the first byte of its serialization.
fn share_x_coord(share: &blahaj::Share) -> u8 {
    Vec::from(share)[0]
}

/// Label suffix marking share `i` (0-based) as duress, given the duress count.
/// The last `duress_count` shares are duress. Empty for normal shares.
fn duress_suffix(cli: &Cli, i: usize) -> &'static str {
    let d = cli.duress.unwrap_or(0) as usize;
    if d > 0 && i >= cli.shares as usize - d {
        " [DURESS]"
    } else {
        ""
    }
}

fn make_format_opts(cli: &Cli, encoding: ShareEncoding, share_number: u8) -> ShareFormatOptions {
    ShareFormatOptions {
        encoding,
        include_crc32: !cli.no_integrity,
        include_envelope: !cli.bare,
        include_metadata: !cli.no_metadata && !cli.bare,
        share_number,
        total_shares: cli.shares,
        threshold: cli.threshold,
    }
}

fn output_stdout(shares: &[blahaj::Share], cli: &Cli, encoding: ShareEncoding) {
    for (i, share) in shares.iter().enumerate() {
        let mut sharks_data: Vec<u8> = Vec::from(share);
        let index = sharks_data[0];
        let opts = make_format_opts(cli, encoding, (i + 1) as u8);
        let mut formatted = share_format::format_share(&sharks_data, &opts);
        eprintln!(
            "Share {} (index {}):{}",
            i + 1,
            index,
            duress_suffix(cli, i)
        );
        println!("{}", formatted);
        sharks_data.zeroize();
        formatted.zeroize();
    }
}

fn output_files(
    shares: &[blahaj::Share],
    cli: &Cli,
    encoding: ShareEncoding,
    dir: &PathBuf,
) -> Result<()> {
    std::fs::create_dir_all(dir).map_err(|e| {
        anyhow::anyhow!("failed to create output directory {}: {}", dir.display(), e)
    })?;

    for (i, share) in shares.iter().enumerate() {
        let mut sharks_data: Vec<u8> = Vec::from(share);
        let index = sharks_data[0];
        let opts = make_format_opts(cli, encoding, (i + 1) as u8);
        let mut formatted = share_format::format_share(&sharks_data, &opts);
        let filename = dir.join(format!("share-{}.txt", i + 1));
        let write_result = std::fs::write(&filename, format!("{}\n", formatted));
        sharks_data.zeroize();
        formatted.zeroize();
        write_result
            .map_err(|e| anyhow::anyhow!("failed to write {}: {}", filename.display(), e))?;
        eprintln!(
            "Share {} (index {}) written to {}{}",
            i + 1,
            index,
            filename.display(),
            duress_suffix(cli, i)
        );
    }

    eprintln!(
        "\n{} share files written to {}",
        shares.len(),
        dir.display()
    );
    eprintln!("Distribute each file to its holder, then delete this directory.");

    Ok(())
}

/// Key ceremony mode: show one share at a time on the controlling terminal,
/// wait for Enter between shares, clear the screen (including scrollback)
/// before and after each.
///
/// Uses /dev/tty directly for both display and confirmation — stdin has
/// already been consumed by the secret, and stdout may be redirected.
fn output_interactive(shares: &[blahaj::Share], cli: &Cli, encoding: ShareEncoding) -> Result<()> {
    use std::io::{BufRead, BufReader, Write};

    let tty = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open("/dev/tty")
        .map_err(|e| {
            anyhow::anyhow!(
                "interactive mode requires a controlling terminal (/dev/tty): {}",
                e
            )
        })?;
    let mut tty_out = tty.try_clone()?;
    let mut tty_in = BufReader::new(tty);

    // 2J clears the visible screen, 3J clears scrollback (where the terminal
    // supports it), H homes the cursor
    const CLEAR: &str = "\x1b[2J\x1b[3J\x1b[H";

    write!(
        tty_out,
        "{}keyquorum key ceremony — {} shares, threshold {}\n\n\
         Each share is shown alone and the screen is cleared between shares.\n\
         Have each holder step up in turn.\n\n\
         Press Enter to show share 1.\n",
        CLEAR,
        shares.len(),
        cli.threshold
    )?;
    tty_out.flush()?;
    let mut line = String::new();
    tty_in.read_line(&mut line)?;

    for (i, share) in shares.iter().enumerate() {
        let mut sharks_data: Vec<u8> = Vec::from(share);
        let index = sharks_data[0];
        let opts = make_format_opts(cli, encoding, (i + 1) as u8);
        let mut formatted = share_format::format_share(&sharks_data, &opts);

        let mut screen = format!(
            "{}Share {} of {} (index {})\n\n{}\n\n\
             When the holder has recorded this share, press Enter to clear the screen.\n",
            CLEAR,
            i + 1,
            shares.len(),
            index,
            formatted
        );

        // Write, wait for confirmation, then wipe our copies — error
        // handling deferred so the plaintext is zeroized on every path
        let write_result = tty_out
            .write_all(screen.as_bytes())
            .and_then(|_| tty_out.flush());
        let read_result = if write_result.is_ok() {
            line.clear();
            tty_in.read_line(&mut line).map(|_| ())
        } else {
            Ok(())
        };

        sharks_data.zeroize();
        formatted.zeroize();
        screen.zeroize();

        write_result?;
        read_result?;

        write!(tty_out, "{}", CLEAR)?;
        tty_out.flush()?;
    }

    writeln!(
        tty_out,
        "{}Ceremony complete: {} shares shown, screen cleared.",
        CLEAR,
        shares.len()
    )?;
    tty_out.flush()?;
    eprintln!("{} shares shown interactively.", shares.len());

    Ok(())
}

fn parse_recipients_file(path: &Path) -> Result<Vec<age::x25519::Recipient>> {
    let contents = std::fs::read_to_string(path)
        .map_err(|e| anyhow::anyhow!("failed to read recipients file {}: {}", path.display(), e))?;

    let mut recipients = Vec::new();
    for (line_num, line) in contents.lines().enumerate() {
        let trimmed = line.trim();
        if trimmed.is_empty() || trimmed.starts_with('#') {
            continue;
        }
        let recipient: age::x25519::Recipient = trimmed.parse().map_err(|e| {
            anyhow::anyhow!(
                "invalid age recipient on line {} of {}: {} ({})",
                line_num + 1,
                path.display(),
                trimmed,
                e
            )
        })?;
        recipients.push(recipient);
    }

    if recipients.is_empty() {
        bail!("no valid age recipients found in {}", path.display());
    }

    Ok(recipients)
}

fn output_age(
    shares: &[blahaj::Share],
    cli: &Cli,
    encoding: ShareEncoding,
    dir: &PathBuf,
    recipients: &[age::x25519::Recipient],
) -> Result<()> {
    std::fs::create_dir_all(dir).map_err(|e| {
        anyhow::anyhow!("failed to create output directory {}: {}", dir.display(), e)
    })?;

    for (i, share) in shares.iter().enumerate() {
        let mut sharks_data: Vec<u8> = Vec::from(share);
        let index = sharks_data[0];
        let opts = make_format_opts(cli, encoding, (i + 1) as u8);
        let mut formatted = share_format::format_share(&sharks_data, &opts);

        if i < recipients.len() {
            // Encrypt to the corresponding recipient.
            // Note: age's internal plaintext buffers are outside our zeroize scope,
            // same boundary as blahaj internals.
            let filename = if cli.armor {
                dir.join(format!("share-{}.txt.age.txt", i + 1))
            } else {
                dir.join(format!("share-{}.txt.age", i + 1))
            };

            // Encrypt first, zeroize the plaintext, and only then handle
            // errors — a `?` before the zeroize calls would drop the
            // plaintext share unwiped on the error path
            let encrypt_result = if cli.armor {
                age::encrypt_and_armor(&recipients[i], formatted.as_bytes()).map(String::into_bytes)
            } else {
                age::encrypt(&recipients[i], formatted.as_bytes())
            };

            sharks_data.zeroize();
            formatted.zeroize();

            let encrypted = encrypt_result
                .map_err(|e| anyhow::anyhow!("failed to encrypt share {}: {}", i + 1, e))?;
            std::fs::write(&filename, &encrypted)
                .map_err(|e| anyhow::anyhow!("failed to write {}: {}", filename.display(), e))?;

            eprintln!(
                "Share {} (index {}) encrypted and written to {}{}",
                i + 1,
                index,
                filename.display(),
                duress_suffix(cli, i)
            );
        } else {
            // No recipient for this share — write unencrypted (warning already emitted)
            let filename = dir.join(format!("share-{}.txt", i + 1));
            let write_result = std::fs::write(&filename, format!("{}\n", formatted));

            sharks_data.zeroize();
            formatted.zeroize();

            write_result
                .map_err(|e| anyhow::anyhow!("failed to write {}: {}", filename.display(), e))?;

            eprintln!(
                "Share {} (index {}) written UNENCRYPTED to {}{}",
                i + 1,
                index,
                filename.display(),
                duress_suffix(cli, i)
            );
        }
    }

    let encrypted_count = std::cmp::min(recipients.len(), shares.len());
    let unencrypted_count = shares.len() - encrypted_count;
    eprintln!(
        "\n{} share files written to {} ({} encrypted, {} unencrypted)",
        shares.len(),
        dir.display(),
        encrypted_count,
        unencrypted_count,
    );
    if unencrypted_count == 0 {
        eprintln!("Distribute each .age file to its holder.");
        eprintln!("Recipients decrypt with: age -d -i identity.txt share-N.txt.age");
        eprintln!("(age CLI: apt install age / brew install age — see https://age-encryption.org)");
    } else {
        eprintln!("Distribute each file to its holder, then delete this directory.");
        eprintln!(
            "Encrypted shares require the age CLI to decrypt \
             (apt install age / brew install age — see https://age-encryption.org)"
        );
    }

    Ok(())
}

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

    fn generate_test_pubkeys(n: usize) -> Vec<String> {
        (0..n)
            .map(|_| {
                let identity = age::x25519::Identity::generate();
                identity.to_public().to_string()
            })
            .collect()
    }

    fn write_temp_file(contents: &str) -> tempfile::NamedTempFile {
        let mut f = tempfile::NamedTempFile::new().unwrap();
        f.write_all(contents.as_bytes()).unwrap();
        f
    }

    #[test]
    fn parse_recipients_valid() {
        let keys = generate_test_pubkeys(3);
        let contents = format!(
            "# Alice\n{}\n\n# Bob\n{}\n# Carol\n{}\n",
            keys[0], keys[1], keys[2]
        );
        let f = write_temp_file(&contents);
        let recipients = parse_recipients_file(f.path()).unwrap();
        assert_eq!(recipients.len(), 3);
    }

    #[test]
    fn parse_recipients_empty() {
        let f = write_temp_file("# only comments\n\n# another comment\n");
        let err = parse_recipients_file(f.path()).unwrap_err();
        assert!(
            err.to_string().contains("no valid age recipients"),
            "unexpected error: {}",
            err
        );
    }

    #[test]
    fn parse_recipients_invalid_key() {
        let keys = generate_test_pubkeys(1);
        let contents = format!("{}\nnot-a-valid-key\n", keys[0]);
        let f = write_temp_file(&contents);
        let err = parse_recipients_file(f.path()).unwrap_err();
        assert!(
            err.to_string().contains("line 2"),
            "expected line number in error: {}",
            err
        );
        assert!(
            err.to_string().contains("not-a-valid-key"),
            "expected offending content in error: {}",
            err
        );
    }

    #[test]
    fn parse_recipients_whitespace() {
        let keys = generate_test_pubkeys(2);
        let contents = format!("  {}  \n\t{}\t\n", keys[0], keys[1]);
        let f = write_temp_file(&contents);
        let recipients = parse_recipients_file(f.path()).unwrap();
        assert_eq!(recipients.len(), 2);
    }
}
