mod support;

use std::cell::{Cell, RefCell};
use std::io;
use std::rc::Rc;

use avril_core::{
    Error, Extension, ModuleLoader, ModuleSource, Op, Runtime, RuntimeOptions, Snapshot,
    SnapshotError, WasmStreamingRegistry, v8,
};

struct Loader(u32);

impl ModuleLoader for Loader {
    fn resolve(&self, specifier: &str, _: &str) -> Result<String, Error> {
        Ok(if specifier == "private-alias" {
            "ext:test/state.js"
        } else {
            specifier
        }
        .into())
    }

    fn load(&self, name: &str) -> Result<ModuleSource, Error> {
        match name {
            "guest:main" => Ok(include_str!("snapshot/guest.js").into()),
            "guest:dep" => Ok(format!(
                "export const seed = {}; export const url = import.meta.url;",
                self.0
            )
            .into()),
            _ => Err(Error::Module(format!("Missing {name}"))),
        }
    }

    fn module_url(&self, name: &str) -> Option<String> {
        Some(format!("file:///snapshot-{}/{name}.js", self.0))
    }

    fn map_source_location(
        &self,
        name: &str,
        line: u32,
        column: u32,
    ) -> Option<(String, u32, u32)> {
        (name == "guest:stack").then(|| (format!("fresh-{}.ts", self.0), line + 10, column + 2))
    }
}

fn options(seed: u32, allowed: Rc<Cell<bool>>) -> RuntimeOptions {
    let mut ops = vec![
        Op::new("staticOnly", move |(): ()| {
            assert_eq!(seed, 0, "Static entry reexecuted during restoration");
            Ok::<_, io::Error>(())
        }),
        Op::new("nativeSeed", move |(): ()| {
            assert_ne!(seed, 0, "Live initialization ran during snapshot creation");
            if !allowed.get() {
                return Err(io::Error::other("native access denied"));
            }
            Ok::<_, io::Error>(seed)
        }),
        Op::with_scope("collect", |scope, (): ()| {
            scope.low_memory_notification();
            Ok::<_, io::Error>(())
        }),
        Op::asynchronous("asyncSeed", move |(): ()| async move {
            Ok::<_, io::Error>(seed)
        }),
        Op::try_raw("registerWasm", |scope, args| {
            let callback =
                v8::Local::<v8::Function>::try_from(args.get(0)).map_err(io::Error::other)?;
            scope
                .get_slot::<Rc<RefCell<WasmStreamingRegistry>>>()
                .unwrap()
                .borrow_mut()
                .register(v8::Global::new(scope, callback))?;
            Ok::<_, io::Error>(v8::undefined(scope).into())
        }),
        Op::with_scope("finishWasm", |scope, id: u64| {
            let registry = scope
                .get_slot::<Rc<RefCell<WasmStreamingRegistry>>>()
                .unwrap();
            let mut registry = registry.borrow_mut();
            registry
                .stream_mut(id)?
                .on_bytes_received(b"\0asm\x01\0\0\0");
            registry.finish(id)
        }),
    ];
    if seed != 0 {
        ops.reverse();
    }
    RuntimeOptions {
        extensions: vec![Extension {
            name: "ext:test",
            dependencies: &[],
            ops,
            modules: &[
                ("ext:test/main.js", include_str!("snapshot/main.js")),
                (
                    "ext:test/state.js",
                    "export const state = {seed: undefined};",
                ),
                ("ext:test/lazy.js", include_str!("snapshot/lazy.js")),
            ],
            entry_point: Some("ext:test/main.js"),
        }],
        module_loader: Some(Rc::new(Loader(seed))),
    }
}

#[test]
fn core_snapshot_cross_process() {
    let path = || support::directory().join("core.snapshot");
    match support::role().as_deref() {
        Some("create") => {
            let snapshot =
                Runtime::create_snapshot(vec![options(0, Rc::new(Cell::new(true)))]).unwrap();
            std::fs::write(path(), snapshot.as_bytes()).unwrap();
        }
        Some("restore") => {
            let bytes = std::fs::read(path()).unwrap();
            // SAFETY: The sibling create role produced this artifact with this executable.
            let snapshot = unsafe { Snapshot::from_bytes(bytes) }.unwrap();
            for seed in [10, 100] {
                let allowed = Rc::new(Cell::new(true));
                let mut runtime =
                    Runtime::from_snapshot(options(seed, allowed.clone()), &snapshot).unwrap();
                runtime
                    .register_extension_module("test:main", "ext:test/main.js")
                    .unwrap();
                runtime
                    .register_extension_module("test:lazy", "ext:test/lazy.js")
                    .unwrap();
                runtime.execute_module("guest:main").unwrap();
                runtime
                    .execute_script::<()>("report events", "afterTasks()")
                    .unwrap();
                assert!(
                    runtime
                        .execute_script::<bool>("complete", "globalThis.complete === true")
                        .unwrap()
                );
                assert_eq!(
                    runtime
                        .execute_script::<u32>("seed", "snapshotState.seed")
                        .unwrap(),
                    seed
                );
                assert_eq!(
                    runtime
                        .execute_script::<Vec<String>>("events", "events")
                        .unwrap(),
                    [
                        "static",
                        "live",
                        "lazy",
                        "task",
                        "finish",
                        "microtask",
                        "rejection"
                    ]
                );
                let stack: String = runtime
                    .execute_script("guest:stack", "new Error('restored stack').stack")
                    .unwrap();
                assert!(stack.contains(&format!("fresh-{seed}.ts:11:")), "{stack}");
                runtime
                    .execute_script::<()>("handled", "rejected.catch(() => {});")
                    .unwrap();
                assert!(
                    runtime
                        .execute_script::<bool>("handled event", "events.at(-1) === 'handled'")
                        .unwrap()
                );
                allowed.set(false);
                assert!(
                    matches!(runtime.execute_script::<u32>("denied", "readNative()"), Err(Error::JavaScript(message)) if message.contains("native access denied"))
                );
            }
        }
        _ => support::roles("core_snapshot_cross_process", &["create", "restore"]),
    }
}

#[test]
fn snapshot_layout_mismatches_are_rejected_before_live_initialization() {
    if support::role().is_none() {
        return support::roles(
            "snapshot_layout_mismatches_are_rejected_before_live_initialization",
            &["run"],
        );
    }
    let snapshot = Runtime::create_snapshot(vec![options(0, Rc::new(Cell::new(true)))]).unwrap();
    for change in 0..3 {
        let mut options = options(10, Rc::new(Cell::new(false)));
        match change {
            0 => options.extensions[0].ops[0].name = "renamed",
            1 => {
                options.extensions[0].modules =
                    &[("ext:test/main.js", "throw Error('changed source');")]
            }
            2 => options.extensions[0].entry_point = None,
            _ => unreachable!(),
        }
        assert!(matches!(
            Runtime::from_snapshot(options, &snapshot),
            Err(Error::Snapshot(SnapshotError::IncompatibleLayout))
        ));
    }
}

#[test]
fn restored_live_failure_does_not_run_later_initializers() {
    if support::role().is_none() {
        return support::roles(
            "restored_live_failure_does_not_run_later_initializers",
            &["run"],
        );
    }
    let called = Rc::new(Cell::new(false));
    let options = || RuntimeOptions {
        extensions: vec![Extension {
            name: "ext:test",
            dependencies: &[],
            ops: vec![
                Op::new("fail", |(): ()| {
                    Err::<(), _>(io::Error::other("failed live setup"))
                }),
                Op::new("later", {
                    let called = called.clone();
                    move |(): ()| {
                        called.set(true);
                        Ok::<_, io::Error>(())
                    }
                }),
            ],
            modules: &[(
                "ext:test/main.js",
                r#"
                import {onRuntimeStart} from 'ext:core/bootstrap';
                import {fail, later} from 'ext:core/ops';
                onRuntimeStart(() => { fail(); });
                onRuntimeStart(() => { later(); });
            "#,
            )],
            entry_point: Some("ext:test/main.js"),
        }],
        ..Default::default()
    };
    let snapshot = Runtime::create_snapshot(vec![options()]).unwrap();
    assert!(matches!(
        Runtime::from_snapshot(options(), &snapshot),
        Err(Error::JavaScript(message)) if message.contains("failed live setup")
    ));
    assert!(!called.get());
}

struct DropProbe(Rc<Cell<bool>>);

impl Drop for DropProbe {
    fn drop(&mut self) {
        self.0.set(true);
    }
}

#[test]
fn abandoned_snapshot_builders_release_native_state_and_return_errors() {
    if support::role().is_none() {
        return support::roles(
            "abandoned_snapshot_builders_release_native_state_and_return_errors",
            &["run"],
        );
    }
    let cases: &[&[(&str, &str)]] = &[
        &[("ext:test/main.js", "throw Error('failed static setup');")],
        &[(
            "ext:test/main.js",
            "import {pending} from 'ext:core/ops'; pending();",
        )],
        &[(
            "ext:test/main.js",
            "import {setEventHandlers} from 'ext:core/event_hooks'; setEventHandlers(() => true, () => true, () => {});",
        )],
    ];
    for (index, modules) in cases.iter().enumerate() {
        let dropped = Rc::new(Cell::new(false));
        let probe = DropProbe(dropped.clone());
        let result = Runtime::create_snapshot(vec![RuntimeOptions {
            extensions: vec![Extension {
                name: "ext:test",
                dependencies: &[],
                ops: vec![
                    Op::new("retain", move |(): ()| {
                        let _retain = &probe;
                        Ok::<_, io::Error>(())
                    }),
                    Op::asynchronous_unreferenced("pending", |(): ()| {
                        futures::future::pending::<Result<(), io::Error>>()
                    }),
                ],
                modules,
                entry_point: Some("ext:test/main.js"),
            }],
            ..Default::default()
        }]);
        if index == 0 {
            assert!(
                matches!(result, Err(Error::JavaScript(message)) if message.contains("failed static setup"))
            );
        } else {
            assert!(matches!(
                result,
                Err(Error::Snapshot(SnapshotError::ActiveState(_)))
            ));
        }
        assert!(dropped.get(), "Native op state leaked");
    }
}