//! WAI Studio, served by `wai dev` at `/_wai/` on the origin it already serves on loopback:
//! no listener of its own. The page's files are Studio's build (`wai/wai-studio`): from
//! `$WAI_STUDIO`, else `share/wai/studio` beside this binary's directory (an installed
//! archive's layout). Beside them, under `/_wai/api/`, the dev stack's own routes for the
//! page:
//!
//! - `GET /_wai/api/state`: the stack as `wai dev` started it: its origin, room service and
//!   relay, each service with its process and the loopback address it listens on, its logs,
//!   and what the local `wai` decodes and encodes;
//! - `GET /_wai/api/log?name=<relay|room|conversation>&from=<byte>`: a log from a byte on;
//! - `GET /_wai/api/sessions`: the relay's session records (`sessions.jsonl`);
//! - `POST /_wai/api/op/<id>`: one of [`LOCAL_OPS`], run by this `wai` on the bytes the page
//!   sends, in a directory of its own under the dev directory, removed after. The answer is
//!   its exit status, what it printed and the files it wrote.
//!
//! Every API request must name this origin as its host (`127.0.0.1:<port>`, `localhost:<port>`
//! or `[::1]:<port>`), which a page reached through another name cannot; must carry
//! `X-WAI-Studio: 1`, a header no page of another origin can send without a preflight this
//! origin never answers; and, when it states an origin or a fetch site, must state this one.

use super::{Head, respond};
use serde_json::{Value, json};
use std::io::Read;
use std::net::TcpStream;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

/// Studio's path on the origin.
pub(super) const PREFIX: &str = "/_wai";
const API: &str = "/_wai/api";

/// The operations the local `wai` runs for the page, by their catalogue ids, each with the
/// command it runs. Studio's drift check (`wai/wai-studio/scripts/actions.mjs`) holds this
/// table to the catalogue and to Studio's bindings.
const LOCAL_OPS: &[(&str, &str)] = &[
    ("envelope.inspect", "inspect"),
    ("media.decode", "decode"),
    ("media.encode", "encode"),
];

/// The codecs `wai encode` takes in this build, and what each encodes from.
fn encodes() -> &'static [&'static str] {
    if cfg!(feature = "codecs") {
        &["png", "jpeg", "avif", "jxl", "jxl-lossless", "opus", "flac", "zstd", "xz"]
    } else if cfg!(feature = "floor") {
        &["png", "zstd"]
    } else {
        &[]
    }
}

const IMAGE: &[&str] = &["png", "jpeg", "avif", "jxl", "jxl-lossless"];
const AUDIO: &[&str] = &["opus", "flac"];

/// The longest body the page sends an operation.
const MAX_INPUT: usize = 64 << 20;
/// The most an operation answers with: what it wrote, all of it.
const MAX_OUTPUT: usize = 64 << 20;
/// What it printed, each stream, at most.
const MAX_PRINTED: usize = 1 << 20;
/// How long an operation may run.
const OP_SECONDS: u64 = 120;

pub(super) struct Studio {
    /// Studio's files (`index.html` and the rest), when they are found.
    pub(super) files: Option<PathBuf>,
    /// Where operations run: `<dev dir>/studio/`.
    work: PathBuf,
    /// The dev directory: keys, configuration, logs and the relay's session records.
    dir: PathBuf,
    /// The origin's port, which every API request must name.
    port: Mutex<u16>,
    state: Mutex<Value>,
    started: Instant,
    exe: Option<PathBuf>,
    one_at_a_time: Mutex<()>,
    next: AtomicU64,
}

impl Studio {
    pub(super) fn new(files: Option<PathBuf>, dir: &Path) -> Studio {
        Studio {
            files,
            work: dir.join("studio"),
            dir: dir.to_path_buf(),
            port: Mutex::new(0),
            state: Mutex::new(Value::Null),
            started: Instant::now(),
            exe: std::env::current_exe().ok(),
            one_at_a_time: Mutex::new(()),
            next: AtomicU64::new(0),
        }
    }

    /// The stack is up on `port`: its state, as `GET /_wai/api/state` gives it.
    pub(super) fn ready(&self, port: u16, mut state: Value) {
        *self.port.lock().unwrap_or_else(|e| e.into_inner()) = port;
        let logs: Vec<&str> = ["relay", "room", "conversation"].into_iter().filter(|n| self.dir.join(format!("{n}.log")).is_file()).collect();
        state["logs"] = json!(logs);
        state["decodes"] = json!(crate::native_capabilities());
        state["encodes"] = json!(encodes());
        state["operations"] = json!(LOCAL_OPS.iter().map(|(id, _)| id).collect::<Vec<_>>());
        *self.state.lock().unwrap_or_else(|e| e.into_inner()) = state;
    }
}

/// Studio's files: `$WAI_STUDIO`, or `share/wai/studio` beside the binary's directory, each
/// taken when it holds `index.html`.
pub(super) fn find(exe_dir: Option<&Path>) -> Option<PathBuf> {
    let candidates: Vec<PathBuf> = match std::env::var_os("WAI_STUDIO") {
        Some(p) => vec![PathBuf::from(p)],
        None => exe_dir.map(|d| vec![d.join("../share/wai/studio"), d.join("share/wai/studio")]).unwrap_or_default(),
    };
    candidates.into_iter().find(|c| c.join("index.html").is_file()).and_then(|c| std::fs::canonicalize(c).ok())
}

/// A request under `/_wai`.
pub(super) fn handle(c: &mut TcpStream, h: &Head, rest: Vec<u8>, path: &str, s: &Studio) -> std::io::Result<()> {
    if path == API || path.starts_with(&format!("{API}/")) {
        return api(c, h, rest, &path[API.len()..], s);
    }
    let head_only = h.method == "HEAD";
    if h.method != "GET" && !head_only {
        return respond(c, "405 Method Not Allowed", "text/plain", b"method\n", false, &[("Allow", "GET, HEAD")]);
    }
    if path == PREFIX {
        return respond(c, "301 Moved Permanently", "text/plain", b"", head_only, &[("Location", "/_wai/")]);
    }
    let Some(root) = &s.files else {
        return respond(
            c,
            "404 Not Found",
            "text/plain; charset=utf-8",
            b"WAI Studio's files are not here: build them (wai/wai-studio: node build.mjs), then give their dist/ as WAI_STUDIO, or install the archive that holds share/wai/studio\n",
            head_only,
            &[],
        );
    };
    let Some(rel) = super::clean(&path[PREFIX.len()..]) else {
        return respond(c, "400 Bad Request", "text/plain", b"path\n", head_only, &[]);
    };
    let mut p = root.join(&rel);
    if p.is_dir() {
        p = p.join("index.html");
    }
    if !p.is_file() {
        return respond(c, "404 Not Found", "text/plain", b"not found\n", head_only, &[]);
    }
    let body = std::fs::read(&p)?;
    respond(c, "200 OK", super::mime(&p), &body, head_only, &[])
}

fn answer(c: &mut TcpStream, status: &str, v: &Value) -> std::io::Result<()> {
    respond(c, status, "application/json", format!("{v}\n").as_bytes(), false, &[])
}

/// Why a request to the API is refused, or `None`: it must name this origin as its host,
/// carry `X-WAI-Studio: 1`, and state no other origin or fetch site.
fn refused(h: &Head, port: u16) -> Option<&'static str> {
    let ours = [format!("127.0.0.1:{port}"), format!("localhost:{port}"), format!("[::1]:{port}")];
    if port == 0 || !h.header("host").is_some_and(|v| ours.iter().any(|o| o.eq_ignore_ascii_case(v.trim()))) {
        return Some("this origin answers requests that name it as their host");
    }
    if h.header("x-wai-studio").map(str::trim) != Some("1") {
        return Some("a request to the dev stack's API carries X-WAI-Studio: 1");
    }
    if let Some(o) = h.header("origin")
        && !ours.iter().any(|x| o.trim().eq_ignore_ascii_case(&format!("http://{x}")))
    {
        return Some("a request from a page of another origin");
    }
    if h.header("sec-fetch-site").is_some_and(|v| !v.trim().eq_ignore_ascii_case("same-origin")) {
        return Some("a request from a page of another origin");
    }
    None
}

fn query<'a>(target: &'a str, key: &str) -> Option<&'a str> {
    target.split_once('?')?.1.split('&').find_map(|kv| kv.split_once('=').filter(|(k, _)| *k == key).map(|(_, v)| v))
}

fn api(c: &mut TcpStream, h: &Head, rest: Vec<u8>, sub: &str, s: &Studio) -> std::io::Result<()> {
    let port = *s.port.lock().unwrap_or_else(|e| e.into_inner());
    if let Some(why) = refused(h, port) {
        return answer(c, "403 Forbidden", &json!({ "error": why }));
    }
    match (h.method.as_str(), sub) {
        ("GET", "/state") => {
            let mut v = s.state.lock().unwrap_or_else(|e| e.into_inner()).clone();
            if v.is_null() {
                return answer(c, "503 Service Unavailable", &json!({ "error": "the stack is starting" }));
            }
            v["uptime_s"] = json!(s.started.elapsed().as_secs_f64());
            answer(c, "200 OK", &v)
        }
        ("GET", "/log") => {
            let name = query(&h.target, "name").unwrap_or("");
            if !["relay", "room", "conversation"].contains(&name) {
                return answer(c, "404 Not Found", &json!({ "error": "no such log" }));
            }
            let from: u64 = query(&h.target, "from").and_then(|v| v.parse().ok()).unwrap_or(0);
            let bytes = std::fs::read(s.dir.join(format!("{name}.log"))).unwrap_or_default();
            let size = bytes.len() as u64;
            // A log read from its last 64 KiB when the page asks from further back.
            let at = if size.saturating_sub(from) > 64 << 10 { size - (64 << 10) } else { from.min(size) };
            let text = String::from_utf8_lossy(&bytes[at as usize..]).into_owned();
            answer(c, "200 OK", &json!({ "name": name, "from": at, "next": size, "text": text }))
        }
        ("GET", "/sessions") => {
            let t = std::fs::read_to_string(s.dir.join("sessions.jsonl")).unwrap_or_default();
            let lines: Vec<Value> = t.lines().filter_map(|l| serde_json::from_str(l).ok()).collect();
            let last = lines[lines.len().saturating_sub(50)..].to_vec();
            answer(c, "200 OK", &json!({ "sessions": last }))
        }
        ("POST", op) if op.starts_with("/op/") => {
            let id = op.trim_start_matches("/op/").split('?').next().unwrap_or("");
            if h.header("transfer-encoding").is_some() {
                return answer(c, "411 Length Required", &json!({ "error": "a body with a length" }));
            }
            let len: usize = h.header("content-length").and_then(|v| v.trim().parse().ok()).unwrap_or(0);
            if len > MAX_INPUT {
                return answer(c, "413 Content Too Large", &json!({ "error": "the input is larger than the page may send" }));
            }
            let mut body = rest;
            body.truncate(len);
            let mut chunk = [0u8; 65536];
            while body.len() < len {
                let n = c.read(&mut chunk[..(len - body.len()).min(65536)])?;
                if n == 0 {
                    break;
                }
                body.extend_from_slice(&chunk[..n]);
            }
            match run_op(s, id, &h.target, &body) {
                Ok(v) => answer(c, "200 OK", &v),
                Err((status, why)) => answer(c, status, &json!({ "error": why })),
            }
        }
        _ => answer(c, "404 Not Found", &json!({ "error": "no such route" })),
    }
}

/// One operation, run by this `wai` in a directory of its own.
fn run_op(s: &Studio, id: &str, target: &str, body: &[u8]) -> Result<Value, (&'static str, String)> {
    let Some((_, verb)) = LOCAL_OPS.iter().find(|(i, _)| *i == id) else {
        return Err(("404 Not Found", format!("{id}: not an operation the local wai runs for the page")));
    };
    let exe = s.exe.clone().ok_or(("500 Internal Server Error", "this binary's path is unknown".to_string()))?;
    let _one = s.one_at_a_time.lock().unwrap_or_else(|e| e.into_inner());
    let n = s.next.fetch_add(1, Ordering::Relaxed);
    let dir = s.work.join(format!("op-{}-{n}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).map_err(|e| ("500 Internal Server Error", e.to_string()))?;
    let r = (|| {
        let write = |name: &str, b: &[u8]| std::fs::write(dir.join(name), b).map_err(|e| ("500 Internal Server Error", e.to_string()));
        let num = |k: &str| query(target, k).and_then(|v| v.parse::<u32>().ok());
        const NOTHING: &[&str] = &[];
        const DECODED: &[&str] = &["out", "out.meta"];
        const ENCODED: &[&str] = &["out.wai"];
        let (argv, outputs): (Vec<String>, &[&str]) = match *verb {
            "inspect" => {
                write("in.wai", body)?;
                (vec!["inspect".into(), "in.wai".into()], NOTHING)
            }
            "decode" => {
                write("in.wai", body)?;
                (vec!["decode".into(), "in.wai".into(), "-o".into(), "out".into()], DECODED)
            }
            _ => {
                let codec = query(target, "codec").unwrap_or("");
                if !encodes().contains(&codec) {
                    return Err(("400 Bad Request", format!("{codec}: this wai encodes {}", encodes().join(", "))));
                }
                let input = if IMAGE.contains(&codec) {
                    let (Some(w), Some(h)) = (num("width"), num("height")) else {
                        return Err(("400 Bad Request", "a picture takes its width and height".into()));
                    };
                    if (w as usize) * (h as usize) * 3 != body.len() {
                        return Err(("400 Bad Request", format!("{} bytes is not {w} × {h} packed RGB8", body.len())));
                    }
                    write("in.rgb", body)?;
                    write("in.rgb.meta", format!("{h} {w}\n").as_bytes())?;
                    "in.rgb"
                } else if AUDIO.contains(&codec) {
                    let Some(rate) = num("sample_rate") else {
                        return Err(("400 Bad Request", "a sound takes its sample rate".into()));
                    };
                    write("in.f32", body)?;
                    write("in.f32.meta", format!("{rate}\n").as_bytes())?;
                    "in.f32"
                } else {
                    write("in.bin", body)?;
                    "in.bin"
                };
                let q = num("quality").unwrap_or(75).clamp(1, 100);
                (vec!["encode".into(), codec.into(), input.into(), "-o".into(), "out.wai".into(), "-q".into(), q.to_string()], ENCODED)
            }
        };
        let out = std::fs::File::create(dir.join(".stdout")).map_err(|e| ("500 Internal Server Error", e.to_string()))?;
        let err = std::fs::File::create(dir.join(".stderr")).map_err(|e| ("500 Internal Server Error", e.to_string()))?;
        let mut child = Command::new(&exe)
            .args(&argv)
            .current_dir(&dir)
            .stdin(Stdio::null())
            .stdout(out)
            .stderr(err)
            .env("NO_COLOR", "1")
            .spawn()
            .map_err(|e| ("500 Internal Server Error", format!("{}: {e}", exe.display())))?;
        let end = Instant::now() + Duration::from_secs(OP_SECONDS);
        let status = loop {
            if let Some(st) = child.try_wait().map_err(|e| ("500 Internal Server Error", e.to_string()))? {
                break st;
            }
            if Instant::now() > end {
                let _ = child.kill();
                let _ = child.wait();
                return Err(("504 Gateway Timeout", format!("{id}: not done in {OP_SECONDS} s")));
            }
            std::thread::sleep(Duration::from_millis(20));
        };
        let printed = |name: &str| {
            let mut b = std::fs::read(dir.join(name)).unwrap_or_default();
            b.truncate(MAX_PRINTED);
            String::from_utf8_lossy(&b).into_owned()
        };
        let mut files = Vec::new();
        let mut total = 0usize;
        for name in outputs {
            if let Ok(b) = std::fs::read(dir.join(name)) {
                total += b.len();
                if total > MAX_OUTPUT {
                    return Err(("413 Content Too Large", format!("what {id} wrote is larger than the page takes; `wai {}` writes it", argv.join(" "))));
                }
                files.push(json!({ "name": name, "data": base64(&b) }));
            }
        }
        let mut shown = vec!["wai".to_string()];
        shown.extend(argv);
        Ok(json!({ "op": id, "ok": status.success(), "exit": status.code(), "argv": shown, "stdout": printed(".stdout"), "stderr": printed(".stderr"), "outputs": files }))
    })();
    let _ = std::fs::remove_dir_all(&dir);
    r
}

/// Standard base64 with padding (RFC 4648 §4).
fn base64(b: &[u8]) -> String {
    const T: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut s = String::with_capacity(b.len().div_ceil(3) * 4);
    for ch in b.chunks(3) {
        let n = (u32::from(ch[0]) << 16) | (u32::from(*ch.get(1).unwrap_or(&0)) << 8) | u32::from(*ch.get(2).unwrap_or(&0));
        for i in 0..4 {
            if i <= ch.len() {
                s.push(T[(n >> (18 - 6 * i) & 63) as usize] as char);
            } else {
                s.push('=');
            }
        }
    }
    s
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn base64_is_rfc_4648s() {
        // RFC 4648 §10's test vectors.
        for (i, o) in [("", ""), ("f", "Zg=="), ("fo", "Zm8="), ("foo", "Zm9v"), ("foob", "Zm9vYg=="), ("fooba", "Zm9vYmE="), ("foobar", "Zm9vYmFy")] {
            assert_eq!(base64(i.as_bytes()), o);
        }
    }

    fn head(method: &str, target: &str, headers: &[(&str, &str)]) -> Head {
        Head { method: method.into(), target: target.into(), headers: headers.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect() }
    }

    #[test]
    fn the_api_answers_only_this_origin_by_its_own_name() {
        let ok = [("Host", "127.0.0.1:8080"), ("X-WAI-Studio", "1")];
        assert_eq!(refused(&head("GET", "/_wai/api/state", &ok), 8080), None);
        assert_eq!(refused(&head("GET", "/", &[("Host", "localhost:8080"), ("X-WAI-Studio", "1"), ("Origin", "http://localhost:8080"), ("Sec-Fetch-Site", "same-origin")]), 8080), None);
        // Another name for this machine (a rebound name), another port, no host at all.
        for host in ["evil.example:8080", "127.0.0.1:8081", "127.0.0.1"] {
            assert!(refused(&head("GET", "/", &[("Host", host), ("X-WAI-Studio", "1")]), 8080).is_some(), "{host}");
        }
        assert!(refused(&head("GET", "/", &[("X-WAI-Studio", "1")]), 8080).is_some());
        // Without the header a page of another origin cannot send unasked.
        assert!(refused(&head("POST", "/", &[("Host", "127.0.0.1:8080")]), 8080).is_some());
        // A page of another origin, by its Origin or its fetch site.
        assert!(refused(&head("POST", "/", &[("Host", "127.0.0.1:8080"), ("X-WAI-Studio", "1"), ("Origin", "https://evil.example")]), 8080).is_some());
        assert!(refused(&head("POST", "/", &[("Host", "127.0.0.1:8080"), ("X-WAI-Studio", "1"), ("Sec-Fetch-Site", "cross-site")]), 8080).is_some());
        // Before the stack is up there is no origin to name.
        assert!(refused(&head("GET", "/", &ok), 0).is_some());
    }

    #[test]
    fn the_local_operations_are_operations_of_the_catalogue_by_name() {
        for (id, verb) in LOCAL_OPS {
            let (noun, v) = id.split_once('.').expect("noun.verb");
            assert!(!noun.is_empty() && !v.contains('.'), "{id}");
            assert!(["inspect", "decode", "encode"].contains(verb), "{id} runs `wai {verb}`");
        }
        assert_eq!(query("/op/x?codec=png&width=3", "width"), Some("3"));
        assert_eq!(query("/op/x?codec=png", "height"), None);
    }
}
