use std::fs::{self, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};

use anyhow::{anyhow, bail, Context, Result};
use base64::engine::general_purpose::STANDARD;
use base64::Engine as _;
use clap::{
    builder::styling::{AnsiColor, Color, Style, Styles},
    Args, Parser, Subcommand, ValueHint,
};

use hdpki::{
    encode_private_key, format_authorized_key, init_seed, last_mnemonic, master,
    Error as HdpkiError, X509Params,
};

#[derive(Parser)]
#[command(
    name = "hdpki",
    about = "Hierarchical deterministic key derivation for SSH and X.509",
    version,
    styles = cli_styles()
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Generate a new master seed
    Master(MasterArgs),
    /// Derive an OpenSSH key from a stored master seed
    Derive(DeriveArgs),
    /// Generate an X.509 certificate from a derived key
    Cert(CertArgs),
}

#[derive(Args)]
struct MasterArgs {
    /// Optional output file for the encoded master seed (base64)
    #[arg(short = 'o', long, value_name = "FILE", value_hint = ValueHint::FilePath)]
    output: Option<PathBuf>,
    /// Generate a BIP-39 mnemonic and print it to stdout
    #[arg(short = 'b', long)]
    bip39: bool,
    /// Overwrite existing files when writing seeds
    #[arg(short = 'f', long)]
    force: bool,
}

#[derive(Args)]
struct DeriveArgs {
    /// Path to the stored master seed
    #[arg(short = 'm', long, value_name = "FILE", value_hint = ValueHint::FilePath)]
    master: PathBuf,
    /// Derivation path (m/ssh/...)
    #[arg(short = 'p', long)]
    path: String,
    /// Optional output file for the OpenSSH private key
    #[arg(short = 'o', long, value_name = "FILE", value_hint = ValueHint::FilePath)]
    output: Option<PathBuf>,
    /// Optional output file for the authorized key line
    #[arg(
        short = 'A',
        long = "authorized-output",
        value_name = "FILE",
        value_hint = ValueHint::FilePath
    )]
    authorized_output: Option<PathBuf>,
    /// Append the public key line to stdout
    #[arg(short = 'a', long)]
    print_authorized: bool,
    /// Print the fingerprint to stdout (defaults to stderr when writing files)
    #[arg(short = 'F', long)]
    print_fingerprint: bool,
    /// Override the default comment appended to keys (defaults to the canonical path)
    #[arg(short = 'c', long)]
    comment: Option<String>,
    /// Overwrite existing files when writing derived material
    #[arg(short = 'f', long)]
    force: bool,
    /// Skip strict path validation and derive using the literal path string
    #[arg(short = 'r', long)]
    raw_path: bool,
}

#[derive(Args)]
struct CertArgs {
    /// Path to the stored master seed
    #[arg(short = 'm', long, value_name = "FILE", value_hint = ValueHint::FilePath)]
    master: PathBuf,
    /// Derivation path (m/ssh/...)
    #[arg(short = 'p', long)]
    path: String,
    /// Common Name (CN) for the certificate subject
    #[arg(long)]
    cn: String,
    /// Organization (O) for the certificate subject
    #[arg(long)]
    org: Option<String>,
    /// Country (C) for the certificate subject (two-letter ISO code)
    #[arg(long)]
    country: Option<String>,
    /// Certificate validity in days (default: 365)
    #[arg(long, default_value = "365")]
    days: u32,
    /// Generate a CA certificate
    #[arg(long)]
    ca: bool,
    /// DNS Subject Alternative Names (can be repeated)
    #[arg(long = "san-dns", value_name = "DOMAIN")]
    san_dns: Vec<String>,
    /// Email Subject Alternative Names (can be repeated)
    #[arg(long = "san-email", value_name = "EMAIL")]
    san_email: Vec<String>,
    /// Output file for the PEM certificate
    #[arg(short = 'o', long, value_name = "FILE", value_hint = ValueHint::FilePath)]
    output: Option<PathBuf>,
    /// Output file for the PEM private key
    #[arg(short = 'k', long = "key-output", value_name = "FILE", value_hint = ValueHint::FilePath)]
    key_output: Option<PathBuf>,
    /// Sign with a CA node (derivation path for the issuer)
    #[arg(long = "issuer-path")]
    issuer_path: Option<String>,
    /// PEM file of the issuer CA certificate (required with --issuer-path)
    #[arg(long = "issuer-cert", value_name = "FILE", value_hint = ValueHint::FilePath)]
    issuer_cert: Option<PathBuf>,
    /// Overwrite existing files
    #[arg(short = 'f', long)]
    force: bool,
    /// Skip strict path validation
    #[arg(short = 'r', long)]
    raw_path: bool,
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    match cli.command {
        Command::Master(args) => handle_master(args),
        Command::Derive(args) => handle_derive(args),
        Command::Cert(args) => handle_cert(args),
    }
}

fn handle_master(args: MasterArgs) -> Result<()> {
    let seed = init_seed(args.bip39)?;
    let mnemonic = last_mnemonic();
    let encoded = STANDARD.encode(&seed);

    if let Some(path) = args.output {
        write_secure_file(&path, &encoded, args.force, 0o600)?;
        println!(
            "stored {} bytes of master seed in {}",
            seed.len(),
            path.display()
        );
    } else {
        println!("{encoded}");
    }

    if let Some(mnemonic) = mnemonic {
        println!("mnemonic:\n{mnemonic}");
    }

    Ok(())
}

fn handle_derive(args: DeriveArgs) -> Result<()> {
    let seed = read_seed(&args.master)?;
    let root = master(&seed)?;
    let node = derive_node(&root, &args.path, args.raw_path)?;
    let material = node.ed25519()?;
    let comment = args.comment.unwrap_or_else(|| node.path().to_string());
    let private_key = encode_private_key(&material, &comment);
    let authorized_line = format_authorized_key(&material, Some(&comment));
    let fingerprint = node.fingerprint()?;

    match &args.output {
        Some(path) => {
            write_secure_file(path, &private_key, args.force, 0o600)?;
            eprintln!("wrote private key to {}", path.display());
        }
        None => {
            print!("{}", private_key);
        }
    }

    if let Some(path) = &args.authorized_output {
        write_secure_file(path, &authorized_line, args.force, 0o644)?;
        eprintln!("wrote authorized key to {}", path.display());
    }

    if args.print_authorized && args.authorized_output.is_none() {
        print!("{}", authorized_line);
    }

    if args.print_fingerprint || args.output.is_none() {
        println!("{fingerprint}");
    } else {
        eprintln!("fingerprint: {fingerprint}");
    }

    Ok(())
}

fn handle_cert(args: CertArgs) -> Result<()> {
    let seed = read_seed(&args.master)?;
    let root = master(&seed)?;
    let node = derive_node(&root, &args.path, args.raw_path)?;

    let params = X509Params {
        common_name: args.cn,
        organization: args.org,
        country: args.country,
        validity_days: args.days,
        is_ca: args.ca,
        san_dns_names: args.san_dns,
        san_emails: args.san_email,
    };

    let cert = match (&args.issuer_path, &args.issuer_cert) {
        (Some(issuer_path), Some(issuer_cert_path)) => {
            let issuer_node = derive_node(&root, issuer_path, args.raw_path)?;
            let issuer_pem = fs::read_to_string(issuer_cert_path)
                .with_context(|| format!("reading issuer cert {issuer_cert_path:?}"))?;
            node.x509_signed_by(&issuer_node, &issuer_pem, &params)?
        }
        (Some(_), None) => bail!("--issuer-cert is required when using --issuer-path"),
        (None, Some(_)) => bail!("--issuer-path is required when using --issuer-cert"),
        (None, None) => node.x509_self_signed(&params)?,
    };

    match &args.output {
        Some(path) => {
            write_secure_file(path, &cert.cert_pem, args.force, 0o644)?;
            eprintln!("wrote certificate to {}", path.display());
        }
        None => {
            print!("{}", cert.cert_pem);
        }
    }

    match &args.key_output {
        Some(path) => {
            write_secure_file(path, &cert.key_pem, args.force, 0o600)?;
            eprintln!("wrote private key to {}", path.display());
        }
        None if args.output.is_some() => {
            // Only print key to stdout if cert was written to a file
            print!("{}", cert.key_pem);
        }
        None => {}
    }

    Ok(())
}

fn derive_node(
    root: &hdpki::Node,
    path: &str,
    raw_path: bool,
) -> Result<hdpki::Node> {
    match root.derive(path) {
        Ok(node) => Ok(node),
        Err(HdpkiError::Path(err)) if raw_path => {
            eprintln!("warning: {err}");
            eprintln!(
                "warning: deriving using raw path {path}; result may not follow hdpki conventions",
            );
            Ok(root.derive_raw(path)?)
        }
        Err(err) => Err(err.into()),
    }
}

fn read_seed(path: &Path) -> Result<Vec<u8>> {
    let mut file = fs::File::open(path).with_context(|| format!("opening {path:?}"))?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .with_context(|| format!("reading {path:?}"))?;

    let filtered: String = contents
        .lines()
        .map(str::trim)
        .filter(|line| !line.is_empty() && !line.starts_with('#'))
        .collect();

    if filtered.is_empty() {
        bail!("seed file {path:?} is empty");
    }

    STANDARD
        .decode(filtered.as_bytes())
        .or_else(|_| hex::decode(filtered.as_bytes()))
        .map_err(|_| anyhow!("failed to decode seed: expected base64 or hex data"))
}

fn write_secure_file(path: &Path, contents: &str, force: bool, mode: u32) -> Result<()> {
    if path.exists() && !force {
        bail!("refusing to overwrite existing file: {}", path.display());
    }
    if let Some(parent) = path.parent() {
        if !parent.as_os_str().is_empty() {
            fs::create_dir_all(parent).with_context(|| format!("creating directory {parent:?}"))?;
        }
    }

    let mut options = OpenOptions::new();
    options.write(true).create(true).truncate(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        options.mode(mode);
    }

    let mut file = options
        .open(path)
        .with_context(|| format!("opening {path:?} for writing"))?;
    file.write_all(contents.as_bytes())
        .with_context(|| format!("writing to {path:?}"))?;
    file.flush().with_context(|| format!("flushing {path:?}"))?;
    Ok(())
}

fn cli_styles() -> Styles {
    Styles::styled()
        .header(Style::new().bold())
        .usage(Style::new().bold())
        .literal(Style::new().bold())
        .error(
            Style::new()
                .bold()
                .fg_color(Some(Color::Ansi(AnsiColor::Red))),
        )
}
