//! Command-line parsing, by hand so the tool stays dependency-free.

use std::time::Duration;

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum Mode {
    /// Sit in front of a local application: accept plaintext, encrypt upstream.
    Client,
    /// Sit in front of a local service: accept encrypted, decrypt downstream.
    Server,
}

pub const DEFAULT_MAX_CONNS: usize = 2;

/// No realistic shared secret runs past this. Anything longer is junk, not a
/// genuine configuration, so we refuse to start rather than honour it.
pub const MAX_PASSWORD_LEN: usize = 64;

/// A client names the target it wants by host and port, and the host travels on
/// the wire behind a single length byte. A DNS name tops out at 253 bytes and an
/// IPv6 literal is far shorter, so 255 is room to spare while keeping the field
/// length-prefixed by one byte.
pub const MAX_TARGET_HOST_LEN: usize = 255;

/// Fallback source for the password when `--password` is omitted. Handy for
/// keeping the secret out of the process table and shell history.
pub const PASSWORD_ENV: &str = "JOKYCRYPT_PASSWORD";

/// Idle read/write timeout applied to each socket unless `--timeout` overrides
/// it. Long enough to ride out a quiet stretch on a real connection, short
/// enough that a wedged or vanished peer cannot pin a slot indefinitely.
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60);

/// Where to bind the listening socket when the listen address omits a host.
pub const DEFAULT_LISTEN_HOST: &str = "0.0.0.0";

/// Where to dial when the target address omits a host.
pub const DEFAULT_TARGET_HOST: &str = "127.0.0.1";

/// A network endpoint: a host (name or IP literal) paired with a port.
#[derive(Clone)]
pub struct Addr {
    pub host: String,
    pub port: u16,
}

impl std::fmt::Display for Addr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}:{}", self.host, self.port)
    }
}

pub struct Config {
    pub listen: Addr,
    /// Where to forward to. For a server this is the allow-list of targets a
    /// client may ask to reach (at least one, each a distinct host:port, and the
    /// first is the default). For a client it is exactly one entry: the server
    /// to dial.
    pub targets: Vec<Addr>,
    /// Client only: which of the server's allowed targets to reach. `None` means
    /// "no preference" — the server picks its first. When set, the host defaults
    /// to 127.0.0.1 just as the server's targets do, so usually this is only a
    /// port; spell out a host to disambiguate two targets that share one.
    pub requested_target: Option<Addr>,
    pub password: Vec<u8>,
    pub mode: Mode,
    pub max_conns: Option<usize>,
    pub timeout: Option<Duration>,
    pub verbose: bool,
}

impl Config {
    /// Server side: resolve a client's requested target against the allow-list.
    /// An empty host with port 0 carries no preference and selects the first
    /// configured target; otherwise the match is on host *and* port, so two
    /// targets may share a port as long as their hosts differ. Returns `None`
    /// when nothing in the allow-list matches.
    pub fn select_target(&self, host: &str, port: u16) -> Option<&Addr> {
        if host.is_empty() && port == 0 {
            self.targets.first()
        } else {
            self.targets
                .iter()
                .find(|t| t.host == host && t.port == port)
        }
    }
}

pub const HELP: &str = "\
jokycrypt — a fast, deliberately insecure encrypting TCP relay

USAGE:
    jow <client|server> <LISTEN> <TARGET> [OPTIONS]

ROLE (first argument, required):
    client    accept plaintext locally and encrypt it toward the relay
    server    accept ciphertext and decrypt it toward a target service

ADDRESSES (required):
    LISTEN    address to listen on, written as [host:]port; the host is optional
              and defaults to all interfaces, so '8000' listens on every IP
    TARGET    where to forward to. A server takes a comma-separated allow-list of
              targets, each written as [host:]port, e.g. '80,443,192.168.1.50:8888';
              the host of each defaults to 127.0.0.1. A client takes a single
              address: the server to dial.

OPTIONS:
    -p, --password <PW>      Shared password, 1-64 bytes (also seeds the cipher)
                             [env: JOKYCRYPT_PASSWORD if --password is omitted]
    -t, --target <TARGET>    (client) which of the server's allowed targets to
                             reach, as [host:]port; usually just a port. Defaults
                             to the server's first target.
    -c, --max-conns <N>      Cap on simultaneous connections [default: 2]
        --timeout <SECS>     Idle read/write timeout per socket [default: 60]
    -v, --verbose            Print the plaintext crossing the relay
    -h, --help               Show this help

EXAMPLES:
    # server: let decrypted traffic reach local :80, local :443, or a remote host
    jow server 9000 80,443,192.168.1.50:8888 --password hunter2

    # client: encrypt local :8000 traffic toward the server, asking for port 443
    jow client 8000 machine-b:9000 --target 443 --password hunter2

A client/server pair forms a tunnel:

    app --plain--> [client LISTEN -> server] ==cipher==> [server -> chosen target] --plain--> service
";

fn parse_port(value: &str, what: &str) -> Result<u16, String> {
    value
        .parse::<u16>()
        .map_err(|_| format!("invalid {what}: '{value}' is not a port number"))
}

/// Parse a positional address written as `[host:]port`. The host is optional;
/// when it is absent — or given as the empty string — the supplied default is
/// used. IPv6 literals must be bracketed, e.g. `[::1]:8000`, so the colons
/// inside them are not mistaken for the host/port separator.
fn parse_addr(value: &str, what: &str, default_host: &str) -> Result<Addr, String> {
    // A bare port with no host at all: the most common, terse form.
    if let Ok(port) = value.parse::<u16>() {
        return Ok(Addr {
            host: default_host.to_string(),
            port,
        });
    }

    let (host, port_str) = if let Some(rest) = value.strip_prefix('[') {
        let close = rest
            .find(']')
            .ok_or_else(|| format!("invalid {what} '{value}': unmatched '['"))?;
        let port_str = rest[close + 1..]
            .strip_prefix(':')
            .ok_or_else(|| format!("invalid {what} '{value}': expected ':<port>' after ']'"))?;
        (&rest[..close], port_str)
    } else {
        value
            .rsplit_once(':')
            .ok_or_else(|| format!("invalid {what} '{value}': expected '[host:]port'"))?
    };

    let port = parse_port(port_str, what)?;
    let host = if host.is_empty() { default_host } else { host };
    Ok(Addr {
        host: host.to_string(),
        port,
    })
}

/// Parse a comma-separated list of targets, each an `[host:]port` understood by
/// `parse_addr`. Empty entries and a port of 0 are rejected — neither names a
/// reachable service.
fn parse_targets(value: &str, what: &str, default_host: &str) -> Result<Vec<Addr>, String> {
    let mut targets = Vec::new();
    for piece in value.split(',') {
        if piece.is_empty() {
            return Err(format!(
                "{what} has an empty entry; write targets as 'port[,host:port,...]'"
            ));
        }
        let addr = parse_addr(piece, what, default_host)?;
        if addr.port == 0 {
            return Err(format!("invalid {what} '{piece}': port 0 is not a target"));
        }
        targets.push(addr);
    }
    Ok(targets)
}

fn parse_mode(value: &str) -> Result<Mode, String> {
    match value {
        "client" => Ok(Mode::Client),
        "server" => Ok(Mode::Server),
        other => Err(format!(
            "invalid role '{other}', expected 'client' or 'server'"
        )),
    }
}

pub fn parse<I: Iterator<Item = String>>(
    mut args: I,
    env_password: Option<String>,
) -> Result<Config, String> {
    let mut positionals: Vec<String> = Vec::new();
    let mut password = None;
    let mut requested_target = None;
    let mut max_conns = Some(DEFAULT_MAX_CONNS);
    let mut timeout = Some(DEFAULT_TIMEOUT);
    let mut verbose = false;

    let take = |args: &mut I, flag: &str| -> Result<String, String> {
        args.next().ok_or_else(|| format!("{flag} expects a value"))
    };

    while let Some(arg) = args.next() {
        match arg.as_str() {
            "-h" | "--help" => return Err(HELP.to_string()),
            "-v" | "--verbose" => verbose = true,
            "-p" | "--password" => password = Some(take(&mut args, &arg)?.into_bytes()),
            "-t" | "--target" => {
                let v = take(&mut args, &arg)?;
                let addr = parse_addr(&v, "target", DEFAULT_TARGET_HOST)?;
                if addr.port == 0 {
                    return Err("--target must name a port between 1 and 65535".to_string());
                }
                if addr.host.len() > MAX_TARGET_HOST_LEN {
                    return Err(format!(
                        "--target host is {} bytes; the maximum is {MAX_TARGET_HOST_LEN}",
                        addr.host.len()
                    ));
                }
                requested_target = Some(addr);
            }
            "-c" | "--max-conns" => {
                let v = take(&mut args, &arg)?;
                let n = v
                    .parse::<usize>()
                    .map_err(|_| format!("invalid max-conns: '{v}'"))?;
                if n == 0 {
                    return Err("max-conns must be at least 1".to_string());
                }
                max_conns = Some(n);
            }
            "--timeout" => {
                let v = take(&mut args, &arg)?;
                let secs = v
                    .parse::<u64>()
                    .map_err(|_| format!("invalid timeout: '{v}'"))?;
                if secs == 0 {
                    return Err("timeout must be at least 1 second".to_string());
                }
                timeout = Some(Duration::from_secs(secs));
            }
            other if other.starts_with('-') => {
                return Err(format!("unknown option '{other}' (try --help)"));
            }
            _ => positionals.push(arg),
        }
    }

    let mut positionals = positionals.into_iter();
    let mode = parse_mode(
        &positionals
            .next()
            .ok_or("missing role: expected 'client' or 'server' as the first argument")?,
    )?;
    let listen_arg = positionals
        .next()
        .ok_or("missing listen address: expected '[host:]port'")?;
    let target_arg = positionals
        .next()
        .ok_or("missing target address: expected '[host:]port'")?;
    if let Some(extra) = positionals.next() {
        return Err(format!("unexpected extra argument '{extra}' (try --help)"));
    }

    let listen = parse_addr(&listen_arg, "listen address", DEFAULT_LISTEN_HOST)?;

    // The TARGET positional means different things to the two roles: a server is
    // handed the allow-list of places a client may reach, a client just the one
    // server to dial. Parse and validate each accordingly.
    let targets = match mode {
        Mode::Server => {
            if requested_target.is_some() {
                return Err("--target is only meaningful for the client".to_string());
            }
            let targets = parse_targets(&target_arg, "allowed targets", DEFAULT_TARGET_HOST)?;
            let mut seen = std::collections::HashSet::new();
            for t in &targets {
                if !seen.insert((t.host.as_str(), t.port)) {
                    return Err(format!("target {t} is listed twice in the allow-list"));
                }
            }
            targets
        }
        Mode::Client => {
            let targets = parse_targets(&target_arg, "target address", DEFAULT_TARGET_HOST)?;
            if targets.len() != 1 {
                return Err(
                    "the client takes a single target — the server to dial, not a list".to_string(),
                );
            }
            targets
        }
    };

    // The command line wins; the environment variable is only a fallback.
    let password = password
        .or_else(|| env_password.map(String::into_bytes))
        .ok_or(format!(
            "missing password: pass --password <PW> or set {PASSWORD_ENV}"
        ))?;
    if password.is_empty() {
        return Err("password must not be empty".to_string());
    }
    if password.len() > MAX_PASSWORD_LEN {
        return Err(format!(
            "password is {} bytes; the maximum is {MAX_PASSWORD_LEN}",
            password.len()
        ));
    }

    Ok(Config {
        listen,
        targets,
        requested_target,
        password,
        mode,
        max_conns,
        timeout,
        verbose,
    })
}

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

    fn parse_args(args: &[&str]) -> Result<Config, String> {
        super::parse(args.iter().map(|s| s.to_string()), None)
    }

    fn parse_with_env(args: &[&str], env_password: Option<&str>) -> Result<Config, String> {
        super::parse(
            args.iter().map(|s| s.to_string()),
            env_password.map(|s| s.to_string()),
        )
    }

    #[test]
    fn parses_a_full_command_line() {
        let c = parse_args(&[
            "client",
            "example.com:8000",
            "10.0.0.5:9000",
            "-p",
            "secret",
            "-c",
            "32",
            "--timeout",
            "60",
            "-v",
        ])
        .unwrap();
        assert_eq!(c.listen.host, "example.com");
        assert_eq!(c.listen.port, 8000);
        assert_eq!(c.targets.len(), 1);
        assert_eq!(c.targets[0].host, "10.0.0.5");
        assert_eq!(c.targets[0].port, 9000);
        assert_eq!(c.password, b"secret");
        assert!(c.mode == Mode::Client);
        assert_eq!(c.max_conns, Some(32));
        assert_eq!(c.timeout, Some(Duration::from_secs(60)));
        assert!(c.verbose);
    }

    #[test]
    fn hosts_default_when_omitted() {
        let c = parse_args(&["server", "8888", "9999", "-p", "x"]).unwrap();
        assert_eq!(c.listen.host, DEFAULT_LISTEN_HOST);
        assert_eq!(c.listen.port, 8888);
        assert_eq!(c.targets[0].host, DEFAULT_TARGET_HOST);
        assert_eq!(c.targets[0].port, 9999);
        assert!(c.mode == Mode::Server);
        assert_eq!(c.max_conns, Some(DEFAULT_MAX_CONNS));
        assert_eq!(c.timeout, Some(DEFAULT_TIMEOUT));
        assert!(!c.verbose);
    }

    #[test]
    fn empty_host_means_all_interfaces() {
        let c = parse_args(&["server", ":8888", "192.168.1.254:9999", "-p", "x"]).unwrap();
        assert_eq!(c.listen.host, DEFAULT_LISTEN_HOST);
        assert_eq!(c.listen.port, 8888);
        assert_eq!(c.targets[0].host, "192.168.1.254");
        assert_eq!(c.targets[0].port, 9999);
    }

    #[test]
    fn parses_bracketed_ipv6() {
        let c = parse_args(&["client", "[::1]:8000", "[2001:db8::1]:9000", "-p", "x"]).unwrap();
        assert_eq!(c.listen.host, "::1");
        assert_eq!(c.listen.port, 8000);
        assert_eq!(c.targets[0].host, "2001:db8::1");
        assert_eq!(c.targets[0].port, 9000);
    }

    #[test]
    fn options_may_precede_the_positionals() {
        let c = parse_args(&["-v", "-p", "x", "client", "8000", "9000"]).unwrap();
        assert!(c.mode == Mode::Client);
        assert_eq!(c.listen.port, 8000);
        assert_eq!(c.targets[0].port, 9000);
        assert!(c.verbose);
    }

    #[test]
    fn missing_required_fields_error() {
        assert!(parse_args(&[]).is_err());
        assert!(parse_args(&["client"]).is_err());
        assert!(parse_args(&["client", "8000"]).is_err());
    }

    #[test]
    fn rejects_extra_positional() {
        assert!(parse_args(&["server", "1", "2", "3", "-p", "x"]).is_err());
    }

    #[test]
    fn rejects_unknown_option() {
        assert!(parse_args(&["server", "1", "2", "-p", "x", "--nope"]).is_err());
    }

    #[test]
    fn rejects_bad_values() {
        assert!(parse_args(&["client", "notaport", "2", "-p", "x"]).is_err());
        assert!(parse_args(&["sideways", "1", "2", "-p", "x"]).is_err());
        assert!(parse_args(&["client", "1", "2", "-p", "x", "-c", "0"]).is_err());
        assert!(parse_args(&["client", "1", "2", "-p", ""]).is_err());
        assert!(parse_args(&["client", "host:", "2", "-p", "x"]).is_err());
        assert!(parse_args(&["client", "[::1]8000", "2", "-p", "x"]).is_err());
    }

    #[test]
    fn password_falls_back_to_the_environment() {
        let c = parse_with_env(&["server", "1", "2"], Some("from-env")).unwrap();
        assert_eq!(c.password, b"from-env");
    }

    #[test]
    fn command_line_password_beats_the_environment() {
        let c = parse_with_env(&["server", "1", "2", "-p", "from-cli"], Some("from-env")).unwrap();
        assert_eq!(c.password, b"from-cli");
    }

    #[test]
    fn missing_password_everywhere_is_an_error() {
        assert!(parse_with_env(&["server", "1", "2"], None).is_err());
    }

    #[test]
    fn environment_password_must_also_pass_validation() {
        let over_limit = "p".repeat(MAX_PASSWORD_LEN + 1);
        assert!(parse_with_env(&["server", "1", "2"], Some("")).is_err());
        assert!(parse_with_env(&["server", "1", "2"], Some(&over_limit)).is_err());
    }

    #[test]
    fn enforces_password_length_bound() {
        let at_limit = "p".repeat(MAX_PASSWORD_LEN);
        let over_limit = "p".repeat(MAX_PASSWORD_LEN + 1);
        assert!(parse_args(&["server", "1", "2", "-p", &at_limit]).is_ok());
        assert!(parse_args(&["server", "1", "2", "-p", &over_limit]).is_err());
    }

    #[test]
    fn server_parses_a_target_allow_list() {
        let c = parse_args(&["server", "9000", "80,443,192.168.1.50:8888", "-p", "x"]).unwrap();
        assert_eq!(c.targets.len(), 3);
        assert_eq!(
            (c.targets[0].host.as_str(), c.targets[0].port),
            ("127.0.0.1", 80)
        );
        assert_eq!(
            (c.targets[1].host.as_str(), c.targets[1].port),
            ("127.0.0.1", 443)
        );
        assert_eq!(
            (c.targets[2].host.as_str(), c.targets[2].port),
            ("192.168.1.50", 8888)
        );
    }

    #[test]
    fn server_allows_one_port_across_several_hosts() {
        let c = parse_args(&["server", "9000", "10.0.0.1:8000,10.0.0.2:8000", "-p", "x"]).unwrap();
        assert_eq!(c.targets.len(), 2);
        assert_eq!(c.targets[0].host, "10.0.0.1");
        assert_eq!(c.targets[1].host, "10.0.0.2");
    }

    #[test]
    fn server_rejects_a_duplicated_target() {
        // An exact host:port repeat is a mistake; differing hosts are fine.
        assert!(parse_args(&["server", "9000", "80,127.0.0.1:80", "-p", "x"]).is_err());
        assert!(parse_args(&["server", "9000", "80,80", "-p", "x"]).is_err());
    }

    #[test]
    fn server_rejects_empty_or_zero_targets() {
        assert!(parse_args(&["server", "9000", "80,,443", "-p", "x"]).is_err());
        assert!(parse_args(&["server", "9000", "0", "-p", "x"]).is_err());
        assert!(parse_args(&["server", "9000", "", "-p", "x"]).is_err());
    }

    #[test]
    fn client_refuses_a_target_list() {
        assert!(parse_args(&["client", "8000", "a:9000,b:9000", "-p", "x"]).is_err());
    }

    #[test]
    fn client_target_flag_defaults_its_host() {
        let c = parse_args(&["client", "8000", "host:9000", "-p", "x", "--target", "443"]).unwrap();
        let req = c.requested_target.unwrap();
        assert_eq!(req.host, DEFAULT_TARGET_HOST);
        assert_eq!(req.port, 443);
    }

    #[test]
    fn client_target_flag_keeps_an_explicit_host() {
        let c = parse_args(&[
            "client",
            "8000",
            "host:9000",
            "-p",
            "x",
            "-t",
            "192.168.1.50:8888",
        ])
        .unwrap();
        let req = c.requested_target.unwrap();
        assert_eq!(req.host, "192.168.1.50");
        assert_eq!(req.port, 8888);
    }

    #[test]
    fn no_target_flag_means_no_preference() {
        let c = parse_args(&["client", "8000", "9000", "-p", "x"]).unwrap();
        assert!(c.requested_target.is_none());
    }

    #[test]
    fn target_flag_is_server_only_error() {
        assert!(parse_args(&["server", "9000", "80,443", "-p", "x", "--target", "443"]).is_err());
    }

    #[test]
    fn target_flag_rejects_port_zero() {
        assert!(parse_args(&["client", "8000", "9000", "-p", "x", "--target", "0"]).is_err());
    }

    #[test]
    fn select_target_resolves_against_the_allow_list() {
        let c = parse_args(&["server", "9000", "80,443,192.168.1.50:8888", "-p", "x"]).unwrap();
        // No preference -> the first configured target.
        assert_eq!(c.select_target("", 0).map(|t| t.port), Some(80));
        // A plain port matches the localhost target on that port.
        assert_eq!(c.select_target("127.0.0.1", 443).map(|t| t.port), Some(443));
        // Host and port together pick the remote target.
        let remote = c.select_target("192.168.1.50", 8888).unwrap();
        assert_eq!((remote.host.as_str(), remote.port), ("192.168.1.50", 8888));
        // A port that exists only on another host does not match.
        assert!(c.select_target("127.0.0.1", 8888).is_none());
        // An unknown port is refused.
        assert!(c.select_target("127.0.0.1", 22).is_none());
    }
}
