mod keys;

use std::{fs::File, io::BufReader, path::Path};

use anyhow::{anyhow, Result};
use clap::{Args, Parser, Subcommand};
use keys::{Bink, Keys};
use serde_json::{from_reader, from_str};

use umskt::{
    bink1998, bink2002, confid,
    crypto::{EllipticCurve, PrivateKey},
};

#[derive(Parser, Debug)]
#[command(author, about, version, long_about = None)]
struct Cli {
    /// Enable verbose output
    #[arg(short, long)]
    verbose: bool,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand, Clone, Debug)]
enum Commands {
    /// Show which products/binks can be loaded
    List(ListArgs),
    /// Generate new product keys
    Generate(GenerateArgs),
    /// Validate a product key
    Validate(ValidateArgs),
    /// Generate a phone activation Confirmation ID from an Installation ID
    Confid(ConfirmationIdArgs),
}

#[derive(Args, Clone, Debug)]
struct ListArgs {
    /// Specify which keys file to load
    #[arg(short = 'f', long = "file")]
    keys_path: Option<String>,
}

#[derive(Args, Clone, Debug)]
struct GenerateArgs {
    /// Specify which BINK identifier to load
    #[arg(short, long, default_value = "2E")]
    binkid: String,

    /// Specify which Channel Identifier to use
    #[arg(short = 'c', long = "channel", default_value = "640")]
    channel_id: u32,

    /// Number of keys to generate
    #[arg(short = 'n', long = "number", default_value = "1")]
    num_keys: u64,

    /// Specify which keys file to load
    #[arg(short = 'f', long = "file")]
    keys_path: Option<String>,
}

#[derive(Args, Clone, Debug)]
struct ValidateArgs {
    /// Specify which BINK identifier to load
    #[arg(short, long, default_value = "2E")]
    binkid: String,

    /// Specify which keys file to load
    #[arg(short = 'f', long = "file")]
    keys_path: Option<String>,

    /// Product key to validate signature
    key_to_check: String,
}

#[derive(Args, Clone, Debug)]
struct ConfirmationIdArgs {
    /// Installation ID used to generate confirmation ID
    instid: String,
}

fn main() -> Result<()> {
    let args = Cli::parse();

    match &args.command {
        Commands::List(list_args) => {
            let keys = load_keys(list_args.keys_path.as_ref(), args.verbose)?;
            for (key, value) in keys.products.iter() {
                println!("{}: {:?}", key, value.bink);
            }

            println!("\n\n** Please note: any BINK ID other than 2E is considered experimental at this time **\n");
        }
        Commands::Generate(generate_args) => {
            if generate_args.channel_id > 999 {
                return Err(anyhow!("Channel ID must be 3 digits or fewer"));
            }
            let keys = load_keys(generate_args.keys_path.as_ref(), args.verbose)?;
            generate(
                &keys,
                &generate_args.binkid,
                generate_args.channel_id,
                generate_args.num_keys,
                args.verbose,
            )?;
        }
        Commands::Validate(validate_args) => {
            let keys = load_keys(validate_args.keys_path.as_ref(), args.verbose)?;
            validate(
                &keys,
                &validate_args.binkid,
                &validate_args.key_to_check,
                args.verbose,
            )?;
        }
        Commands::Confid(confirmation_id_args) => {
            confirmation_id(&confirmation_id_args.instid)?;
        }
    }

    Ok(())
}

fn load_keys<P: AsRef<Path> + std::fmt::Display>(path: Option<P>, verbose: bool) -> Result<Keys> {
    let keys = {
        if let Some(path) = path {
            if verbose {
                println!("Loading keys file {}", path);
            }

            let file = File::open(&path)?;
            let reader = BufReader::new(file);
            let keys: Keys = from_reader(reader)?;

            if verbose {
                println!("Loaded keys from {} successfully", path);
            }

            keys
        } else {
            from_str(std::include_str!("../../../keys.json"))?
        }
    };

    Ok(keys)
}

fn generate(keys: &Keys, bink_id: &str, channel_id: u32, count: u64, verbose: bool) -> Result<()> {
    let bink_id = bink_id.to_ascii_uppercase();
    let bink = &keys.bink[&bink_id];

    // We cannot produce a valid key without knowing the private key k. The reason for this is that
    // we need the result of the function K(x; y) = kG(x; y).
    let private_key = &bink.private;

    // We can, however, validate any given key using the available public key: {p, a, b, G, K}.
    // genOrder the order of the generator G, a value we have to reverse -> Schoof's Algorithm.
    let gen_order = &bink.n;

    let curve = initialize_curve(bink, &bink_id, verbose)?;
    if verbose {
        println!(" n: {gen_order}");
        println!(" k: {private_key}");
        println!();
    }
    let private_key = PrivateKey::new(gen_order, private_key)?;

    if u32::from_str_radix(&bink_id, 16)? < 0x40 {
        bink1998_generate(&curve, &private_key, channel_id, count, verbose)?;
    } else {
        bink2002_generate(&curve, &private_key, channel_id, count, verbose)?;
    }

    Ok(())
}

fn validate(keys: &Keys, bink_id: &str, key: &str, verbose: bool) -> Result<()> {
    let bink_id = bink_id.to_ascii_uppercase();
    let bink = &keys.bink[&bink_id];
    let curve = initialize_curve(bink, &bink_id, verbose)?;

    if u32::from_str_radix(&bink_id, 16)? < 0x40 {
        bink1998_validate(&curve, key, verbose)?;
    } else {
        bink2002_validate(&curve, key, verbose)?;
    }

    Ok(())
}

fn initialize_curve(bink: &Bink, bink_id: &str, verbose: bool) -> Result<EllipticCurve> {
    let p = &bink.p;
    let a = &bink.a;
    let b = &bink.b;
    let gx = &bink.g.x;
    let gy = &bink.g.y;
    let kx = &bink.public.x;
    let ky = &bink.public.y;

    if verbose {
        println!("-----------------------------------------------------------");
        println!(
            "Loaded the following elliptic curve parameters: BINK[{}]",
            bink_id
        );
        println!("-----------------------------------------------------------");
        println!(" P: {p}");
        println!(" a: {a}");
        println!(" b: {b}");
        println!("Gx: {gx}");
        println!("Gy: {gy}");
        println!("Kx: {kx}");
        println!("Ky: {ky}");
    }

    EllipticCurve::new(p, a, b, gx, gy, kx, ky)
}

fn bink1998_generate(
    curve: &EllipticCurve,
    private_key: &PrivateKey,
    channel_id: u32,
    count: u64,
    verbose: bool,
) -> Result<()> {
    for _ in 0..count {
        let product_key = bink1998::ProductKey::new(curve, private_key, channel_id, None, None)?;
        if verbose {
            println!("{:?}", product_key);
        }
        println!("{product_key}");
    }
    Ok(())
}

fn bink2002_generate(
    curve: &EllipticCurve,
    private_key: &PrivateKey,
    channel_id: u32,
    count: u64,
    verbose: bool,
) -> Result<()> {
    for _ in 0..count {
        let product_key = bink2002::ProductKey::new(curve, private_key, channel_id, None, None)?;
        if verbose {
            println!("{:?}", product_key);
        }
        println!("{product_key}");
    }
    Ok(())
}

fn bink1998_validate(curve: &EllipticCurve, key: &str, verbose: bool) -> Result<()> {
    let product_key = bink1998::ProductKey::from_key(curve, key)?;
    if verbose {
        println!("{:?}", product_key);
    }
    println!("{product_key}");
    println!("Key validated successfully!");
    Ok(())
}

fn bink2002_validate(curve: &EllipticCurve, key: &str, verbose: bool) -> Result<()> {
    let product_key = bink2002::ProductKey::from_key(curve, key)?;
    if verbose {
        println!("{:?}", product_key);
    }
    println!("{product_key}");
    println!("Key validated successfully!");
    Ok(())
}

fn confirmation_id(installation_id: &str) -> Result<()> {
    let confirmation_id = confid::generate(installation_id)?;
    println!("Confirmation ID: {confirmation_id}");
    Ok(())
}