use anyhow::{bail, ensure, Context, Result};
use axum::body::Body;
use clap::{parser::ValueSource, Args, CommandFactory, FromArgMatches, Parser, Subcommand};
use http::{Method, Request, StatusCode};
use http_body_util::{BodyExt, Limited};
use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder};
use hyper_util::{
    client::legacy::{connect::HttpConnector, Client},
    rt::TokioExecutor,
};
use rust_proxy::{
    auth,
    config::{self, Config, Envelope, ValidationPolicy},
    db::Store,
    runtime::App,
};
use rustls::{
    client::{
        danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
        WebPkiServerVerifier,
    },
    pki_types::{CertificateDer, ServerName, UnixTime},
    ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme,
};
use sha2::{Digest, Sha256};
use std::io::{self, BufRead, IsTerminal, Read, Write};
use std::{
    net::SocketAddr,
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};
use tokio::{
    net::{TcpListener, UdpSocket},
    task::JoinSet,
};
use tracing_subscriber::prelude::*;
use zeroize::Zeroizing;

#[derive(Parser)]
#[command(
    name = "rust-proxy",
    version,
    about = "RustProxy - independent Rust reverse proxy and static Web Server 0.27.0"
)]
struct Cli {
    #[arg(
        long,
        global = true,
        env = "RUST_PROXY_DATA_DIR",
        default_value = "./data"
    )]
    data_dir: PathBuf,
    #[arg(
        long,
        global = true,
        env = "RUST_PROXY_ADMIN_LISTEN",
        default_value = "127.0.0.1:9090"
    )]
    admin_listen: SocketAddr,
    #[arg(long, global = true, env = "RUST_PROXY_UPDATE_REPOSITORY_URL")]
    update_repository_url: Option<String>,
    #[arg(long, global = true, env = "RUST_PROXY_UPDATE_VERSION_URL")]
    update_version_url: Option<String>,
    /// Seconds between release checks. Set to 0 to disable automatic checks.
    #[arg(long, global = true, env = "RUST_PROXY_UPDATE_CHECK_INTERVAL_SECS")]
    update_check_interval_secs: Option<u64>,
    /// Deprecated compatibility no-op kept so older systemd units continue to start.
    #[arg(long = "accept-remote-admin-risk", global = true, hide = true)]
    _accept_remote_admin_risk: bool,
    /// Interactive varnishlog-style live request view from all vhosts.
    #[arg(long, global = true)]
    livelogs: bool,
    #[command(subcommand)]
    command: Option<Command>,
}
#[derive(Subcommand)]
enum Command {
    /// Start proxy and local administration listener.
    Serve(ServeOptions),
    /// Compatibility alias for `admin setup-token`; never an API access token.
    Token,
    /// Local administrator setup, password recovery and session management.
    Admin {
        #[command(subcommand)]
        action: AdminCommand,
    },
    /// Manage local certificate files and hot-reload certificate bindings.
    Tls {
        #[command(subcommand)]
        action: TlsCommand,
    },
    /// Read or replace configuration through the running admin API.
    Config {
        #[command(subcommand)]
        action: ConfigCommand,
    },
    /// Create or restore a password-encrypted full application backup.
    Backup {
        #[command(subcommand)]
        action: BackupCommand,
    },
}
#[derive(Subcommand)]
enum AdminCommand {
    /// Create the sole administrator. Requires ownership of the private data directory.
    Init {
        #[arg(long)]
        username: Option<String>,
        #[arg(long)]
        password_stdin: bool,
    },
    /// Replace the administrator password and revoke all sessions, online or offline.
    ResetPassword {
        #[arg(long)]
        password_stdin: bool,
    },
    /// Generate or print the one-time setup capability, only before an administrator exists.
    SetupToken,
    /// Report whether administrator setup is complete (no credentials printed).
    Status,
    /// Revoke all browser/CLI sessions without changing the password.
    LogoutAll,
}
#[derive(Subcommand)]
enum ConfigCommand {
    /// Print desired configuration as JSON (suitable for `config apply`).
    Get,
    /// Revalidate and hot-publish the current configuration and certificate files.
    Reload,
    /// Atomically validate, save and apply a full configuration document.
    Apply { file: PathBuf },
    /// Validate JSON, ACLs, headers and backend policy without changing the database.
    Validate { file: PathBuf },
}
#[derive(Subcommand)]
enum BackupCommand {
    /// Create an encrypted full backup. Safe while RustProxy is serving.
    Create {
        /// Destination .rpb file. Existing files are never overwritten.
        #[arg(long)]
        output: Option<PathBuf>,
    },
    /// Restore an encrypted full backup atomically. Stop the RustProxy service first.
    Restore { archive: PathBuf },
}

#[derive(Subcommand)]
enum TlsCommand {
    /// Generate a self-signed development certificate. Existing files are never overwritten.
    Generate {
        #[arg(long)]
        name: String,
        #[arg(long, value_delimiter = ',', required = true)]
        hosts: Vec<String>,
    },
    /// Import a matching PEM certificate/key pair under a new name.
    Import {
        #[arg(long)]
        name: String,
        #[arg(long)]
        cert: PathBuf,
        #[arg(long)]
        key: PathBuf,
    },
    /// Stage an externally issued pair (for example Certbot) and atomically
    /// switch an existing manual certificate id to it through the running API.
    Deploy {
        #[arg(long)]
        id: String,
        #[arg(long)]
        cert: PathBuf,
        #[arg(long)]
        key: PathBuf,
    },
    /// Reload configured certificate files without restarting listeners.
    Reload,
}
#[derive(Args)]
struct ServeOptions {
    /// Optional HTTPS listener. dev.sh enables 127.0.0.1:8443 unless RUST_PROXY_TLS=0.
    #[arg(long, env = "RUST_PROXY_TLS_LISTEN")]
    tls_listen: Option<SocketAddr>,
    #[arg(
        long,
        env = "RUST_PROXY_PROXY_LISTEN",
        default_value = "127.0.0.1:8080"
    )]
    proxy_listen: SocketAddr,
    #[arg(long)]
    demo: bool,
    #[arg(long, env = "RUST_PROXY_DEMO_LISTEN", default_value = "127.0.0.1:9000")]
    demo_listen: SocketAddr,
    #[arg(long, default_value_t = 512)]
    max_connections: usize,
    #[arg(long, default_value_t = 256)]
    max_inflight: usize,
    /// Includes upload time until upstream response headers arrive.
    #[arg(long, default_value_t = 120)]
    upstream_timeout_secs: u64,
    #[arg(long, default_value_t = 60)]
    body_idle_timeout_secs: u64,
    /// Inactivity in either direction; independent of ordinary HTTP bodies.
    #[arg(long, default_value_t = 300)]
    websocket_idle_timeout_secs: u64,
    #[arg(long, default_value_t = 10)]
    shutdown_timeout_secs: u64,
    /// Permit an administration listener on an unspecified address such as 0.0.0.0.
    #[arg(long)]
    allow_remote_admin: bool,
}
#[cfg(target_os = "linux")]
fn disable_transparent_huge_pages() {
    // Keep THP policy local to RustProxy. Large 2 MiB anonymous pages amplify
    // allocator fragmentation/retention for the long-running proxy workload, so
    // opt this process out before Tokio creates worker threads and their heaps.
    let result = unsafe { libc::prctl(libc::PR_SET_THP_DISABLE, 1, 0, 0, 0) };
    if result != 0 {
        eprintln!(
            "RustProxy warning: could not disable Transparent Huge Pages for this process: {}",
            io::Error::last_os_error()
        );
    }
}

#[cfg(not(target_os = "linux"))]
fn disable_transparent_huge_pages() {}

fn main() -> Result<()> {
    disable_transparent_huge_pages();
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()?;
    let result = runtime.block_on(application());
    // OS getaddrinfo workers cannot be killed by cancelling their future.
    // Bound runtime shutdown after listeners and application tasks have drained.
    runtime.shutdown_timeout(Duration::from_secs(2));
    result
}

fn installed_data_dir_candidates() -> &'static [&'static str] {
    #[cfg(any(
        target_os = "freebsd",
        target_os = "openbsd",
        target_os = "netbsd",
        target_os = "dragonfly"
    ))]
    {
        return &["/var/db/rustproxy", "/var/lib/rustproxy"];
    }
    #[cfg(not(any(
        target_os = "freebsd",
        target_os = "openbsd",
        target_os = "netbsd",
        target_os = "dragonfly"
    )))]
    {
        &["/var/lib/rustproxy", "/var/db/rustproxy"]
    }
}

fn resolve_installed_data_dir(configured: PathBuf, explicitly_configured: bool) -> PathBuf {
    if explicitly_configured || configured.join("config.sqlite3").is_file() {
        return configured;
    }
    installed_data_dir_candidates()
        .iter()
        .map(PathBuf::from)
        .find(|candidate| candidate.join("config.sqlite3").is_file())
        .unwrap_or(configured)
}

async fn application() -> Result<()> {
    let error_log = rust_proxy::logging::ErrorLogController::start();
    rust_proxy::logging::install_error_log_controller(error_log.clone())?;
    let filter = tracing_subscriber::EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| "rust_proxy=info,warn".into());
    tracing_subscriber::registry()
        .with(filter)
        .with(
            tracing_subscriber::fmt::layer()
                .with_target(false)
                .with_writer(rust_proxy::logging::RedactingStderr),
        )
        .with(rust_proxy::logging::ErrorLogLayer::new(error_log))
        .init();
    let matches = Cli::command().get_matches();
    let explicit = |matches: &clap::ArgMatches, id: &str| {
        matches.value_source(id).is_some_and(|source| {
            matches!(source, ValueSource::CommandLine | ValueSource::EnvVariable)
        })
    };
    let serve_matches = matches.subcommand_matches("serve");
    let startup_locks = rust_proxy::settings::StartupLocks {
        administration: explicit(&matches, "admin_listen").then(|| {
            *matches
                .get_one::<SocketAddr>("admin_listen")
                .expect("parsed admin listener")
        }),
        http: serve_matches
            .filter(|sub| explicit(sub, "proxy_listen"))
            .map(|sub| {
                *sub.get_one::<SocketAddr>("proxy_listen")
                    .expect("parsed proxy listener")
            }),
        https: serve_matches
            .filter(|sub| explicit(sub, "tls_listen"))
            .and_then(|sub| sub.get_one::<SocketAddr>("tls_listen").copied()),
    };
    let cli = Cli::from_arg_matches(&matches)?;
    let data_dir_explicit = explicit(&matches, "data_dir");
    if cli.livelogs {
        ensure!(
            cli.command.is_none(),
            "--livelogs cannot be combined with a subcommand"
        );
        let data_dir = resolve_installed_data_dir(cli.data_dir, data_dir_explicit);
        return live_logs_command(data_dir, cli.admin_listen).await;
    }
    let command = cli
        .command
        .context("a subcommand is required (or use --livelogs)")?;
    match command {
        Command::Token => {
            admin_command(
                resolve_installed_data_dir(cli.data_dir, data_dir_explicit),
                AdminCommand::SetupToken,
            )
            .await
        }
        Command::Admin { action } => {
            admin_command(
                resolve_installed_data_dir(cli.data_dir, data_dir_explicit),
                action,
            )
            .await
        }
        Command::Tls { action } => {
            tls_command(
                resolve_installed_data_dir(cli.data_dir, data_dir_explicit),
                cli.admin_listen,
                action,
            )
            .await
        }
        Command::Config { action } => {
            config_command(
                resolve_installed_data_dir(cli.data_dir, data_dir_explicit),
                cli.admin_listen,
                action,
            )
            .await
        }
        Command::Backup { action } => {
            backup_command(
                resolve_installed_data_dir(cli.data_dir, data_dir_explicit),
                action,
            )
            .await
        }
        Command::Serve(options) => {
            let update_check = rust_proxy::update_check::Config::embedded()?.with_overrides(
                cli.update_repository_url,
                cli.update_version_url,
                cli.update_check_interval_secs,
            )?;
            // Server startup never guesses a data directory. Package/service/plugin
            // installation must pass its configured path explicitly.
            run(
                cli.data_dir,
                cli.admin_listen,
                options,
                startup_locks,
                update_check,
            )
            .await
        }
    }
}
async fn backup_command(data_dir: PathBuf, action: BackupCommand) -> Result<()> {
    match action {
        BackupCommand::Create { output } => {
            let output = output.unwrap_or_else(|| {
                PathBuf::from(format!("rustproxy-backup-{}.rpb", rust_proxy::auth::now()))
            });
            let mut generated = false;
            let password = if io::stdin().is_terminal() {
                let value = Zeroizing::new(rpassword::prompt_password(
                    "Backup password (leave empty to generate): ",
                )?);
                if value.is_empty() {
                    generated = true;
                    Zeroizing::new(rust_proxy::backup::generated_password())
                } else {
                    let confirmation =
                        Zeroizing::new(rpassword::prompt_password("Confirm backup password: ")?);
                    ensure!(
                        value.as_str() == confirmation.as_str(),
                        "backup passwords do not match"
                    );
                    value
                }
            } else {
                generated = true;
                Zeroizing::new(rust_proxy::backup::generated_password())
            };
            let summary = rust_proxy::backup::create(&data_dir, &output, password.as_str()).await?;
            println!(
                "Encrypted backup created: {} ({} files, configuration revision {})",
                summary.path.display(),
                summary.files,
                summary.configuration_revision
            );
            if generated {
                println!(
                    "Generated backup password (shown once): {}",
                    password.as_str()
                );
                println!("Store this password securely; restore is impossible without it.");
            }
            Ok(())
        }
        BackupCommand::Restore { archive } => {
            ensure!(
                io::stdin().is_terminal(),
                "backup restore requires an interactive TTY for the password"
            );
            let password = Zeroizing::new(rpassword::prompt_password("Backup password: ")?);
            ensure!(!password.is_empty(), "backup password cannot be empty");
            let summary =
                rust_proxy::backup::restore(&archive, &data_dir, password.as_str()).await?;
            println!(
                "Backup restored to {} ({} files, configuration revision {})",
                data_dir.display(),
                summary.files,
                summary.configuration_revision
            );
            Ok(())
        }
    }
}

fn validate_admin_binding(admin_addr: SocketAddr, tls: bool, allow_remote: bool) -> Result<bool> {
    let remote_admin = !admin_addr.ip().is_loopback();
    ensure!(
        !remote_admin || tls || allow_remote,
        "remote plain-HTTP administration requires --allow-remote-admin"
    );
    Ok(remote_admin)
}