use std::ffi::OsString;
use std::fs::File;
use std::os::fd::{AsFd, AsRawFd};
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::time::{Duration, Instant};

use lazaret_core::codec::{decode, encode};
use lazaret_core::identity::{
    ConfigurationFingerprint, CoordinatorId, DomainId, GenerationId, NetworkSidecarId,
};
use lazaret_core::network::{
    DnsQueryExpectation, ReachabilityExpectation, ReadinessExpectation, accept_network_readiness,
    derive_policy, validate_network_attestation,
};
use lazaret_core::records::{
    AddressFamily, AddressFamilyDisposition, AddressFamilyDispositionValue, DefaultAction,
    DerivedNetworkPolicyFields, HostProcessOwnershipSlot, IpAddress, IpNetwork,
    NetworkReadinessAttestation, PolicyAction, PrefixRule, ProcessOwnershipCommitAck, Resolver,
};
use lazaret_linux::filesystem::{atomic_replace, require_write_sealed, sealed_memfd};
use lazaret_linux::process::{
    ProcessOwnershipContext, TargetCommand, child_broker, current_process_identity, recover_process,
};
use lazaret_runtime_supervisor::network::{
    DnsProbeTarget, NetworkAttestationContext, NetworkProbeRequest, TcpProbeTarget,
    enter_generation_filesystem_namespaces, prepare_network, probe_current_network,
};
use rustix::io::{FdFlags, fcntl_setfd};
use rustix::process::Signal;

const CHILD_ROLE: &str = "--network-runtime-child";
const SKIP: u8 = 77;

fn main() -> ExitCode {
    let arguments = std::env::args_os().skip(1).collect::<Vec<_>>();
    if let Some(outcome) = lazaret_runtime_supervisor::private_roles::REGISTRY.dispatch(&arguments)
    {
        return exit(outcome.result);
    }
    if arguments.first().and_then(|value| value.to_str()) == Some(CHILD_ROLE) {
        return exit(run_child(&arguments[1..]));
    }
    match run_parent() {
        Ok(()) => {
            println!("PASS real rootless network namespace and supervised slirp4netns");
            ExitCode::SUCCESS
        }
        Err(error) if prerequisite_failure(&error) => {
            println!("SKIP real network runtime: {error}");
            ExitCode::from(SKIP)
        }
        Err(error) => {
            eprintln!("network runtime failed: {error}");
            ExitCode::FAILURE
        }
    }
}

fn run_parent() -> std::io::Result<()> {
    for executable in ["/usr/bin/slirp4netns", "/usr/bin/ip"] {
        if !Path::new(executable).is_file() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::NotFound,
                format!("required executable is absent: {executable}"),
            ));
        }
    }
    let directory = tempfile::tempdir()?;
    let policy = policy()?;
    let policy_file = sealed_memfd("network-runtime-policy", policy.canonical_bytes())?;
    fcntl_setfd(policy_file.as_fd(), FdFlags::empty())?;
    let verifier_executable = std::env::current_exe()?;
    let target = TargetCommand::new(verifier_executable)
        .argument(CHILD_ROLE)
        .argument(directory.path())
        .argument(policy_file.as_raw_fd().to_string())
        .argument(hex(policy.digest().as_bytes()));
    let parent = current_process_identity()?;
    let runtime = child_broker().acquire(
        parent,
        ProcessOwnershipContext {
            domain_id: DomainId::new([1; 32]),
            generation_id: GenerationId::new(7),
            coordinator_id: CoordinatorId::new([2; 32]),
            slot: HostProcessOwnershipSlot::SandboxRuntimeSupervisor,
        },
        &target,
        |_| Ok(()),
        |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
    )?;
    let attestation_path = directory.path().join("attestation.postcard");
    wait_for_path(&attestation_path, &runtime)?;
    let attestation: NetworkReadinessAttestation = decode(&std::fs::read(&attestation_path)?)
        .map_err(|error| std::io::Error::other(error.to_string()))?;
    if attestation.runtime_supervisor != runtime.identity()
        || !attestation.observation.helper_ready
        || !attestation.observation.routes_attested
        || attestation.observation.observed_routes.len() < 3
    {
        return cleanup_runtime(runtime, std::io::Error::other("attestation is incomplete"));
    }
    let probe_path = directory.path().join("probe.postcard");
    wait_for_path(&probe_path, &runtime)?;
    let probe: lazaret_core::records::NetworkReachabilityProbe =
        decode(&std::fs::read(&probe_path)?)
            .map_err(|error| std::io::Error::other(error.to_string()))?;
    let expectation = ReadinessExpectation {
        domain_id: attestation.domain_id,
        generation_id: attestation.generation_id,
        configuration_fingerprint: attestation.configuration_fingerprint,
        network_sidecar_id: attestation.network_sidecar_id,
        enforcement_namespace: attestation.enforcement_namespace,
        runtime_supervisor: &attestation.runtime_supervisor,
        connectivity_helper: &attestation.connectivity_helper,
    };
    accept_network_readiness(
        &policy,
        &expectation,
        attestation.clone(),
        probe,
        &reachability(),
    )
    .map_err(std::io::Error::other)?;
    let mut mutated = attestation.clone();
    mutated.derived_policy_digest = lazaret_core::identity::DerivedPolicyDigest::new([9; 32]);
    if validate_network_attestation(&policy, &expectation, &mutated).is_ok() {
        return cleanup_runtime(
            runtime,
            std::io::Error::other("mutated attestation was accepted"),
        );
    }
    let helper = recover_process(attestation.connectivity_helper)?;
    helper.send_signal(Signal::KILL)?;
    wait_for_path(&directory.path().join("helper-lost"), &runtime)?;
    runtime.send_signal(Signal::KILL)?;
    child_broker().reap(runtime)
}

fn run_child(arguments: &[OsString]) -> std::io::Result<()> {
    let directory = PathBuf::from(
        arguments
            .first()
            .ok_or_else(|| std::io::Error::other("network child directory is absent"))?,
    );
    let policy_fd = arguments
        .get(1)
        .and_then(|value| value.to_str())
        .and_then(|value| value.parse::<i32>().ok())
        .filter(|value| *value >= 0)
        .ok_or_else(|| std::io::Error::other("network child policy fd is invalid"))?;
    let digest = parse_digest(
        arguments
            .get(2)
            .ok_or_else(|| std::io::Error::other("network child digest is absent"))?,
    )?;
    let policy_file = File::open(format!("/proc/self/fd/{policy_fd}"))?;
    require_write_sealed(&policy_file)?;
    enter_generation_filesystem_namespaces()?;
    let runtime = current_process_identity()?;
    let prepared = prepare_network(
        &policy_file,
        digest,
        NetworkAttestationContext {
            domain_id: DomainId::new([1; 32]),
            generation_id: GenerationId::new(7),
            configuration_fingerprint: ConfigurationFingerprint::new([3; 32]),
            network_sidecar_id: NetworkSidecarId::new([4; 32]),
            runtime_supervisor: runtime,
            coordinator_id: CoordinatorId::new([2; 32]),
        },
        |_| Ok(()),
        |staged| Ok(ProcessOwnershipCommitAck::v1(staged)),
    )?;
    let attestation = prepared.attestation().clone();
    atomic_replace(
        &directory.join("attestation.postcard"),
        &encode(&attestation).map_err(std::io::Error::other)?,
    )?;
    let probe = probe_current_network(&probe_request(attestation.enforcement_namespace))?;
    atomic_replace(
        &directory.join("probe.postcard"),
        &encode(&probe).map_err(std::io::Error::other)?,
    )?;
    loop {
        if prepared.helper_exited()? {
            prepared.reap_exited()?;
            std::fs::write(directory.join("helper-lost"), b"lost")?;
            loop {
                std::thread::sleep(Duration::from_secs(1));
            }
        }
        std::thread::sleep(Duration::from_millis(5));
    }
}

fn policy() -> std::io::Result<lazaret_core::network::CanonicalDerivedNetworkPolicy> {
    derive_policy(DerivedNetworkPolicyFields {
        address_families: vec![
            AddressFamilyDisposition {
                family: AddressFamily::Ipv4,
                disposition: AddressFamilyDispositionValue::Enabled,
            },
            AddressFamilyDisposition {
                family: AddressFamily::Ipv6,
                disposition: AddressFamilyDispositionValue::Disabled,
            },
        ],
        default_actions: vec![DefaultAction {
            family: AddressFamily::Ipv4,
            action: PolicyAction::Allow,
        }],
        prefix_rules: [
            ([10, 0, 0, 0], 8),
            ([172, 16, 0, 0], 12),
            ([192, 168, 0, 0], 16),
        ]
        .into_iter()
        .map(|(address, prefix_length)| PrefixRule {
            family: AddressFamily::Ipv4,
            network: IpNetwork {
                address: IpAddress::V4(address),
                prefix_length,
            },
            action: PolicyAction::Deny,
        })
        .collect(),
        resolvers: vec![Resolver {
            family: AddressFamily::Ipv4,
            address: IpAddress::V4([10, 0, 2, 3]),
        }],
    })
    .map_err(std::io::Error::other)
}

fn probe_request(namespace: lazaret_core::identity::NamespaceIdentity) -> NetworkProbeRequest {
    NetworkProbeRequest {
        domain_id: DomainId::new([1; 32]),
        generation_id: GenerationId::new(7),
        enforcement_namespace: namespace,
        denied: vec![TcpProbeTarget {
            address: IpAddress::V4([10, 1, 2, 3]),
            port: 443,
        }],
        permitted: vec![TcpProbeTarget {
            address: IpAddress::V4([1, 1, 1, 1]),
            port: 443,
        }],
        dns: vec![DnsProbeTarget {
            resolver: IpAddress::V4([10, 0, 2, 3]),
            query_name: "example.com".into(),
        }],
    }
}

fn reachability() -> ReachabilityExpectation {
    ReachabilityExpectation {
        denied_destinations: vec![IpAddress::V4([10, 1, 2, 3])],
        permitted_destinations: vec![IpAddress::V4([1, 1, 1, 1])],
        dns_queries: vec![DnsQueryExpectation {
            resolver: IpAddress::V4([10, 0, 2, 3]),
            query_name: "example.com".into(),
        }],
    }
}

fn wait_for_path(
    path: &Path,
    process: &lazaret_linux::process::RetainedProcess,
) -> std::io::Result<()> {
    let deadline = Instant::now() + Duration::from_secs(30);
    while !path.exists() {
        if process.has_exited()? {
            return Err(std::io::Error::other(
                "network runtime exited before evidence",
            ));
        }
        if Instant::now() >= deadline {
            return Err(std::io::Error::other("network runtime evidence timed out"));
        }
        std::thread::sleep(Duration::from_millis(5));
    }
    Ok(())
}

fn cleanup_runtime<T>(
    runtime: lazaret_linux::process::RetainedProcess,
    error: std::io::Error,
) -> std::io::Result<T> {
    let _ = runtime.send_signal(Signal::KILL);
    let _ = child_broker().reap(runtime);
    Err(error)
}

fn parse_digest(value: &OsString) -> std::io::Result<lazaret_core::identity::DerivedPolicyDigest> {
    let value = value
        .to_str()
        .filter(|value| value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()))
        .ok_or_else(|| std::io::Error::other("network child digest is invalid"))?;
    let mut bytes = [0_u8; 32];
    for (index, byte) in bytes.iter_mut().enumerate() {
        *byte = u8::from_str_radix(&value[index * 2..index * 2 + 2], 16)
            .map_err(|_| std::io::Error::other("network child digest is invalid"))?;
    }
    Ok(lazaret_core::identity::DerivedPolicyDigest::new(bytes))
}

fn hex(bytes: &[u8]) -> String {
    use std::fmt::Write as _;
    bytes.iter().fold(String::new(), |mut output, byte| {
        write!(output, "{byte:02x}").unwrap();
        output
    })
}

fn prerequisite_failure(error: &std::io::Error) -> bool {
    matches!(
        error.kind(),
        std::io::ErrorKind::NotFound | std::io::ErrorKind::PermissionDenied
    ) || matches!(error.raw_os_error(), Some(1 | 13 | 95))
}

fn exit(result: std::io::Result<()>) -> ExitCode {
    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("network child failed: {error}");
            ExitCode::FAILURE
        }
    }
}