use std::collections::BTreeMap;
use std::fs::{self, File, OpenOptions};
use std::net::{Ipv4Addr, SocketAddrV4, TcpListener, TcpStream, ToSocketAddrs};
use std::os::linux::net::SocketAddrExt;
use std::os::unix::ffi::OsStringExt;
use std::os::unix::fs::symlink;
use std::os::unix::net::{SocketAddr, UnixListener, UnixStream};
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::sync::Arc;
use std::thread;
use std::time::Duration;

use lazaret_core::configuration::{
    ConfigurationLoadRequest, ResolvedDomainConfiguration, load_trusted_configuration,
};
use lazaret_core::identity::{
    BundleDigest, ClientId, ContainerId, CoordinatorId, GenerationId, HostPath, NetworkSidecarId,
    PidfsProcessIdentityV1, RunscBuildId, RunscDigest, SandboxId, SessionAnchorId, SessionId,
};
use lazaret_core::network::derive_ipv4_policy;
use lazaret_core::records::{
    Domain, DurableDecision, Generation, HostProcessOwnershipSlot, ParticipantRole,
    ParticipantState, ProcessOwnershipCommitAck, Session, SessionPhase, SessionTerminalEventV2,
    SessionTerminalMessage, TerminalBreak, TerminalInput, TerminalModesV2, TerminalResize,
};
use lazaret_host_coordinator::domain_store::{DomainStoreMutation, HostCoordinatorDomainStore};
use lazaret_host_coordinator::git_routing::GitRoutingCoordinator;
use lazaret_host_coordinator::network::NetworkCoordinator;
use lazaret_host_coordinator::runtime::{
    AwaitingAttachedRunsc, CommittedSandboxRuntime, SandboxCoordinator, SandboxIntent,
};
use lazaret_host_coordinator::terminal::{
    NoActivePrompt, RunscSessionTerminalFactory, SessionTerminalEngine, TerminalAcquisition,
    TerminalRoute,
};
use lazaret_linux::filesystem::{current_mount_identity, try_exclusive_flock};
use lazaret_linux::platform::{PlatformObservation, check_platform};
use lazaret_linux::process::{
    ProcessOwnershipContext, RetainedProcess, TargetCommand, child_broker,
    current_process_identity, read_start_time,
};
use lazaret_linux::user_namespace::UserNamespaceMap;
use lazaret_runtime_supervisor::sandbox::{
    NETWORK_NAMESPACE_PATH, RootProcess, RunscCommand, SandboxBuildRequest, SandboxBundle,
    build_sandbox_bundle, build_sandbox_bundle_with_workspace, default_sandbox_layout,
};
use lazaret_runtime_supervisor::supervision::{
    AttachedRunscSignal, RuntimeGitRoutingAuthority, RuntimeGitRoutingLaunch, RuntimeNetworkLaunch,
    RuntimeSupervisorLaunch,
};
use rustix::process::Signal;
use tempfile::TempDir;

const SKIP: u8 = 77;
const HOST_COORDINATOR_DRIVER: &str = "--runtime-composition-host-coordinator";
const SANDBOX_NETWORK_PROBE: &str = "--sandbox-network-probe";
const NETWORK_MATRIX: &str = "--network-matrix";
const HANDOFF_MATRIX: &str = "--handoff-matrix";
const SESSION_TERMINAL_MATRIX: &str = "--session-terminal-matrix";
const BENCHMARK_DRIVER: &str = "--runtime-composition-benchmark-command";
const BENCHMARK_GUEST: &str = "--runtime-composition-benchmark-guest";
const BENCHMARK_COMPLETE: &str = "benchmark-complete";
const BENCHMARK_ERROR: &str = "benchmark-error";

fn main() -> ExitCode {
    let arguments = std::env::args_os().skip(1).collect::<Vec<_>>();
    if let Some(outcome) = lazaret_host_coordinator::private_roles::REGISTRY.dispatch(&arguments) {
        return exit(outcome.result);
    }
    if arguments.first().and_then(|argument| argument.to_str()) == Some(BENCHMARK_GUEST) {
        return exit(run_benchmark_guest(&arguments[1..]));
    }
    if arguments.first().and_then(|argument| argument.to_str()) == Some(BENCHMARK_DRIVER) {
        return exit(run_benchmark_driver(&arguments[1..]));
    }
    if arguments.first().and_then(|argument| argument.to_str()) == Some(HOST_COORDINATOR_DRIVER) {
        return exit(run_host_coordinator_driver(&arguments[1..]));
    }
    if arguments.first().and_then(|argument| argument.to_str()) == Some(SANDBOX_NETWORK_PROBE) {
        return exit(run_sandbox_network_probe(&arguments[1..]));
    }
    if arguments.first().and_then(|argument| argument.to_str()) == Some(NETWORK_MATRIX) {
        return match run_network_matrix() {
            Ok(()) => {
                println!("PASS network composition matrix");
                ExitCode::SUCCESS
            }
            Err(MatrixError::Skip(reason)) => {
                println!("SKIP network composition: {reason}");
                ExitCode::from(SKIP)
            }
            Err(MatrixError::Failure(error)) => {
                eprintln!("network composition failed: {error}");
                ExitCode::FAILURE
            }
        };
    }
    if arguments.first().and_then(|argument| argument.to_str()) == Some(HANDOFF_MATRIX) {
        return match run_handoff_matrix() {
            Ok(()) => {
                println!("PASS final handoff replacement refusal");
                ExitCode::SUCCESS
            }
            Err(MatrixError::Skip(reason)) => {
                println!("SKIP final handoff replacement refusal: {reason}");
                ExitCode::from(SKIP)
            }
            Err(MatrixError::Failure(error)) => {
                eprintln!("final handoff replacement refusal failed: {error}");
                ExitCode::FAILURE
            }
        };
    }
    if arguments.first().and_then(|argument| argument.to_str()) == Some(SESSION_TERMINAL_MATRIX) {
        return match run_session_terminal_matrix() {
            Ok(()) => {
                println!("PASS Session Terminal composition matrix");
                ExitCode::SUCCESS
            }
            Err(MatrixError::Skip(reason)) => {
                println!("SKIP Session Terminal composition matrix: {reason}");
                ExitCode::from(SKIP)
            }
            Err(MatrixError::Failure(error)) => {
                eprintln!("Session Terminal composition matrix failed: {error}");
                ExitCode::FAILURE
            }
        };
    }
    match run_matrix() {
        Ok(()) => {
            println!("PASS foreground runtime composition and causal-loss matrix");
            ExitCode::SUCCESS
        }
        Err(MatrixError::Skip(reason)) => {
            println!("SKIP foreground runtime composition: {reason}");
            ExitCode::from(SKIP)
        }
        Err(MatrixError::Failure(error)) => {
            eprintln!("runtime composition failed: {error}");
            ExitCode::FAILURE
        }
    }
}

fn run_benchmark_driver(arguments: &[std::ffi::OsString]) -> std::io::Result<()> {
    let separator = arguments
        .iter()
        .position(|argument| argument == "--")
        .ok_or_else(|| std::io::Error::other("benchmark command separator is missing"))?;
    if separator != 8 || arguments.len() == separator + 1 {
        return Err(std::io::Error::other(
            "benchmark composition arguments are malformed",
        ));
    }
    let paths = arguments[..separator]
        .iter()
        .map(PathBuf::from)
        .collect::<Vec<_>>();
    let fixture = Fixture::build_benchmark(
        &paths[0],
        &paths[1],
        &paths[2],
        &paths[3],
        &paths[4],
        &paths[5],
        &paths[6],
        &paths[7],
        &arguments[separator + 1..],
    )
    .map_err(matrix_io)?;
    let coordinator = SandboxCoordinator::new(&fixture.store);
    let runtime = fixture.launch_committed().map_err(matrix_io)?;
    let results = fixture
        .benchmark_results
        .as_ref()
        .ok_or_else(|| std::io::Error::other("benchmark composition result authority is absent"))?;
    let complete = results.join(BENCHMARK_COMPLETE);
    let error = results.join(BENCHMARK_ERROR);
    for _ in 0..24_000 {
        if complete.is_file() || error.is_file() {
            break;
        }
        thread::sleep(Duration::from_millis(5));
    }
    let outcome = if complete.is_file() {
        Ok(())
    } else if error.is_file() {
        Err(std::io::Error::other(fs::read_to_string(&error)?))
    } else {
        Err(std::io::Error::other("benchmark command timed out"))
    };
    runtime
        .stop(&coordinator)
        .map_err(|error| std::io::Error::other(error.to_string()))?;
    outcome
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn run_benchmark_guest(arguments: &[std::ffi::OsString]) -> std::io::Result<()> {
    if arguments.len() < 3 {
        return Err(std::io::Error::other(
            "benchmark guest arguments are malformed",
        ));
    }
    let results = PathBuf::from(&arguments[0]);
    let output = std::process::Command::new(&arguments[1])
        .args(&arguments[2..])
        .output();
    match output {
        Ok(output) if output.status.success() => {
            fs::write(results.join(BENCHMARK_COMPLETE), b"ok")?;
        }
        Ok(output) => fs::write(
            results.join(BENCHMARK_ERROR),
            format!(
                "benchmark command exited with {}: {}",
                output.status,
                String::from_utf8_lossy(&output.stderr)
            ),
        )?,
        Err(error) => fs::write(results.join(BENCHMARK_ERROR), error.to_string())?,
    }
    loop {
        thread::park();
    }
}

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

enum MatrixError {
    Skip(String),
    Failure(String),
}

impl From<std::io::Error> for MatrixError {
    fn from(error: std::io::Error) -> Self {
        Self::Failure(error.to_string())
    }
}

fn run_matrix() -> Result<(), MatrixError> {
    require_matrix_prerequisites()?;
    lost_acknowledgment_refuses_target_work()?;
    replaced_handoff_authority_refuses_target_work()?;
    git_routing_clean_stop()?;
    normal_stop_and_stop_resume()?;
    attached_runsc_loss_stops_guest()?;
    runtime_supervisor_loss_stops_guest()?;
    host_coordinator_loss_stops_guest()
}

fn run_network_matrix() -> Result<(), MatrixError> {
    require_matrix_prerequisites()?;
    run_network_cases()
}

#[allow(clippy::too_many_lines)]
fn run_session_terminal_matrix() -> Result<(), MatrixError> {
    require_matrix_prerequisites()?;
    let fixture = Fixture::new()?;
    let runtime_domain = host_path(
        fixture
            .configuration
            .runtime_domain_root()
            .canonical_label(),
    );
    fs::copy(
        env!("CARGO_BIN_EXE_lazaret-host-coordinator"),
        runtime_domain
            .join("sandbox-control")
            .join(hex(ContainerId::new([6; 32]).as_bytes()))
            .join("bin/lazaret-sandbox-init"),
    )?;
    set_fixture_generation_decision(&fixture.store, DurableDecision::Preparing)?;
    let runtime = fixture.launch_committed()?;
    set_fixture_generation_decision(&fixture.store, DurableDecision::Published)?;
    let session_id = SessionId::new([31; 32]);
    let mut session = Session {
        generation_id: GenerationId::new(7),
        session_id,
        cli_id: ClientId::new([32; 32]),
        session_anchor_id: SessionAnchorId::new([33; 32]),
        session_phase: SessionPhase::Registering,
    };
    fixture
        .store
        .apply(&[DomainStoreMutation::PutSession(session.clone())])
        .map_err(failure)?;
    let second_session_id = SessionId::new([41; 32]);
    let mut second_session = Session {
        generation_id: GenerationId::new(7),
        session_id: second_session_id,
        cli_id: ClientId::new([42; 32]),
        session_anchor_id: SessionAnchorId::new([43; 32]),
        session_phase: SessionPhase::Registering,
    };
    fixture
        .store
        .apply(&[DomainStoreMutation::PutSession(second_session.clone())])
        .map_err(failure)?;
    second_session.session_phase = SessionPhase::Attached;
    fixture
        .store
        .apply(&[DomainStoreMutation::PutSession(second_session)])
        .map_err(failure)?;
    session.session_phase = SessionPhase::Attached;
    fixture
        .store
        .apply(&[DomainStoreMutation::PutSession(session.clone())])
        .map_err(failure)?;
    let command = fixture.bundle.command();
    let run = command
        .arguments()
        .iter()
        .position(|argument| argument == "run")
        .ok_or_else(|| failure("runsc command has no run action"))?;
    let workspace = host_path(fixture.configuration.workspace().canonical_label());
    let count = workspace.join("session-count");
    let finish = workspace.join("session-finish");
    let signals = workspace.join("session-signals");
    let second_count = workspace.join("session-second-count");
    let second_finish = workspace.join("session-second-finish");
    let second_signals = workspace.join("session-second-signals");
    let line = "x".repeat(1024);
    let script = format!(
        "test \"$PROFILE_SENTINEL\" = profile-value || exit 97; \
         command -v sleep >/dev/null || exit 96; \
         trap 'echo INT >> /workspace/session-signals' INT; \
         trap 'echo QUIT >> /workspace/session-signals' QUIT; \
         trap 'echo TSTP >> /workspace/session-signals' TSTP; \
         trap 'echo WINCH >> /workspace/session-signals' WINCH; \
         n=0; while test ! -e /workspace/session-finish; do \
         echo '{line}'; n=$((n+1)); echo $n > /workspace/session-count; done; exit 23"
    );
    let second_line = "y".repeat(1024);
    let second_script = format!(
        "trap 'echo INT >> /workspace/session-second-signals' INT; \
         trap 'echo QUIT >> /workspace/session-second-signals' QUIT; \
         trap 'echo TSTP >> /workspace/session-second-signals' TSTP; \
         trap 'echo WINCH >> /workspace/session-second-signals' WINCH; \
         n=0; while test ! -e /workspace/session-second-finish; do \
         echo '{second_line}'; n=$((n+1)); echo $n > /workspace/session-second-count; done; exit 42"
    );
    let request = TerminalAcquisition {
        domain_id: fixture.configuration.domain_id(),
        generation_id: GenerationId::new(7),
        coordinator_id: CoordinatorId::new([3; 32]),
        session_id,
        container_id: ContainerId::new([6; 32]),
        host_coordinator: current_process_identity()?,
        runsc_program: command.program().to_owned(),
        runsc_global_arguments: command.arguments()[..run].to_vec(),
        environment: fixture.configuration.ephemeral_environment().clone(),
        internal_pid_file: runtime_domain.join("session-internal.pid"),
        argv: vec![b"/bin/sh".to_vec(), b"-c".to_vec(), script.into_bytes()],
        initial_size: TerminalResize {
            session_id,
            rows: std::num::NonZeroU16::new(24).unwrap(),
            columns: std::num::NonZeroU16::new(80).unwrap(),
            pixel_width: 0,
            pixel_height: 0,
        },
        initial_modes: standard_terminal_modes(),
    };
    let mut engine = SessionTerminalEngine::new(
        Arc::clone(&fixture.store),
        Arc::new(RunscSessionTerminalFactory),
        Arc::new(NoActivePrompt),
    );
    engine.acquire(&request).map_err(failure)?;
    engine
        .acquire(&TerminalAcquisition {
            session_id: second_session_id,
            internal_pid_file: runtime_domain.join("session-second-internal.pid"),
            argv: vec![
                b"/bin/sh".to_vec(),
                b"-c".to_vec(),
                second_script.into_bytes(),
            ],
            ..request.clone()
        })
        .map_err(failure)?;
    wait_for_path(&count)?;
    wait_for_path(&second_count)?;
    let running = fixture
        .store
        .snapshot()
        .map_err(failure)?
        .session_terminals
        .into_iter()
        .find(|terminal| terminal.session_id() == session_id)
        .ok_or_else(|| failure("running Session Terminal record is absent"))?;
    let guest = running
        .guest_process()
        .ok_or_else(|| failure("running Session Terminal guest identity is absent"))?;
    if guest.leader_pid.get() != guest.process_group_id.get() {
        return Err(failure("guest leader does not own its process group"));
    }
    engine
        .handle(
            GenerationId::new(7),
            SessionTerminalEventV2::Resize(TerminalResize {
                session_id,
                rows: std::num::NonZeroU16::new(41).unwrap(),
                columns: std::num::NonZeroU16::new(121).unwrap(),
                pixel_width: 0,
                pixel_height: 0,
            }),
        )
        .map_err(failure)?;
    for byte in [3, 28, 26] {
        engine
            .handle(
                GenerationId::new(7),
                SessionTerminalEventV2::Input(TerminalInput {
                    session_id,
                    bytes: vec![byte],
                }),
            )
            .map_err(failure)?;
    }
    engine
        .handle(
            GenerationId::new(7),
            SessionTerminalEventV2::Break(TerminalBreak { session_id }),
        )
        .map_err(failure)?;
    wait_for_signal_lines(&mut engine, session_id, &signals, 4)?;
    if second_signals.exists() {
        return Err(failure("terminal signals crossed Session IDs"));
    }
    let mut forwarded_before_prompt = 0_usize;
    for _ in 0..2_000 {
        let messages = engine
            .poll(session_id, GenerationId::new(7))
            .map_err(failure)?;
        if messages.is_empty() {
            if forwarded_before_prompt > 0 {
                break;
            }
            thread::sleep(Duration::from_millis(5));
            continue;
        }
        for message in messages {
            if let SessionTerminalMessage::Output(chunk) = message {
                forwarded_before_prompt += chunk.bytes.len();
            }
        }
    }
    if forwarded_before_prompt == 0 {
        return Err(failure(
            "sandbox route did not forward buffered output before Prompt",
        ));
    }
    session.session_phase = SessionPhase::Prompting;
    fixture
        .store
        .apply(&[DomainStoreMutation::PutSession(session.clone())])
        .map_err(failure)?;
    engine
        .set_route(session_id, GenerationId::new(7), TerminalRoute::Prompt)
        .map_err(failure)?;
    thread::sleep(Duration::from_millis(300));
    let before = fs::read_to_string(&count)?;
    let second_before = fs::read_to_string(&second_count)?;
    for _ in 0..60 {
        engine
            .poll(second_session_id, GenerationId::new(7))
            .map_err(failure)?;
        thread::sleep(Duration::from_millis(5));
    }
    let after = fs::read_to_string(&count)?;
    let second_after = fs::read_to_string(&second_count)?;
    if before != after {
        return Err(failure("prompt route did not apply PTY backpressure"));
    }
    if second_before == second_after {
        return Err(failure("one Prompt paused another Session Terminal"));
    }
    session.session_phase = SessionPhase::Attached;
    fixture
        .store
        .apply(&[DomainStoreMutation::PutSession(session)])
        .map_err(failure)?;
    engine
        .set_route(session_id, GenerationId::new(7), TerminalRoute::Sandbox)
        .map_err(failure)?;
    let mut output = 0_usize;
    for _ in 0..2_000 {
        for message in engine
            .poll(session_id, GenerationId::new(7))
            .map_err(failure)?
        {
            if let SessionTerminalMessage::Output(chunk) = message {
                output += chunk.bytes.len();
            }
        }
        if fs::read_to_string(&count).ok().as_deref() != Some(after.as_str()) {
            break;
        }
        thread::sleep(Duration::from_millis(5));
    }
    if output == 0 || fs::read_to_string(&count).ok().as_deref() == Some(after.as_str()) {
        return Err(failure("sandbox route did not resume ordered PTY output"));
    }
    fs::write(&finish, b"finish")?;
    let mut exit = None;
    for _ in 0..2_000 {
        for message in engine
            .poll(session_id, GenerationId::new(7))
            .map_err(failure)?
        {
            if let SessionTerminalMessage::Exit(outcome) = message {
                exit = Some(outcome);
            }
        }
        if exit.is_some() {
            break;
        }
        thread::sleep(Duration::from_millis(5));
    }
    if exit.and_then(lazaret_core::records::SessionExit::exit_code) != Some(23)
        || engine.contains(session_id)
        || !engine.contains(second_session_id)
    {
        return Err(failure(
            "Session Terminal exit or no-reattachment cleanup failed",
        ));
    }
    let second_before_first_exit = fs::read_to_string(&second_count)?;
    for _ in 0..200 {
        engine
            .poll(second_session_id, GenerationId::new(7))
            .map_err(failure)?;
        if fs::read_to_string(&second_count).ok().as_deref()
            != Some(second_before_first_exit.as_str())
        {
            break;
        }
        thread::sleep(Duration::from_millis(5));
    }
    if fs::read_to_string(&second_count).ok().as_deref() == Some(second_before_first_exit.as_str())
    {
        return Err(failure("first Session exit paused the second Session"));
    }
    fs::write(&second_finish, b"finish")?;
    let mut second_exit = None;
    for _ in 0..2_000 {
        for message in engine
            .poll(second_session_id, GenerationId::new(7))
            .map_err(failure)?
        {
            if let SessionTerminalMessage::Exit(outcome) = message {
                second_exit = Some(outcome);
            }
        }
        if second_exit.is_some() {
            break;
        }
        thread::sleep(Duration::from_millis(5));
    }
    if second_exit.and_then(lazaret_core::records::SessionExit::exit_code) != Some(42)
        || engine.contains(second_session_id)
    {
        return Err(failure("second Session exit routing or cleanup failed"));
    }
    runtime
        .stop(&SandboxCoordinator::new(&fixture.store))
        .map_err(failure)
}

fn set_fixture_generation_decision(
    store: &HostCoordinatorDomainStore,
    decision: DurableDecision,
) -> Result<(), MatrixError> {
    let mut generation = store
        .snapshot()
        .map_err(failure)?
        .generations
        .into_iter()
        .find(|generation| generation.generation_id == GenerationId::new(7))
        .ok_or_else(|| failure("Session Terminal fixture generation is absent"))?;
    generation.durable_decision = decision;
    store
        .apply(&[DomainStoreMutation::PutGeneration(generation)])
        .map_err(failure)
}

fn wait_for_path(path: &Path) -> Result<(), MatrixError> {
    for _ in 0..2_000 {
        if path.is_file() {
            return Ok(());
        }
        thread::sleep(Duration::from_millis(5));
    }
    Err(failure("Session Terminal guest did not become ready"))
}

fn standard_terminal_modes() -> TerminalModesV2 {
    TerminalModesV2 {
        input_modes: 1280,
        output_modes: 5,
        control_modes: 983_231,
        local_modes: 35_387,
        line_discipline: 0,
        input_speed: 38_400,
        output_speed: 38_400,
        interrupt: 3,
        quit: 28,
        erase: 127,
        kill: 21,
        eof: 4,
        time: 0,
        minimum: 1,
        switch: 0,
        start: 17,
        stop: 19,
        suspend: 26,
        end_of_line: 0,
        reprint: 18,
        discard: 15,
        word_erase: 23,
        literal_next: 22,
        end_of_line_2: 0,
    }
}

fn wait_for_signal_lines(
    engine: &mut SessionTerminalEngine,
    session_id: SessionId,
    path: &Path,
    expected: usize,
) -> Result<(), MatrixError> {
    for _ in 0..2_000 {
        engine
            .poll(session_id, GenerationId::new(7))
            .map_err(failure)?;
        if fs::read_to_string(path)
            .ok()
            .is_some_and(|value| value.lines().count() >= expected)
        {
            return Ok(());
        }
        thread::sleep(Duration::from_millis(5));
    }
    Err(failure(
        "Session Terminal guest did not observe exact signals",
    ))
}

#[path = "runtime-composition/fixtures.rs"]
mod fixtures;
#[path = "runtime-composition/handoff.rs"]
mod handoff;
#[path = "runtime-composition/network.rs"]
mod network;

use fixtures::{Fixture, NetworkFixture};
use handoff::{
    attached_runsc_loss_stops_guest, failure, git_routing_clean_stop,
    host_coordinator_loss_stops_guest, lost_acknowledgment_refuses_target_work, matrix_io,
    normal_stop_and_stop_resume, replaced_handoff_authority_refuses_target_work,
    require_matrix_prerequisites, run_handoff_matrix, run_host_coordinator_driver,
    run_network_cases, runtime_supervisor_loss_stops_guest,
};
use network::{
    hex, host_path, lock_file, networked_sandbox_view_isolated, run_sandbox_network_probe,
    sandbox_participant, stage_benchmark_rootfs_and_control, stage_rootfs_and_control,
    unreachable_permitted_destination_does_not_block_start,
};