use std::collections::BTreeMap;
use std::fs::{self, File, OpenOptions};
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::os::unix::fs::{MetadataExt, PermissionsExt};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use lazaret_core::configuration::{ConfigurationLoadRequest, load_trusted_configuration};
use lazaret_core::identity::{
    ContainerId, CoordinatorId, DeviceId, GenerationId, HostPath, IdMapDigest, InodeId, PolicyId,
    PreparedArtifactId, SandboxId, SelectedManifestDigest, SourceIdentity,
};
use lazaret_core::records::{ArtifactState, Domain, DurableDecision, Generation, PreparedArtifact};
use lazaret_host_coordinator::cold_start::{ColdStart, ColdStartInputs};
use lazaret_host_coordinator::domain_store::HostCoordinatorDomainStore;
use lazaret_linux::filesystem::try_exclusive_flock;
use lazaret_linux::user_namespace::{
    REQUIRED_SUBORDINATE_IDS, UserNamespaceIdRange, UserNamespaceMap,
};
use lazaret_runtime_supervisor::sandbox::{RootProcess, default_sandbox_layout};
use lazaret_state_anchor::StateAnchorLayout;
use lazaret_storage_catalog::Catalog;
use rustix::process::{getgid, getuid};

#[test]
fn validated_inputs_stage_one_complete_plan_and_drop_exact_runtime_artifacts() {
    let fixture = ColdStartFixture::new();
    let staged = ColdStart::stage(fixture.inputs(&fixture.escalation_client)).unwrap();
    let bundle = fixture.bundle();
    let control = fixture.control();
    let overlay = fixture.overlay();
    assert_private_artifacts(&bundle, &control);
    assert!(overlay.is_dir());
    drop(staged);
    assert!(!bundle.exists());
    assert!(!control.exists());
    assert!(!overlay.exists());
    assert!(ColdStart::stage(fixture.inputs(&fixture.invalid_escalation_client)).is_err());
    assert!(!bundle.exists());
    assert!(!control.exists());
}

#[test]
fn non_git_workspace_stages_without_a_routing_runtime() {
    let fixture = ColdStartFixture::new();

    let staged = ColdStart::stage(fixture.inputs(&fixture.escalation_client)).unwrap();
    assert!(!fixture.runtime_domain.join("git-routing").exists());
    drop(staged);
}

#[test]
fn production_stage_composes_derived_resolver_outside_the_prepared_lower() {
    let fixture = ColdStartFixture::new();
    let image_resolver = fixture.prepared_lower.join("etc/resolv.conf");
    let upper_resolver = fixture
        .persistent_domain
        .join("state-anchor/upper/etc/resolv.conf");
    let empty_runtime_resolver = StateAnchorLayout::create(&fixture.configuration).unwrap();
    assert_eq!(
        fs::read(empty_runtime_resolver.runtime_files().join("resolv.conf")).unwrap(),
        b""
    );
    drop(empty_runtime_resolver);

    let staged = ColdStart::stage(fixture.inputs(&fixture.escalation_client)).unwrap();

    assert_eq!(fs::read(&image_resolver).unwrap(), b"image resolver\n");
    assert!(!upper_resolver.exists());
    assert_eq!(
        fs::read(
            fixture
                .runtime_domain
                .join("state-anchor/runtime-files/resolv.conf")
        )
        .unwrap(),
        b"nameserver 10.137.54.3\n"
    );
    assert_eq!(
        fs::read(
            fixture
                .runtime_domain
                .join("state-anchor/runtime-files/hosts")
        )
        .unwrap(),
        b"127.0.0.1 localhost\n"
    );
    assert_eq!(
        fs::read(
            fixture
                .runtime_domain
                .join("state-anchor/runtime-files/hostname")
        )
        .unwrap(),
        b"lazaret\n"
    );
    let specification: serde_json::Value =
        serde_json::from_slice(&fs::read(fixture.bundle().join("config.json")).unwrap()).unwrap();
    let resolver_mount = specification["mounts"]
        .as_array()
        .unwrap()
        .iter()
        .find(|mount| mount["destination"] == "/etc/resolv.conf")
        .unwrap();
    assert_eq!(
        resolver_mount["source"],
        fixture
            .runtime_domain
            .join("state-anchor/runtime-files/resolv.conf")
            .display()
            .to_string()
    );
    drop(staged);
}

#[test]
fn cleanup_leaves_a_replacement_directory_untouched() {
    let fixture = ColdStartFixture::new();
    let staged = ColdStart::stage(fixture.inputs(&fixture.escalation_client)).unwrap();
    let bundle = fixture.bundle();
    let moved = bundle.with_extension("moved");
    fs::rename(&bundle, &moved).unwrap();
    fs::create_dir(&bundle).unwrap();
    fs::write(bundle.join("replacement"), b"trusted replacement").unwrap();
    drop(staged);
    assert_eq!(
        fs::read(bundle.join("replacement")).unwrap(),
        b"trusted replacement"
    );
    assert!(moved.is_dir());
}

#[test]
fn fixed_runtime_parents_refuse_symlinks_and_wrong_modes() {
    let fixture = ColdStartFixture::new();
    let bundles = fixture.runtime_domain.join("bundles");
    fs::create_dir(&bundles).unwrap();
    fs::set_permissions(&bundles, fs::Permissions::from_mode(0o755)).unwrap();
    assert!(ColdStart::stage(fixture.inputs(&fixture.escalation_client)).is_err());

    fs::remove_dir(&bundles).unwrap();
    let controls = fixture.runtime_domain.join("sandbox-control");
    let redirect = fixture.tree.path().join("redirect");
    fs::create_dir(&redirect).unwrap();
    std::os::unix::fs::symlink(&redirect, &controls).unwrap();
    assert!(ColdStart::stage(fixture.inputs(&fixture.escalation_client)).is_err());
    assert!(redirect.read_dir().unwrap().next().is_none());
}

#[test]
fn replaced_runtime_domain_authority_is_not_used_for_staging() {
    let fixture = ColdStartFixture::new();
    let moved = fixture.runtime_domain.with_extension("moved");
    fs::rename(&fixture.runtime_domain, &moved).unwrap();
    fs::create_dir(&fixture.runtime_domain).unwrap();
    fs::set_permissions(&fixture.runtime_domain, fs::Permissions::from_mode(0o700)).unwrap();
    assert!(ColdStart::stage(fixture.inputs(&fixture.escalation_client)).is_err());
    assert!(fixture.runtime_domain.read_dir().unwrap().next().is_none());
}

struct ColdStartFixture {
    tree: tempfile::TempDir,
    configuration: lazaret_core::configuration::ResolvedDomainConfiguration,
    persistent_domain: PathBuf,
    runtime_domain: PathBuf,
    prepared_lower: PathBuf,
    catalog: Catalog,
    store: Arc<HostCoordinatorDomainStore>,
    artifact: PreparedArtifact,
    artifact_authority: File,
    generation: Generation,
    container_id: ContainerId,
    escalation_client: File,
    invalid_escalation_client: File,
}

impl ColdStartFixture {
    fn new() -> Self {
        let tree = tempfile::tempdir().unwrap();
        let configuration = configuration(tree.path());
        let persistent_domain = host_path(
            configuration
                .persistent_domain_root()
                .canonical_label()
                .as_bytes(),
        );
        let runtime_domain = host_path(
            configuration
                .runtime_domain_root()
                .canonical_label()
                .as_bytes(),
        );
        let catalog = Catalog::new(
            tree.path().join("catalog.redb"),
            tree.path().join("catalog.lock"),
        );
        catalog.initialize().unwrap();
        let domain = Domain {
            domain_id: configuration.domain_id(),
            profile_name: configuration.selected_profile().to_owned(),
            canonical_workspace: configuration.workspace().canonical_label().clone(),
        };
        let lock = OpenOptions::new()
            .read(true)
            .write(true)
            .create_new(true)
            .open(persistent_domain.join("domain.lock"))
            .unwrap();
        let store = Arc::new(
            HostCoordinatorDomainStore::open(
                try_exclusive_flock(lock).unwrap(),
                &persistent_domain.join("domain.redb"),
                &domain,
            )
            .unwrap(),
        );
        let lower = tree.path().join("prepared-lower");
        fs::create_dir(&lower).unwrap();
        fs::create_dir(lower.join("etc")).unwrap();
        fs::write(lower.join("etc/resolv.conf"), b"image resolver\n").unwrap();
        fs::write(lower.join("etc/hosts"), b"image hosts\n").unwrap();
        fs::write(lower.join("etc/hostname"), b"image hostname\n").unwrap();
        let metadata = lower.metadata().unwrap();
        let artifact_authority = File::open(&lower).unwrap();
        let artifact = PreparedArtifact {
            prepared_artifact_id: PreparedArtifactId::new([8; 32]),
            source_identity: SourceIdentity::new(b"synthetic-spike".to_vec()),
            selected_manifest_digest: SelectedManifestDigest::new([9; 32]),
            policy_id: PolicyId::new([10; 32]),
            id_map_digest: IdMapDigest::new([11; 32]),
            path: HostPath::new(lower.as_os_str().as_bytes().to_vec()),
            device: DeviceId::new(metadata.dev()),
            inode: InodeId::new(metadata.ino()),
            state: ArtifactState::Available,
        };
        let generation = Generation {
            domain_id: configuration.domain_id(),
            generation_id: GenerationId::new(1),
            configuration_fingerprint: configuration.configuration_fingerprint(),
            coordinator_id: CoordinatorId::new([2; 32]),
            network_sidecar_id: lazaret_core::identity::NetworkSidecarId::new([3; 32]),
            sandbox_id: SandboxId::new([4; 32]),
            durable_decision: DurableDecision::None,
        };
        Self {
            tree,
            configuration,
            persistent_domain,
            runtime_domain,
            prepared_lower: lower,
            catalog,
            store,
            artifact,
            artifact_authority,
            generation,
            container_id: ContainerId::new([5; 32]),
            escalation_client: File::open("/bin/true").unwrap(),
            invalid_escalation_client: File::open("/dev/null").unwrap(),
        }
    }

    fn inputs<'a>(&'a self, escalation_client: &'a File) -> ColdStartInputs<'a> {
        ColdStartInputs {
            store: Arc::clone(&self.store),
            catalog: &self.catalog,
            configuration: &self.configuration,
            generation: self.generation.clone(),
            prepared_artifact: &self.artifact,
            prepared_artifact_authority: &self.artifact_authority,
            prepared_artifact_label: Path::new(
                std::str::from_utf8(self.artifact.path.as_bytes()).unwrap(),
            ),
            root_process: RootProcess::new("/bin/sh", ["-c".into(), "sleep infinity".into()])
                .unwrap(),
            container_id: self.container_id,
            invoking_uid: getuid().as_raw(),
            escalation_client,
            sandbox_init: escalation_client,
            user_namespace_map: test_user_namespace_map(),
        }
    }

    fn bundle(&self) -> PathBuf {
        self.runtime_domain
            .join("bundles")
            .join(hex(self.container_id.as_bytes()))
    }

    fn control(&self) -> PathBuf {
        self.runtime_domain
            .join("sandbox-control")
            .join(hex(self.container_id.as_bytes()))
    }

    fn overlay(&self) -> PathBuf {
        self.runtime_domain
            .join("runsc-overlay")
            .join(hex(self.container_id.as_bytes()))
    }
}

fn test_user_namespace_map() -> UserNamespaceMap {
    let uid = getuid().as_raw();
    let gid = getgid().as_raw();
    let subordinate_user_start = if uid < 131_072 { 131_072 } else { 1 };
    let subordinate_group_start = if gid < 131_072 { 131_072 } else { 1 };
    UserNamespaceMap {
        uid: vec![
            UserNamespaceIdRange {
                container_id: 0,
                parent_id: uid,
                size: 1,
            },
            UserNamespaceIdRange {
                container_id: 1,
                parent_id: subordinate_user_start,
                size: REQUIRED_SUBORDINATE_IDS,
            },
        ],
        gid: vec![
            UserNamespaceIdRange {
                container_id: 0,
                parent_id: gid,
                size: 1,
            },
            UserNamespaceIdRange {
                container_id: 1,
                parent_id: subordinate_group_start,
                size: REQUIRED_SUBORDINATE_IDS,
            },
        ],
    }
}

fn assert_private_artifacts(bundle: &Path, control: &Path) {
    let config = bundle.join("config.json");
    let escalation_client = control.join("bin/lazaret-escalate");
    assert!(config.is_file());
    assert!(escalation_client.is_file());
    assert!(control.join("credentials").is_dir());
    assert_eq!(mode(bundle), 0o700);
    assert_eq!(mode(control), 0o700);
    assert_eq!(mode(&config), 0o600);
    assert_eq!(mode(&escalation_client), 0o700);
}

fn mode(path: &Path) -> u32 {
    path.metadata().unwrap().permissions().mode() & 0o777
}

fn configuration(root: &Path) -> lazaret_core::configuration::ResolvedDomainConfiguration {
    let configuration_directory = root.join("configuration");
    let workspace = root.join("workspace");
    let runtime = root.join("runtime");
    fs::create_dir(&configuration_directory).unwrap();
    fs::create_dir(&workspace).unwrap();
    fs::create_dir(&runtime).unwrap();
    let configuration_path = configuration_directory.join("lazaret.toml");
    fs::write(
        &configuration_path,
        format!(
            "schema = \"lazaret-trusted-host-config-v1\"\n\n[storage]\ndata_root = \"{}\"\n\n[network]\ndefault_allow = [\"192.0.2.10/32\"]\n\n[profiles.default]\nimage = \"example.invalid/factory:v1\"\npersistent_dirs = []\n
",
            root.join("data").display()
        ),
    )
    .unwrap();
    load_trusted_configuration(ConfigurationLoadRequest {
        explicit_configuration_path: Some(configuration_path),
        selected_profile: "default".into(),
        workspace,
        home: root.join("home"),
        xdg_config_home: None,
        xdg_data_home: None,
        xdg_runtime_dir: runtime,
        environment: BTreeMap::new(),
        sandbox_layout: default_sandbox_layout().unwrap(),
    })
    .unwrap()
}

fn host_path(bytes: &[u8]) -> PathBuf {
    PathBuf::from(std::ffi::OsString::from_vec(bytes.to_vec()))
}

fn hex(bytes: &[u8]) -> String {
    use std::fmt::Write as _;

    bytes.iter().fold(String::new(), |mut output, byte| {
        write!(output, "{byte:02x}").unwrap();
        output
    })
}