//! soliton — CLI for post-quantum cryptographic operations.

use clap::{Parser, Subcommand};
use std::fs;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::process;
use zeroize::Zeroize;

const STREAM_CHUNK_SIZE: usize = 1_048_576; // 1 MiB

#[derive(Parser)]
#[command(
    name = "soliton",
    version,
    about = "Post-quantum cryptographic toolkit"
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Generate an identity keypair (X-Wing + Ed25519 + ML-DSA-65).
    Keygen {
        /// Output directory (default: current directory).
        #[arg(short, long, default_value = ".")]
        output: PathBuf,
    },

    /// Print the SHA3-256 fingerprint of a public key.
    Fingerprint {
        /// Path to public key file.
        pk: PathBuf,
    },

    /// Hybrid sign a file or stdin.
    Sign {
        /// Path to secret key file.
        sk: PathBuf,
        /// File to sign (reads stdin if omitted).
        file: Option<PathBuf>,
        /// Output signature file (default: <file>.sig or stdout).
        #[arg(short, long)]
        output: Option<PathBuf>,
    },

    /// Verify a hybrid signature.
    Verify {
        /// Path to public key file.
        pk: PathBuf,
        /// File that was signed.
        file: PathBuf,
        /// Signature file (default: <file>.sig).
        #[arg(short, long)]
        sig: Option<PathBuf>,
    },

    /// Generate an X-Wing keypair (for SPK/OPK pre-keys).
    XwingKeygen {
        /// Output directory (default: current directory).
        #[arg(short, long, default_value = ".")]
        output: PathBuf,
    },

    /// Sign a pre-key with an identity key.
    SignPrekey {
        /// Path to identity secret key.
        sk: PathBuf,
        /// Path to X-Wing pre-key public key.
        spk_pub: PathBuf,
        /// Output signature file.
        #[arg(short, long)]
        output: Option<PathBuf>,
    },

    /// Generate a verification phrase from two public keys.
    Phrase {
        /// Path to first public key.
        pk_a: PathBuf,
        /// Path to second public key.
        pk_b: PathBuf,
    },

    /// Encrypt a file with streaming AEAD (XChaCha20-Poly1305).
    Encrypt {
        /// 32-byte key file.
        #[arg(short, long, group = "key_source")]
        key: Option<PathBuf>,
        /// Derive key from passphrase via Argon2id. Requires --salt.
        #[arg(long, group = "key_source")]
        derive: bool,
        /// Hex-encoded salt for --derive (16 bytes = 32 hex chars).
        #[arg(long, requires = "derive")]
        salt: Option<String>,
        /// Input file (reads stdin if omitted).
        file: Option<PathBuf>,
        /// Output file (writes stdout if omitted).
        #[arg(short, long)]
        output: Option<PathBuf>,
    },

    /// Decrypt a file with streaming AEAD.
    Decrypt {
        /// 32-byte key file.
        #[arg(short, long, group = "key_source")]
        key: Option<PathBuf>,
        /// Derive key from passphrase via Argon2id. Requires --salt.
        #[arg(long, group = "key_source", requires = "salt")]
        derive: bool,
        /// Hex-encoded salt for --derive (16 bytes = 32 hex chars).
        #[arg(long, requires = "derive")]
        salt: Option<String>,
        /// Input file (reads stdin if omitted).
        file: Option<PathBuf>,
        /// Output file (writes stdout if omitted).
        #[arg(short, long)]
        output: Option<PathBuf>,
    },

    /// Derive key material from a passphrase via Argon2id.
    Argon2id {
        /// Memory cost in KiB (default: 65536 = 64 MiB).
        #[arg(short, long, default_value = "65536")]
        m_cost: u32,
        /// Time cost / passes (default: 3).
        #[arg(short, long, default_value = "3")]
        t_cost: u32,
        /// Parallelism / lanes (default: 4).
        #[arg(short, long, default_value = "4")]
        p_cost: u32,
        /// Output length in bytes (default: 32).
        #[arg(short, long, default_value = "32")]
        length: usize,
    },

    /// Print the library version.
    Version,
}

fn main() {
    let cli = Cli::parse();
    if let Err(e) = run(cli) {
        eprintln!("error: {e}");
        process::exit(1);
    }
}

fn run(cli: Cli) -> Result<(), String> {
    match cli.command {
        Command::Keygen { output } => cmd_keygen(&output),
        Command::Fingerprint { pk } => cmd_fingerprint(&pk),
        Command::Sign { sk, file, output } => cmd_sign(&sk, file.as_deref(), output.as_deref()),
        Command::Verify { pk, file, sig } => cmd_verify(&pk, &file, sig.as_deref()),
        Command::XwingKeygen { output } => cmd_xwing_keygen(&output),
        Command::SignPrekey {
            sk,
            spk_pub,
            output,
        } => cmd_sign_prekey(&sk, &spk_pub, output.as_deref()),
        Command::Phrase { pk_a, pk_b } => cmd_phrase(&pk_a, &pk_b),
        Command::Encrypt {
            key,
            derive,
            salt,
            file,
            output,
        } => cmd_encrypt(
            key.as_deref(),
            derive,
            salt.as_deref(),
            file.as_deref(),
            output.as_deref(),
        ),
        Command::Decrypt {
            key,
            derive,
            salt,
            file,
            output,
        } => cmd_decrypt(
            key.as_deref(),
            derive,
            salt.as_deref(),
            file.as_deref(),
            output.as_deref(),
        ),
        Command::Argon2id {
            m_cost,
            t_cost,
            p_cost,
            length,
        } => cmd_argon2id(m_cost, t_cost, p_cost, length),
        Command::Version => {
            println!("soliton {}", soliton::VERSION);
            Ok(())
        }
    }
}

fn cmd_keygen(output: &Path) -> Result<(), String> {
    let id = soliton::identity::generate_identity().map_err(|e| e.to_string())?;
    let pk_path = output.join("identity.pk");
    let sk_path = output.join("identity.sk");
    fs::write(&pk_path, id.public_key.as_bytes()).map_err(|e| e.to_string())?;
    write_secret_file(&sk_path, id.secret_key.as_bytes())?;
    let fp = soliton::primitives::sha3_256::fingerprint_hex(id.public_key.as_bytes());
    eprintln!("Public key:  {}", pk_path.display());
    eprintln!("Secret key:  {}", sk_path.display());
    eprintln!("Fingerprint: {fp}");
    Ok(())
}

fn cmd_fingerprint(pk_path: &Path) -> Result<(), String> {
    let pk_bytes = fs::read(pk_path).map_err(|e| e.to_string())?;
    let fp = soliton::primitives::sha3_256::fingerprint_hex(&pk_bytes);
    println!("{fp}");
    Ok(())
}

fn cmd_sign(sk_path: &Path, file: Option<&Path>, output: Option<&Path>) -> Result<(), String> {
    let sk_bytes = fs::read(sk_path).map_err(|e| e.to_string())?;
    let sk =
        soliton::identity::IdentitySecretKey::from_bytes(sk_bytes).map_err(|e| e.to_string())?;
    let message = read_input(file)?;
    let sig = soliton::identity::hybrid_sign(&sk, &message).map_err(|e| e.to_string())?;

    match output {
        Some(path) => fs::write(path, sig.as_bytes()).map_err(|e| e.to_string())?,
        None => match file {
            Some(f) => {
                let sig_path = format!("{}.sig", f.display());
                fs::write(&sig_path, sig.as_bytes()).map_err(|e| e.to_string())?;
                eprintln!("Signature: {sig_path}");
            }
            None => {
                io::stdout()
                    .write_all(sig.as_bytes())
                    .map_err(|e| e.to_string())?;
            }
        },
    }
    Ok(())
}

fn cmd_verify(pk_path: &Path, file: &Path, sig_path: Option<&Path>) -> Result<(), String> {
    let pk_bytes = fs::read(pk_path).map_err(|e| e.to_string())?;
    let pk =
        soliton::identity::IdentityPublicKey::from_bytes(pk_bytes).map_err(|e| e.to_string())?;
    let message = fs::read(file).map_err(|e| e.to_string())?;
    let sig_file = sig_path
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from(format!("{}.sig", file.display())));
    let sig_bytes = fs::read(&sig_file).map_err(|e| e.to_string())?;
    let sig =
        soliton::identity::HybridSignature::from_bytes(sig_bytes).map_err(|e| e.to_string())?;

    soliton::identity::hybrid_verify(&pk, &message, &sig).map_err(|e| e.to_string())?;
    eprintln!("Signature OK");
    Ok(())
}

fn cmd_xwing_keygen(output: &Path) -> Result<(), String> {
    let (pk, sk) = soliton::primitives::xwing::keygen().map_err(|e| e.to_string())?;
    let pk_path = output.join("xwing.pk");
    let sk_path = output.join("xwing.sk");
    fs::write(&pk_path, pk.as_bytes()).map_err(|e| e.to_string())?;
    write_secret_file(&sk_path, sk.as_bytes())?;
    eprintln!("Public key: {}", pk_path.display());
    eprintln!("Secret key: {}", sk_path.display());
    Ok(())
}

fn cmd_sign_prekey(
    sk_path: &Path,
    spk_pub_path: &Path,
    output: Option<&Path>,
) -> Result<(), String> {
    let sk_bytes = fs::read(sk_path).map_err(|e| e.to_string())?;
    let sk =
        soliton::identity::IdentitySecretKey::from_bytes(sk_bytes).map_err(|e| e.to_string())?;
    let spk_bytes = fs::read(spk_pub_path).map_err(|e| e.to_string())?;
    let spk =
        soliton::primitives::xwing::PublicKey::from_bytes(spk_bytes).map_err(|e| e.to_string())?;
    let sig = soliton::kex::sign_prekey(&sk, &spk).map_err(|e| e.to_string())?;

    let out = output
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("spk.sig"));
    fs::write(&out, sig.as_bytes()).map_err(|e| e.to_string())?;
    eprintln!("Pre-key signature: {}", out.display());
    Ok(())
}

fn cmd_phrase(pk_a_path: &Path, pk_b_path: &Path) -> Result<(), String> {
    let pk_a = fs::read(pk_a_path).map_err(|e| e.to_string())?;
    let pk_b = fs::read(pk_b_path).map_err(|e| e.to_string())?;
    let phrase =
        soliton::verification::verification_phrase(&pk_a, &pk_b).map_err(|e| e.to_string())?;
    println!("{phrase}");
    Ok(())
}

fn cmd_encrypt(
    key_path: Option<&Path>,
    derive: bool,
    salt_hex: Option<&str>,
    file: Option<&Path>,
    output: Option<&Path>,
) -> Result<(), String> {
    let mut key = load_or_derive_key(key_path, derive, salt_hex)?;
    let key_arr: &[u8; 32] = key
        .as_slice()
        .try_into()
        .map_err(|_| "key must be 32 bytes".to_string())?;
    let plaintext = read_input(file)?;

    let mut enc =
        soliton::streaming::stream_encrypt_init(key_arr, &[], false).map_err(|e| e.to_string())?;
    let header = enc.header();

    let mut out_data = Vec::with_capacity(header.len() + plaintext.len() + 16);
    out_data.extend_from_slice(&header);

    // Chunk the plaintext into STREAM_CHUNK_SIZE pieces.
    if plaintext.is_empty() {
        let chunk = enc.encrypt_chunk(&[], true).map_err(|e| e.to_string())?;
        out_data.extend_from_slice(&chunk);
    } else {
        let chunks: Vec<&[u8]> = plaintext.chunks(STREAM_CHUNK_SIZE).collect();
        let last_idx = chunks.len() - 1;
        for (i, chunk_data) in chunks.iter().enumerate() {
            let is_last = i == last_idx;
            let ct = enc
                .encrypt_chunk(chunk_data, is_last)
                .map_err(|e| e.to_string())?;
            out_data.extend_from_slice(&ct);
        }
    }

    write_output(output, &out_data)?;
    key.zeroize();
    Ok(())
}

fn cmd_decrypt(
    key_path: Option<&Path>,
    derive: bool,
    salt_hex: Option<&str>,
    file: Option<&Path>,
    output: Option<&Path>,
) -> Result<(), String> {
    let mut key = load_or_derive_key(key_path, derive, salt_hex)?;
    let key_arr: &[u8; 32] = key
        .as_slice()
        .try_into()
        .map_err(|_| "key must be 32 bytes".to_string())?;
    let data = read_input(file)?;

    if data.len() < 26 {
        return Err("input too short (missing stream header)".to_string());
    }
    let header: &[u8; 26] = data[..26].try_into().unwrap();
    let ciphertext = &data[26..];

    let mut dec =
        soliton::streaming::stream_decrypt_init(key_arr, header, &[]).map_err(|e| e.to_string())?;

    let mut plaintext = Vec::new();
    // Each encrypted chunk is: tag_byte (1) + ciphertext (variable) + AEAD tag (16).
    // Non-final uncompressed: 1 + STREAM_CHUNK_SIZE + 16 = 1,048,593 bytes.
    // Final: 1 + plaintext_len + 16 bytes.
    // We must pass exactly one chunk per decrypt_chunk call.
    let mut offset = 0;
    loop {
        if offset >= ciphertext.len() {
            break;
        }
        // Determine this chunk's size from the tag byte and position.
        let tag_byte = ciphertext[offset];
        let is_final_chunk = tag_byte == 0x01;
        let chunk_size = if is_final_chunk {
            // Final chunk: everything remaining.
            ciphertext.len() - offset
        } else {
            // Non-final: 1 (tag) + STREAM_CHUNK_SIZE + 16 (AEAD tag).
            1 + STREAM_CHUNK_SIZE + 16
        };
        if offset + chunk_size > ciphertext.len() {
            return Err("truncated chunk in stream".to_string());
        }
        let (pt, is_last) = dec
            .decrypt_chunk(&ciphertext[offset..offset + chunk_size])
            .map_err(|e| e.to_string())?;
        plaintext.extend_from_slice(&pt);
        offset += chunk_size;
        if is_last {
            break;
        }
    }

    // Truncation detection: an attacker who strips the final chunk leaves
    // the decryptor in a non-finalized state with all prior chunks
    // authentically decrypted. Without this check, truncated files would
    // be silently accepted.
    if !dec.is_finalized() {
        return Err("stream truncated: final chunk missing".to_string());
    }

    write_output(output, &plaintext)?;
    key.zeroize();
    Ok(())
}

fn cmd_argon2id(m_cost: u32, t_cost: u32, p_cost: u32, length: usize) -> Result<(), String> {
    let mut password = read_passphrase("Passphrase: ")?;
    let mut salt = [0u8; 16];
    soliton::primitives::random::random_bytes(&mut salt);

    let params = soliton::primitives::argon2::Argon2Params {
        m_cost,
        t_cost,
        p_cost,
    };
    let mut out = vec![0u8; length];
    soliton::primitives::argon2::argon2id(password.as_bytes(), &salt, params, &mut out)
        .map_err(|e| e.to_string())?;

    eprintln!("Salt: {}", hex::encode(salt));
    println!("{}", hex::encode(&out));
    out.zeroize();
    password.zeroize();
    Ok(())
}

// ── Helpers ─────────────────────────────────────────────────────────────

fn read_input(file: Option<&Path>) -> Result<Vec<u8>, String> {
    match file {
        Some(path) => fs::read(path).map_err(|e| e.to_string()),
        None => {
            let mut buf = Vec::new();
            io::stdin()
                .read_to_end(&mut buf)
                .map_err(|e| e.to_string())?;
            Ok(buf)
        }
    }
}

fn write_output(output: Option<&Path>, data: &[u8]) -> Result<(), String> {
    match output {
        Some(path) => fs::write(path, data).map_err(|e| e.to_string()),
        None => io::stdout().write_all(data).map_err(|e| e.to_string()),
    }
}

/// Write a file with mode 0o600 (owner-only read/write).
/// Uses OpenOptions on Unix to set permissions atomically (no TOCTOU window).
fn write_secret_file(path: &Path, data: &[u8]) -> Result<(), String> {
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .mode(0o600)
            .open(path)
            .map_err(|e| e.to_string())?;
        file.write_all(data).map_err(|e| e.to_string())?;
        Ok(())
    }
    #[cfg(not(unix))]
    {
        fs::write(path, data).map_err(|e| e.to_string())
    }
}

fn load_or_derive_key(
    key_path: Option<&Path>,
    derive: bool,
    salt_hex: Option<&str>,
) -> Result<Vec<u8>, String> {
    if derive {
        let mut password = read_passphrase("Passphrase: ")?;
        let salt = match salt_hex {
            Some(hex_str) => hex::decode(hex_str).map_err(|_| "invalid hex salt".to_string())?,
            None => {
                // Generate new salt for encrypt; print it so the user can save it.
                let mut s = [0u8; 16];
                soliton::primitives::random::random_bytes(&mut s);
                eprintln!("Salt: {} (save this to decrypt later)", hex::encode(s));
                s.to_vec()
            }
        };
        let params = soliton::primitives::argon2::Argon2Params {
            m_cost: 65536,
            t_cost: 3,
            p_cost: 4,
        };
        let mut key = vec![0u8; 32];
        soliton::primitives::argon2::argon2id(password.as_bytes(), &salt, params, &mut key)
            .map_err(|e| e.to_string())?;
        password.zeroize();
        Ok(key)
    } else {
        let path = key_path.ok_or("provide --key <file> or --derive")?;
        fs::read(path).map_err(|e| e.to_string())
    }
}

/// Read a passphrase from the terminal without echoing.
fn read_passphrase(prompt: &str) -> Result<String, String> {
    eprint!("{prompt}");
    io::stderr().flush().map_err(|e| e.to_string())?;

    // Disable echo on Unix terminals.
    #[cfg(unix)]
    {
        use std::os::unix::io::AsRawFd;
        let fd = io::stdin().as_raw_fd();
        let mut termios = unsafe {
            let mut t = std::mem::zeroed();
            if libc::tcgetattr(fd, &mut t) != 0 {
                // Not a terminal — fall back to plain read.
                let mut line = String::new();
                io::stdin()
                    .read_line(&mut line)
                    .map_err(|e| e.to_string())?;
                let result = line.trim_end().to_string();
                line.zeroize();
                return Ok(result);
            }
            t
        };
        let orig = termios;
        termios.c_lflag &= !libc::ECHO;
        unsafe { libc::tcsetattr(fd, libc::TCSANOW, &termios) };
        let mut line = String::new();
        let result = io::stdin().read_line(&mut line);
        unsafe { libc::tcsetattr(fd, libc::TCSANOW, &orig) };
        eprintln!(); // newline after hidden input
        result.map_err(|e| e.to_string())?;
        let trimmed = line.trim_end().to_string();
        line.zeroize();
        Ok(trimmed)
    }

    #[cfg(not(unix))]
    {
        let mut line = String::new();
        io::stdin()
            .read_line(&mut line)
            .map_err(|e| e.to_string())?;
        let trimmed = line.trim_end().to_string();
        line.zeroize();
        Ok(trimmed)
    }
}