//! Serve fresh world files and stream the /api mount through standard WASI HTTP.
//! Stdin EOF stops the server, including waits for connections and HTTP.
use clap::Parser;
use std::{io, net::SocketAddr};
use wasi::{
    http::{
        outgoing_handler,
        types::{Fields, Method, OutgoingBody, OutgoingRequest, Scheme},
    },
    io::{
        poll::{Pollable, poll},
        streams::{InputStream, OutputStream, StreamError},
    },
    sockets::{
        instance_network::instance_network, network::*, tcp_create_socket::create_tcp_socket,
    },
};
#[derive(Parser)]
struct Options {
    #[arg(long, default_value = "0.0.0.0:8080")]
    listen: SocketAddr,
    #[arg(long, default_value = "/srv/www")]
    root: std::path::PathBuf,
}
fn error(e: impl std::fmt::Debug) -> io::Error {
    io::Error::other(format!("{e:?}"))
}
fn wait(stdin: &InputStream, event: &Pollable) -> io::Result<()> {
    let stop = stdin.subscribe();
    loop {
        let ready = poll(&[&stop, event]);
        match stdin.read(65536) {
            Err(StreamError::Closed) => std::process::exit(0),
            Err(e) => return Err(error(e)),
            Ok(_) => {}
        }
        if ready.contains(&1) {
            return Ok(());
        }
    }
}
fn write(stdin: &InputStream, stream: &OutputStream, mut bytes: &[u8]) -> io::Result<()> {
    while !bytes.is_empty() {
        wait(stdin, &stream.subscribe())?;
        let count = (stream.check_write().map_err(error)? as usize).min(bytes.len());
        stream.write(&bytes[..count]).map_err(error)?;
        bytes = &bytes[count..];
    }
    stream.flush().map_err(error)?;
    loop {
        wait(stdin, &stream.subscribe())?;
        if stream.check_write().map_err(error)? > 0 {
            return Ok(());
        }
    }
}
fn wait_response(stdin: &InputStream, event: &Pollable, client: &InputStream) -> io::Result<()> {
    let stop = stdin.subscribe();
    let disconnect = client.subscribe();
    loop {
        let ready = poll(&[&stop, event, &disconnect]);
        match stdin.read(65536) {
            Err(StreamError::Closed) => std::process::exit(0),
            Err(e) => return Err(error(e)),
            Ok(_) => {}
        }
        if ready.contains(&2) {
            client.read(65536).map_err(error)?;
        }
        if ready.contains(&1) {
            return Ok(());
        }
    }
}
fn proxy(
    stdin: &InputStream,
    client: &InputStream,
    method: &str,
    path: &str,
    headers: (Option<&str>, Option<&str>, Option<&str>),
    body: &[u8],
    output: &OutputStream,
) -> io::Result<()> {
    let (content_type, origin, host) = headers;
    let headers = Fields::new();
    headers
        .append("content-length", &body.len().to_string().into_bytes())
        .map_err(error)?;
    if let Some(value) = content_type {
        headers
            .append("content-type", value.as_bytes())
            .map_err(error)?;
    }
    for (key, value) in [("origin", origin), ("x-forwarded-host", host)] {
        if let Some(value) = value {
            headers.append(key, value.as_bytes()).map_err(error)?;
        }
    }
    let request = OutgoingRequest::new(headers);
    let method = match method {
        "GET" => Method::Get,
        "POST" => Method::Post,
        "PUT" => Method::Put,
        "DELETE" => Method::Delete,
        "PATCH" => Method::Patch,
        "HEAD" => Method::Head,
        "OPTIONS" => Method::Options,
        other => Method::Other(other.into()),
    };
    request.set_method(&method).map_err(error)?;
    request.set_scheme(Some(&Scheme::Http)).map_err(error)?;
    request.set_authority(Some("host.sausage")).map_err(error)?;
    request
        .set_path_with_query(Some(&format!("/{path}")))
        .map_err(error)?;
    let outgoing = request.body().map_err(error)?;
    let future = outgoing_handler::handle(request, None).map_err(error)?;
    if !body.is_empty() {
        let stream = outgoing.write().map_err(error)?;
        write(stdin, &stream, body)?;
    }
    OutgoingBody::finish(outgoing, None).map_err(error)?;
    wait_response(stdin, &future.subscribe(), client)?;
    let response = future
        .get()
        .ok_or_else(|| error("missing HTTP response"))?
        .map_err(error)?
        .map_err(error)?;
    let status = response.status();
    let content_type = response
        .headers()
        .get("content-type")
        .first()
        .map(|s| String::from_utf8_lossy(s).into_owned())
        .unwrap_or_else(|| "application/octet-stream".into());
    let incoming = response.consume().map_err(error)?;
    let stream = incoming.stream().map_err(error)?;
    write(stdin, output, format!("HTTP/1.1 {status} Response\r\nContent-Type: {content_type}\r\nConnection: close\r\n\r\n").as_bytes())?;
    loop {
        wait_response(stdin, &stream.subscribe(), client)?;
        match stream.read(65536) {
            Ok(chunk) => write(stdin, output, &chunk)?,
            Err(StreamError::Closed) => break,
            Err(e) => return Err(error(e)),
        }
    }
    Ok(())
}
fn serve(
    stdin: &InputStream,
    input: &InputStream,
    output: &OutputStream,
    root: &std::path::Path,
) -> io::Result<()> {
    let mut bytes = Vec::new();
    let end = loop {
        if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") {
            break index + 4;
        }
        wait(stdin, &input.subscribe())?;
        bytes.extend(input.read(65536).map_err(error)?);
    };
    let headers = std::str::from_utf8(&bytes[..end])
        .map_err(error)?
        .to_owned();
    let mut lines = headers.lines();
    let mut line = lines.next().unwrap_or("").split_whitespace();
    let method = line.next().unwrap_or("");
    let path = line.next().unwrap_or("");
    let headers: Vec<_> = lines
        .filter_map(|s| s.split_once(':'))
        .map(|(k, v)| (k.to_ascii_lowercase(), v.trim()))
        .collect();
    let header = |key: &str| headers.iter().find(|(k, _)| k == key).map(|(_, v)| *v);
    let length = header("content-length")
        .unwrap_or("0")
        .parse::<usize>()
        .map_err(error)?;
    while bytes.len() < end + length {
        wait(stdin, &input.subscribe())?;
        bytes.extend(input.read(65536).map_err(error)?);
    }
    if let Some(path) = path.strip_prefix("/api/") {
        return proxy(
            stdin,
            input,
            method,
            path,
            (header("content-type"), header("origin"), header("host")),
            &bytes[end..end + length],
            output,
        );
    }
    let mut file = root.to_path_buf();
    let path = path.split('?').next().unwrap_or("");
    for part in std::path::Path::new(path).components() {
        match part {
            std::path::Component::Normal(name) => file.push(name),
            std::path::Component::RootDir | std::path::Component::CurDir => {}
            _ => return not_found(stdin, output),
        }
    }
    if file.is_dir() {
        file.push("index.html");
    }
    let body = match std::fs::read(&file) {
        Ok(bytes) => bytes,
        Err(_) => return not_found(stdin, output),
    };
    let kind = match file.extension().and_then(|s| s.to_str()).unwrap_or("") {
        "html" => "text/html; charset=utf-8",
        "css" => "text/css; charset=utf-8",
        "js" | "mjs" => "text/javascript; charset=utf-8",
        "json" => "application/json",
        "svg" => "image/svg+xml",
        "png" => "image/png",
        "jpg" | "jpeg" => "image/jpeg",
        "ico" => "image/x-icon",
        "txt" | "md" => "text/plain; charset=utf-8",
        "wasm" => "application/wasm",
        _ => "application/octet-stream",
    };
    let status = 200;
    if method == "HEAD" {
        return write(stdin, output, format!("HTTP/1.1 200 OK\r\nContent-Type: {kind}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes());
    }
    write(stdin, output,format!("HTTP/1.1 {status} Response\r\nContent-Type: {kind}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",body.len()).as_bytes())?;
    write(stdin, output, &body)
}
fn not_found(stdin: &InputStream, output: &OutputStream) -> io::Result<()> {
    write(stdin, output, b"HTTP/1.1 404 Not Found\r\nContent-Type: text/plain\r\nContent-Length: 9\r\nConnection: close\r\n\r\nNot found")
}
fn main() -> io::Result<()> {
    let options = Options::parse();
    let address = options.listen;
    let stdin = wasi::cli::stdin::get_stdin();
    let (family, address) = match address {
        SocketAddr::V4(a) => (
            IpAddressFamily::Ipv4,
            IpSocketAddress::Ipv4(Ipv4SocketAddress {
                port: a.port(),
                address: a.ip().octets().into(),
            }),
        ),
        SocketAddr::V6(a) => (
            IpAddressFamily::Ipv6,
            IpSocketAddress::Ipv6(Ipv6SocketAddress {
                port: a.port(),
                flow_info: a.flowinfo(),
                address: a.ip().segments().into(),
                scope_id: a.scope_id(),
            }),
        ),
    };
    let socket = create_tcp_socket(family).map_err(error)?;
    socket
        .start_bind(&instance_network(), address)
        .map_err(error)?;
    wait(&stdin, &socket.subscribe())?;
    socket.finish_bind().map_err(error)?;
    socket.start_listen().map_err(error)?;
    wait(&stdin, &socket.subscribe())?;
    socket.finish_listen().map_err(error)?;
    loop {
        match socket.accept() {
            Ok((_connection, input, output)) => {
                let _ = serve(&stdin, &input, &output, &options.root);
            }
            Err(ErrorCode::WouldBlock) => wait(&stdin, &socket.subscribe())?,
            Err(e) => return Err(error(e)),
        }
    }
}