#[path = "supervision.rs"]
pub(super) mod supervision;

use super::sys;
use avril_subprocess::{Options, Status, StdioSpec};
use std::cell::RefCell;
use std::collections::{BTreeMap, HashMap};
use std::convert::Infallible;
use std::io;
use std::path::Path;
use std::rc::Rc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use avril_caps::Capabilities;
use avril_core::serde_v8::{JsBuffer, ToJsBuffer};
use avril_core::{Extension, Op};
use avril_host::environment::Environment;
use serde::Serialize;
use tokio::runtime::Handle;
use tokio_util::sync::CancellationToken;

use avril_ext_file::{Files, process_file};
use avril_host::io::response;

const MAX_ID: u64 = 9_007_199_254_740_991;

fn invalid(message: &str) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidInput, message)
}
fn missing() -> io::Error {
    io::Error::new(io::ErrorKind::NotConnected, "Subprocess handle is closed")
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct SpawnInfo {
    id: u64,
    pid: u32,
    readable: Vec<u32>,
    writable: Vec<u32>,
    terminal: bool,
}
#[derive(Serialize)]
struct SyncOutput {
    status: Status,
    stdout: Option<ToJsBuffer>,
    stderr: Option<ToJsBuffer>,
}

struct Process {
    native: Arc<sys::Process>,
    /// Set while a thread is blocked inside run_sync on this process, so the
    /// REPL can interrupt exactly the subprocess its evaluation is waiting on.
    sync_wait: AtomicBool,
    worker: bool,
}
impl std::ops::Deref for Process {
    type Target = sys::Process;
    fn deref(&self) -> &Self::Target {
        &self.native
    }
}

struct Registry {
    runtime: Handle,
    caps: Arc<Capabilities>,
    files: Arc<Files>,
    processes: Mutex<HashMap<u64, Arc<Process>>>,
    next: AtomicU64,
    owner: CancellationToken,
    tree: Arc<supervision::Tree>,
    worker: bool,
    signal: super::SignalMode,
    signals: Arc<super::signals::State>,
}
impl Registry {
    fn get(&self, id: u64) -> io::Result<Arc<Process>> {
        self.processes
            .lock()
            .unwrap()
            .get(&id)
            .cloned()
            .ok_or_else(missing)
    }
    fn release(&self, id: u64) {
        if let Some(process) = self.processes.lock().unwrap().remove(&id) {
            process.request("dispose", None, false);
        }
    }
    fn spawn(&self, options: Options, environment: &Environment) -> io::Result<SpawnInfo> {
        if self.owner.is_cancelled() || self.tree.stopped.is_cancelled() || self.tree.is_closed() {
            return Err(missing());
        }
        if self.worker && options.stdin.mode == "inherit" {
            return Err(io::Error::new(
                io::ErrorKind::Unsupported,
                "Worker subprocesses cannot inherit host stdin",
            ));
        }
        // Only run mode owns SIGINT; serve and the REPL leave it to their
        // own handlers. SIGTERM fate belongs to whoever claims it, if anyone.
        self.tree
            .install(&self.runtime, matches!(self.signal, super::SignalMode::Run))?;
        let id = self
            .next
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| {
                (id <= MAX_ID).then_some(id + 1)
            })
            .map_err(|_| io::Error::other("Subprocess handle space exhausted"))?;
        let terminal = options.terminal.is_some();
        let native = sys::Launcher {
            runtime: self.runtime.clone(),
            caps: self.caps.clone(),
            owners: vec![self.owner.clone(), self.tree.stopped.clone()],
        }
        .spawn(
            options,
            environment
                .to_object(&self.caps)
                .into_iter()
                .map(|(k, v)| (k.into(), v.into())),
            None,
            |spec, read, cwd| self.file(spec, read, cwd),
        )?;
        let process = Arc::new(Process {
            native,
            sync_wait: AtomicBool::new(false),
            worker: self.worker,
        });
        let info = SpawnInfo {
            id,
            pid: process.pid,
            terminal,
            readable: process.outputs.keys().copied().collect(),
            writable: process.inputs.keys().copied().collect(),
        };
        self.processes.lock().unwrap().insert(id, process.clone());
        self.tree.register(&process);
        Ok(info)
    }
    fn file(&self, spec: &StdioSpec, read: bool, cwd: &Path) -> io::Result<std::fs::File> {
        if let Some(id) = spec.handle {
            return self.files.duplicate_for_process(id, read, !read);
        }
        let path = cwd.join(
            spec.file
                .as_ref()
                .ok_or_else(|| invalid("Missing stdio file"))?,
        );
        let flags = if read {
            (true, false, false, false, false, false, None)
        } else {
            match spec.file_mode.as_deref() {
                Some("append") => (false, true, true, false, true, false, None),
                Some("truncate") => (false, true, false, true, true, false, None),
                Some("create-new") => (false, true, false, false, false, true, None),
                _ => return Err(invalid("Invalid stdio file mode")),
            }
        };
        process_file(
            self.caps.clone(),
            path.to_str()
                .ok_or_else(|| invalid("Stdio paths must be UTF-8"))?
                .into(),
            flags,
        )
    }
}
impl Drop for Registry {
    fn drop(&mut self) {
        // Release signal bindings first, so a dying isolate stops pinning
        // listened signals process-wide before processes drain.
        self.signals.shutdown();
        let processes: Vec<_> = self
            .processes
            .get_mut()
            .unwrap()
            .drain()
            .map(|(_, process)| process)
            .collect();
        for process in &processes {
            process.request("owner", None, true);
        }
        let deadline = Instant::now() + Duration::from_secs(2);
        for process in processes {
            process.wait_blocking(deadline.saturating_duration_since(Instant::now()));
        }
    }
}

pub fn init(
    runtime: Handle,
    caps: Arc<Capabilities>,
    environment: Rc<RefCell<Environment>>,
    files: Arc<Files>,
    owner: CancellationToken,
    worker: bool,
    signal: super::SignalMode,
) -> Extension {
    let tree = supervision::Tree::global();
    let signals = super::signals::State::create(runtime.clone(), worker);
    let registry = Arc::new(Registry {
        runtime,
        caps,
        files,
        processes: Mutex::new(HashMap::new()),
        next: AtomicU64::new(1),
        owner,
        tree: tree.clone(),
        worker,
        signal,
        signals: signals.clone(),
    });
    let mut ops = Vec::new();
    ops.extend(super::signals::ops(&signals));
    ops.push(Op::new("op_process_spawn", {
        let registry = registry.clone();
        let environment = environment.clone();
        move |options: Options| {
            Ok::<_, Infallible>(response(registry.spawn(options, &environment.borrow())))
        }
    }));
    ops.push(Op::new("op_process_run_sync", {
        let registry = registry.clone();
        let environment = environment.clone();
        move |(mut options, input): (Options, Option<JsBuffer>)| {
            let generation = registry.tree.sync_generation();
            // The blocking reader below drains both streams; backpressure keeps
            // a burst from outrunning it into the unread limit.
            options.backpressure = [(1, &options.stdout), (2, &options.stderr)]
                .into_iter()
                .filter(|(_, spec)| spec.mode == "piped")
                .map(|(fd, _)| fd)
                .collect();
            let input = input.map(|bytes| bytes.to_vec());
            let result = (|| {
                let info = registry.spawn(options, &environment.borrow())?;
                let process = registry.get(info.id)?;
                process.sync_wait.store(true, Ordering::SeqCst);
                if !registry.worker && registry.tree.sync_generation() != generation {
                    process.request("terminated", None, false);
                }
                let output = async {
                    let write = async {
                        if let Some(pipe) = process.inputs.get(&0) {
                            if let Some(bytes) = input
                                && !bytes.is_empty()
                                && let Err(error) = pipe.write(bytes).await
                                && error.kind() != io::ErrorKind::BrokenPipe
                            {
                                process.request("io-error", Some(error.to_string()), false);
                            }
                            pipe.close().await;
                        }
                    };
                    let reads = futures::future::join_all(process.outputs.iter().map(
                        |(fd, pipe)| async move {
                            let mut bytes = Vec::new();
                            while let Some(chunk) = pipe.read().await {
                                bytes.extend_from_slice(&chunk);
                            }
                            (*fd, bytes)
                        },
                    ));
                    let (_, bytes) = futures::future::join(write, reads).await;
                    let status = process.wait(true, None).await?.ok_or_else(missing)?;
                    let mut bytes: BTreeMap<_, _> = bytes.into_iter().collect();
                    Ok::<_, io::Error>(SyncOutput {
                        status,
                        stdout: bytes.remove(&1).map(ToJsBuffer::from),
                        stderr: bytes.remove(&2).map(ToJsBuffer::from),
                    })
                };
                // Sync ops may run inside a Tokio-driven JS event loop.
                let _enter = registry.runtime.enter();
                let output = avril_host::wait(output, || false).unwrap();
                process.sync_wait.store(false, Ordering::SeqCst);
                registry.release(info.id);
                output
            })();
            Ok::<_, Infallible>(response(result))
        }
    }));
    ops.push(Op::asynchronous_unreferenced("op_process_read", {
        let registry = registry.clone();
        move |(id, fd): (u64, u32)| {
            let output = registry
                .get(id)
                .and_then(|process| process.outputs.get(&fd).cloned().ok_or_else(missing));
            async move {
                Ok::<_, Infallible>(response(match output {
                    Ok(output) => Ok(output
                        .read()
                        .await
                        .map(|bytes| ToJsBuffer::from(bytes.to_vec()))),
                    Err(error) => Err(error),
                }))
            }
        }
    }));
    ops.push(Op::asynchronous_unreferenced("op_process_write", {
        let registry = registry.clone();
        move |(id, fd, bytes): (u64, u32, JsBuffer)| {
            let bytes = bytes.to_vec();
            let input = registry
                .get(id)
                .and_then(|process| process.inputs.get(&fd).cloned().ok_or_else(missing));
            async move {
                Ok::<_, Infallible>(response(match input {
                    Ok(input) => {
                        let result = input.write(bytes).await;
                        if result
                            .as_ref()
                            .is_err_and(|error| error.kind() == io::ErrorKind::BrokenPipe)
                        {
                            input.close().await;
                        }
                        result
                    }
                    Err(error) => Err(error),
                }))
            }
        }
    }));
    ops.push(Op::asynchronous_unreferenced("op_process_end", {
        let registry = registry.clone();
        move |(id, fd): (u64, u32)| {
            let input = registry
                .get(id)
                .and_then(|process| process.inputs.get(&fd).cloned().ok_or_else(missing));
            async move {
                Ok::<_, Infallible>(response(match input {
                    Ok(input) => {
                        input.close().await;
                        Ok(())
                    }
                    Err(error) => Err(error),
                }))
            }
        }
    }));
    ops.push(Op::new("op_process_cancel_read", {
        let registry = registry.clone();
        move |(id, fd): (u64, u32)| {
            let result = registry.get(id).and_then(|process| {
                process.outputs.get(&fd).ok_or_else(missing)?.cancel();
                Ok(())
            });
            Ok::<_, Infallible>(response(result))
        }
    }));
    ops.push(Op::asynchronous_unreferenced("op_process_wait", {
        let registry = registry.clone();
        move |(id, completed): (u64, bool)| {
            let process = registry.get(id);
            async move {
                Ok::<_, Infallible>(response(match process {
                    Ok(process) => process
                        .wait(completed, None)
                        .await
                        .and_then(|status| status.ok_or_else(missing)),
                    Err(error) => Err(error),
                }))
            }
        }
    }));
    ops.push(Op::asynchronous("op_process_keepalive", {
        let registry = registry.clone();
        move |(id, epoch): (u64, u32)| {
            let process = registry.get(id);
            async move {
                Ok::<_, Infallible>(response(match process {
                    Ok(process) => process.wait(true, Some(epoch)).await.map(|_| ()),
                    Err(error) => Err(error),
                }))
            }
        }
    }));
    ops.push(Op::new("op_process_ref", {
        let registry = registry.clone();
        move |(id, _referenced): (u64, bool)| {
            let result = registry.get(id).map(|process| process.invalidate_waiters());
            Ok::<_, Infallible>(response(result))
        }
    }));
    ops.push(Op::new("op_process_kill", {
        let registry = registry.clone();
        move |(id, signal): (u64, String)| {
            Ok::<_, Infallible>(response(
                registry.get(id).and_then(|process| process.kill(&signal)),
            ))
        }
    }));
    ops.push(Op::new("op_process_stop", {
        let registry = registry.clone();
        move |(id, reason): (u64, String)| {
            let result = registry
                .get(id)
                .map(|process| process.request(&reason, None, false));
            Ok::<_, Infallible>(response(result))
        }
    }));
    ops.push(Op::new("op_process_release", {
        let registry = registry.clone();
        move |id: u64| {
            registry.release(id);
            Ok::<_, Infallible>(())
        }
    }));
    ops.push(Op::new("op_process_resize", {
        let registry = registry.clone();
        move |(id, columns, rows): (u64, u16, u16)| {
            let result = registry
                .get(id)
                .and_then(|process| process.resize(columns, rows));
            Ok::<_, Infallible>(response(result))
        }
    }));
    Extension {
        name: "ext:process",
        dependencies: &["ext:encoding", "ext:file", "ext:web", "ext:webstreams"],
        ops,
        modules: &[
            (
                "ext:process/main.js",
                include_str!(concat!(env!("OUT_DIR"), "/process.js")),
            ),
            (
                "ext:process/signals.js",
                include_str!(concat!(env!("OUT_DIR"), "/signals.js")),
            ),
        ],
        entry_point: None,
    }
}