//! Production-facing kboxlike runtime assembly helpers.
//!
//! These are the reusable pieces proven by the live Chromium restore smoke:
//! materialize stopped-process memory and synthesize a replayable fd model from
//! a captured process tree. Higher-level backend wiring should call through
//! here instead of test-only helpers.

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::fd::{FdError, KboxlikeFdSystem, ProcessId, ShadowBacking, ShadowKind, ShadowObject};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::restore::{
    restore_stopped_kboxlike_snapshot_with_replacements, KboxlikeFdKernelReplayInputs,
    KboxlikeSnapshotRestoreResult, KboxlikeTraceeRestoreInputs,
};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::shadow::HostShadowRebuilder;
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::snapshot::{
    apply_fd_shadow_backing_plans, capture_dump_candidate_memory, capture_shared_object_memory,
    classify_memory_restore_actions, materialize_memory_restore_actions,
    plan_epoll_shadow_backings, plan_eventfd_shadow_backings, plan_fd_queued_data_replay,
    plan_fd_scm_rights_replay, plan_fd_socket_options_replay, plan_fd_socket_shutdown_replay,
    plan_live_socket_shadow_backings, plan_pipe_shadow_backings, plan_socketpair_shadow_backings,
    plan_timerfd_shadow_backings, KboxlikeFdTargetKind, KboxlikeMaterializedMemoryRestore,
    KboxlikeMaterializedMemoryRestoreAction, KboxlikeMemoryDumpCapture,
    KboxlikeMemoryRestoreAction, KboxlikeSharedMemoryDump, KboxlikeStoppedTraceeSnapshot,
};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use super::tracee::{
    ProcessVmMemoryWriter, PtraceRegisterWriter, PtraceScratchSyscallExecutor,
    RestoredTraceeResumer, StoppedTraceeReplacementFactory, TraceeRestoreAccess,
    TraceeRestoreAccessFactory,
};
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use std::collections::BTreeSet;
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
use std::path::PathBuf;

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Debug)]
pub(crate) struct KboxlikeRuntimeMemoryMaterialization {
    pub(crate) materialized: KboxlikeMaterializedMemoryRestore,
    pub(crate) shared_dumps: Vec<KboxlikeSharedMemoryDump>,
    pub(crate) unique_processes: usize,
    pub(crate) duplicate_map_sets: usize,
    pub(crate) actions: usize,
    pub(crate) dump_anonymous_private: usize,
    pub(crate) dump_private_file_cow: usize,
    pub(crate) remap_clean_file: usize,
    pub(crate) recreate_shared_object: usize,
    pub(crate) skip_kernel_synthetic: usize,
    pub(crate) skip_inaccessible: usize,
    pub(crate) byte_backed_bytes: u64,
    pub(crate) byte_dump_bytes: u64,
    pub(crate) raw_dump_bytes: u64,
    pub(crate) shared_bytes: u64,
    pub(crate) shared_dump_bytes: u64,
    pub(crate) shared_duplicate_extents: usize,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Debug)]
pub(crate) struct KboxlikeRuntimeRestoreOptions {
    pub(crate) shadow_root: PathBuf,
    pub(crate) queued_data_max_bytes_per_fd: usize,
    pub(crate) scm_rights_max_data_bytes: usize,
    pub(crate) scm_rights_max_fds: usize,
    pub(crate) clear_existing_replacement_mappings: bool,
    pub(crate) trace_label: Option<&'static str>,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
impl Default for KboxlikeRuntimeRestoreOptions {
    fn default() -> Self {
        Self {
            shadow_root: std::env::temp_dir(),
            queued_data_max_bytes_per_fd: 4096,
            scm_rights_max_data_bytes: 256,
            scm_rights_max_fds: 16,
            clear_existing_replacement_mappings: true,
            trace_label: None,
        }
    }
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Debug)]
pub(crate) struct KboxlikeRuntimeFdReplaySummary {
    pub(crate) epoll: usize,
    pub(crate) queued_data: usize,
    pub(crate) socket_options: usize,
    pub(crate) socket_shutdown: usize,
    pub(crate) scm_rights: usize,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
#[derive(Debug)]
pub(crate) struct KboxlikeRuntimeRestoreResult {
    pub(crate) restore: KboxlikeSnapshotRestoreResult,
    pub(crate) fd_replay: KboxlikeRuntimeFdReplaySummary,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
pub(crate) fn restore_runtime_snapshot_with_replacements(
    snapshot: &KboxlikeStoppedTraceeSnapshot,
    current_exe: &str,
    options: KboxlikeRuntimeRestoreOptions,
    replacement_factory: &mut impl StoppedTraceeReplacementFactory,
    resumer: &mut impl RestoredTraceeResumer,
) -> Result<KboxlikeRuntimeRestoreResult, FdError> {
    let memory = materialize_runtime_memory(snapshot)?;
    if let Some(label) = options.trace_label {
        memory.trace_summary(snapshot, &format!("{label}_MEMORY_PLAN"));
    }

    let fd_system = synthesize_runtime_fd_system(snapshot, current_exe)?;
    let epoll_plans = plan_epoll_shadow_backings(&snapshot.fd_tables, &fd_system);
    let queued_data_plans = plan_fd_queued_data_replay(
        &snapshot.fd_tables,
        &fd_system,
        options.queued_data_max_bytes_per_fd,
    );
    let socket_options_plans = plan_fd_socket_options_replay(&snapshot.fd_tables, &fd_system);
    let socket_shutdown_plans = plan_fd_socket_shutdown_replay(&snapshot.fd_tables, &fd_system);
    let scm_rights_plans = plan_fd_scm_rights_replay(
        &snapshot.fd_tables,
        &fd_system,
        options.scm_rights_max_data_bytes,
        options.scm_rights_max_fds,
    );
    if let Some(label) = options.trace_label {
        eprintln!(
            "{label}_FD_PLAN tracees={} fd_tables={} epoll={} queued={} socket_options={} socket_shutdown={} scm_rights={}",
            snapshot.threads.len(),
            snapshot.fd_tables.len(),
            epoll_plans.len(),
            queued_data_plans.len(),
            socket_options_plans.len(),
            socket_shutdown_plans.len(),
            scm_rights_plans.len(),
        );
    }

    let syscall_page = choose_restore_syscall_page(&memory.materialized)?;
    replacement_factory.prepare_syscall_page(syscall_page)?;
    let mut rebuilder = HostShadowRebuilder::new(options.shadow_root);
    let mut access_factory = ScratchTraceeRestoreAccessFactory {
        syscall_page,
        clear_existing_mappings: options.clear_existing_replacement_mappings,
        cleared_pids: BTreeSet::new(),
    };
    let restore = restore_stopped_kboxlike_snapshot_with_replacements(
        fd_system.snapshot(),
        KboxlikeTraceeRestoreInputs {
            memory_restore: &memory.materialized,
            threads: &snapshot.threads,
            kernel_states: &snapshot.thread_kernel_states,
            shared_dumps: &memory.shared_dumps,
            supervisor_pid: std::process::id() as i32,
        },
        KboxlikeFdKernelReplayInputs {
            epoll_plans: &epoll_plans,
            socket_options_plans: &socket_options_plans,
            queued_data_plans: &queued_data_plans,
            socket_shutdown_plans: &socket_shutdown_plans,
            scm_rights_plans: &scm_rights_plans,
        },
        &mut rebuilder,
        replacement_factory,
        &mut access_factory,
        resumer,
    )?;

    Ok(KboxlikeRuntimeRestoreResult {
        restore,
        fd_replay: KboxlikeRuntimeFdReplaySummary {
            epoll: epoll_plans.len(),
            queued_data: queued_data_plans.len(),
            socket_options: socket_options_plans.len(),
            socket_shutdown: socket_shutdown_plans.len(),
            scm_rights: scm_rights_plans.len(),
        },
    })
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
impl KboxlikeRuntimeMemoryMaterialization {
    pub(crate) fn trace_summary(&self, snapshot: &KboxlikeStoppedTraceeSnapshot, label: &str) {
        eprintln!(
            "{label} tracees={} unique_processes={} duplicate_map_sets={} fd_tables={} actions={} dump_anon={} dump_cow={} remap_clean={} shared={} skip_kernel={} skip_inaccessible={} byte_action_bytes={} byte_dump_bytes={} raw_dump_bytes={} shared_bytes={} shared_dump_bytes={} shared_duplicate_extents={} materialized_actions={}",
            snapshot.threads.len(),
            self.unique_processes,
            self.duplicate_map_sets,
            snapshot.fd_tables.len(),
            self.actions,
            self.dump_anonymous_private,
            self.dump_private_file_cow,
            self.remap_clean_file,
            self.recreate_shared_object,
            self.skip_kernel_synthetic,
            self.skip_inaccessible,
            self.byte_backed_bytes,
            self.byte_dump_bytes,
            self.raw_dump_bytes,
            self.shared_bytes,
            self.shared_dump_bytes,
            self.shared_duplicate_extents,
            self.materialized.actions.len(),
        );
    }
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
pub(crate) fn materialize_runtime_memory(
    snapshot: &KboxlikeStoppedTraceeSnapshot,
) -> Result<KboxlikeRuntimeMemoryMaterialization, FdError> {
    let mut seen_model_pids = BTreeSet::new();
    let unique_memory_maps = snapshot
        .memory_maps
        .iter()
        .filter(|tracee_maps| seen_model_pids.insert(tracee_maps.model_pid))
        .cloned()
        .collect::<Vec<_>>();
    let duplicate_map_sets = snapshot.memory_maps.len() - unique_memory_maps.len();
    let plan = classify_memory_restore_actions(&unique_memory_maps);

    let mut dump_anonymous_private = 0usize;
    let mut dump_private_file_cow = 0usize;
    let mut remap_clean_file = 0usize;
    let mut recreate_shared_object = 0usize;
    let mut skip_kernel_synthetic = 0usize;
    let mut skip_inaccessible = 0usize;
    let mut byte_backed_bytes = 0u64;
    let mut shared_bytes = 0u64;
    for action in &plan.actions {
        match action {
            KboxlikeMemoryRestoreAction::DumpAnonymousPrivate { len, .. } => {
                dump_anonymous_private += 1;
                byte_backed_bytes += *len;
            }
            KboxlikeMemoryRestoreAction::DumpPrivateFileCow { len, .. } => {
                dump_private_file_cow += 1;
                byte_backed_bytes += *len;
            }
            KboxlikeMemoryRestoreAction::RemapCleanFile { .. } => {
                remap_clean_file += 1;
            }
            KboxlikeMemoryRestoreAction::RecreateSharedObject { len, .. } => {
                recreate_shared_object += 1;
                shared_bytes += *len;
            }
            KboxlikeMemoryRestoreAction::SkipKernelSynthetic { .. } => {
                skip_kernel_synthetic += 1;
            }
            KboxlikeMemoryRestoreAction::SkipInaccessible { .. } => {
                skip_inaccessible += 1;
            }
        }
    }

    let raw_dumps = capture_dump_candidate_memory(&KboxlikeStoppedTraceeSnapshot {
        threads: Vec::new(),
        thread_kernel_states: Vec::new(),
        memory_maps: unique_memory_maps,
        fd_tables: Vec::new(),
    });
    let byte_dumps = filter_byte_backed_memory_dumps(&plan.actions, &raw_dumps);
    let shared_dumps = capture_shared_object_memory(&plan);
    if !byte_dumps.failures.is_empty() || !shared_dumps.failures.is_empty() {
        eprintln!(
            "KBOXLIVE_MEMORY_PLAN_CAPTURE_FAILURE byte_failures={} byte_failed_bytes={} shared_failures={} shared_failed_bytes={}",
            byte_dumps.failures.len(),
            byte_dumps.failed_bytes(),
            shared_dumps.failures.len(),
            shared_dumps.failed_bytes(),
        );
        return Err(FdError::TraceeInstallFailed(
            "kboxlike runtime memory capture had restore-required failures",
        ));
    }
    let materialized = materialize_memory_restore_actions(&plan, &byte_dumps)
        .map_err(|_| FdError::TraceeInstallFailed("kboxlike runtime memory materialize failed"))?;
    let shared_dump_bytes = shared_dumps.captured_bytes();
    let shared_duplicate_extents = shared_dumps.skipped_duplicate_extents;

    Ok(KboxlikeRuntimeMemoryMaterialization {
        materialized,
        shared_dumps: shared_dumps.dumps,
        unique_processes: seen_model_pids.len(),
        duplicate_map_sets,
        actions: plan.actions.len(),
        dump_anonymous_private,
        dump_private_file_cow,
        remap_clean_file,
        recreate_shared_object,
        skip_kernel_synthetic,
        skip_inaccessible,
        byte_backed_bytes,
        byte_dump_bytes: byte_dumps.captured_bytes(),
        raw_dump_bytes: raw_dumps.captured_bytes(),
        shared_bytes,
        shared_dump_bytes,
        shared_duplicate_extents,
    })
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
pub(crate) fn synthesize_runtime_fd_system(
    snapshot: &KboxlikeStoppedTraceeSnapshot,
    current_exe: &str,
) -> Result<KboxlikeFdSystem, FdError> {
    let mut fd_system = KboxlikeFdSystem::new();
    let mut model_pids = snapshot
        .fd_tables
        .iter()
        .map(|table| table.model_pid)
        .chain(snapshot.threads.iter().map(|thread| thread.model_pid))
        .collect::<BTreeSet<_>>();
    for model_pid in model_pids.split_off(&0) {
        let spawned = fd_system.spawn_init_with_exe(current_exe);
        if spawned != model_pid {
            return Err(FdError::TraceeInstallFailed(
                "captured model pids were not contiguous from init",
            ));
        }
    }

    for table in &snapshot.fd_tables {
        for fd in &table.fds {
            if fd.kind == KboxlikeFdTargetKind::RegularPath
                && fd.target.starts_with('/')
                && !fd.target.ends_with(" (deleted)")
            {
                fd_system.insert_shadow_fd_with_backing(
                    table.model_pid,
                    fd.fd,
                    None,
                    ShadowObject {
                        kind: ShadowKind::HostPassthrough,
                        supervisor_fd: fd.fd,
                    },
                    ShadowBacking::HostPath {
                        path: fd.target.clone(),
                        writable: captured_fd_is_writable(fd.flags),
                    },
                    fd.cloexec,
                )?;
            } else {
                fd_system.insert_plain_fd(table.model_pid, fd.fd, None, fd.cloexec)?;
            }
        }
    }

    let mut fd_plans = plan_pipe_shadow_backings(&snapshot.fd_tables);
    fd_plans.extend(plan_eventfd_shadow_backings(&snapshot.fd_tables));
    fd_plans.extend(plan_timerfd_shadow_backings(&snapshot.fd_tables));
    fd_plans.extend(plan_live_socket_shadow_backings(&snapshot.fd_tables));
    fd_plans.extend(plan_socketpair_shadow_backings(&snapshot.fd_tables));
    apply_fd_shadow_backing_plans(&mut fd_system, &fd_plans)?;
    let epoll_plans = plan_epoll_shadow_backings(&snapshot.fd_tables, &fd_system);
    let epoll_fd_plans = epoll_plans
        .iter()
        .map(|plan| plan.fd_plan.clone())
        .collect::<Vec<_>>();
    apply_fd_shadow_backing_plans(&mut fd_system, &epoll_fd_plans)?;

    Ok(fd_system)
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn captured_fd_is_writable(flags: Option<u64>) -> bool {
    flags.is_some_and(|flags| flags & libc::O_ACCMODE as u64 != libc::O_RDONLY as u64)
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn byte_backed_memory_action_key(
    action: &KboxlikeMemoryRestoreAction,
) -> Option<(i32, ProcessId, u64)> {
    match action {
        KboxlikeMemoryRestoreAction::DumpAnonymousPrivate {
            host_pid,
            model_pid,
            start,
            ..
        }
        | KboxlikeMemoryRestoreAction::DumpPrivateFileCow {
            host_pid,
            model_pid,
            start,
            ..
        } => Some((*host_pid, *model_pid, *start)),
        _ => None,
    }
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn filter_byte_backed_memory_dumps(
    plan_actions: &[KboxlikeMemoryRestoreAction],
    dumps: &KboxlikeMemoryDumpCapture,
) -> KboxlikeMemoryDumpCapture {
    let required = plan_actions
        .iter()
        .filter_map(byte_backed_memory_action_key)
        .collect::<BTreeSet<_>>();

    KboxlikeMemoryDumpCapture {
        dumps: dumps
            .dumps
            .iter()
            .filter(|dump| required.contains(&(dump.host_pid, dump.model_pid, dump.start)))
            .cloned()
            .collect(),
        failures: dumps
            .failures
            .iter()
            .filter(|failure| {
                required.contains(&(failure.host_pid, failure.model_pid, failure.start))
            })
            .cloned()
            .collect(),
    }
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn materialized_memory_action_range(
    action: &KboxlikeMaterializedMemoryRestoreAction,
) -> Option<(u64, u64)> {
    match action {
        KboxlikeMaterializedMemoryRestoreAction::ByteBacked { start, bytes, .. } => {
            Some((*start, bytes.len() as u64))
        }
        KboxlikeMaterializedMemoryRestoreAction::RemapCleanFile { start, len, .. }
        | KboxlikeMaterializedMemoryRestoreAction::RecreateSharedObject { start, len, .. }
        | KboxlikeMaterializedMemoryRestoreAction::SkipKernelSynthetic { start, len, .. }
        | KboxlikeMaterializedMemoryRestoreAction::SkipInaccessible { start, len, .. } => {
            Some((*start, *len))
        }
    }
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn ranges_overlap(start_a: u64, len_a: u64, start_b: u64, len_b: u64) -> bool {
    let end_a = start_a.saturating_add(len_a);
    let end_b = start_b.saturating_add(len_b);
    start_a < end_b && start_b < end_a
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn choose_restore_syscall_page(
    restore: &KboxlikeMaterializedMemoryRestore,
) -> Result<u64, FdError> {
    const PAGE_LEN: u64 = 8192;
    const CANDIDATES: &[u64] = &[
        0x7000_0000_0000,
        0x6800_0000_0000,
        0x6000_0000_0000,
        0x5800_0000_0000,
        0x5000_0000_0000,
    ];

    CANDIDATES
        .iter()
        .copied()
        .find(|candidate| {
            restore.actions.iter().all(|action| {
                materialized_memory_action_range(action)
                    .is_none_or(|(start, len)| !ranges_overlap(*candidate, PAGE_LEN, start, len))
            })
        })
        .ok_or(FdError::TraceeInstallFailed(
            "no scratch syscall page outside restore ranges",
        ))
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
struct ScratchTraceeRestoreAccessFactory {
    syscall_page: u64,
    clear_existing_mappings: bool,
    cleared_pids: BTreeSet<i32>,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
impl TraceeRestoreAccessFactory for ScratchTraceeRestoreAccessFactory {
    type Writer = ProcessVmMemoryWriter;
    type Executor = PtraceScratchSyscallExecutor;
    type RegisterWriter = PtraceRegisterWriter;

    fn access_for_restore(
        &mut self,
        host_pid: i32,
    ) -> Result<TraceeRestoreAccess<Self::Writer, Self::Executor, Self::RegisterWriter>, FdError>
    {
        let mut executor = PtraceScratchSyscallExecutor::new(host_pid, self.syscall_page)?;
        if self.clear_existing_mappings && self.cleared_pids.insert(host_pid) {
            let unmapped = executor.unmap_existing_mappings_except_scratch()?;
            if std::env::var_os("SUPERMACHINE_KBOXLIKE_TRACE_MEMORY_RESTORE").is_some() {
                eprintln!(
                    "KBOXLIVE_REPLACEMENT_CLEARED pid={} unmapped={unmapped}",
                    host_pid,
                );
            }
        }
        Ok(TraceeRestoreAccess {
            writer: ProcessVmMemoryWriter::new(host_pid)?,
            executor,
            register_writer: PtraceRegisterWriter::new(host_pid)?,
        })
    }
}
