use axum::extract::{Path, Query};
use axum::http::{header, HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use serde::Deserialize;

static STUB_X86_64: &[u8] = include_bytes!(env!("STUB_X86_64_PATH"));
static STUB_AARCH64: &[u8] = include_bytes!(env!("STUB_AARCH64_PATH"));

const SCRIPT_TEMPLATE: &str = r#"#!/bin/sh
set -eu
SESSION='__SESSION__'
BASE='__BASE__'
WS='__WS__'
A=$(uname -m)
case "$A" in
  x86_64|amd64) ARCH=x86_64 ;;
  aarch64|arm64) ARCH=aarch64 ;;
  *) echo "rsh: unsupported arch $A" >&2; exit 1 ;;
esac
TMP=$(mktemp /tmp/rsh.XXXXXX)
trap 'rm -f "$TMP"' EXIT
if command -v curl >/dev/null 2>&1; then
  curl -fsSL "$BASE/rsh/$ARCH" -o "$TMP"
elif command -v wget >/dev/null 2>&1; then
  wget -qO "$TMP" "$BASE/rsh/$ARCH"
else
  echo "rsh: need curl or wget" >&2; exit 1
fi
chmod +x "$TMP"
printf 'Password (empty for none): ' >&2
stty -echo 2>/dev/null || true
IFS= read -r PW || PW=''
stty echo 2>/dev/null || true
echo >&2
trap - EXIT
if [ -n "$PW" ]; then
  exec "$TMP" --url "$WS/ws/stub" --session "$SESSION" --password "$PW"
else
  exec "$TMP" --url "$WS/ws/stub" --session "$SESSION"
fi
"#;

#[derive(Deserialize)]
pub struct RunQuery {
    s: Option<String>,
}

pub async fn run_sh(headers: HeaderMap, Query(q): Query<RunQuery>) -> Response {
    let session = q.s.unwrap_or_else(|| "default".into());
    if !is_valid_session(&session) {
        return (StatusCode::BAD_REQUEST, "invalid session id").into_response();
    }
    let (scheme, ws_scheme) = detect_scheme(&headers);
    let host = headers
        .get(header::HOST)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("localhost");
    let base = format!("{scheme}://{host}");
    let ws = format!("{ws_scheme}://{host}");
    let script = SCRIPT_TEMPLATE
        .replace("__SESSION__", &session)
        .replace("__BASE__", &base)
        .replace("__WS__", &ws);
    (
        [(header::CONTENT_TYPE, "text/x-shellscript; charset=utf-8")],
        script,
    )
        .into_response()
}

pub async fn stub(Path(arch): Path<String>) -> Response {
    let bytes: &[u8] = match arch.as_str() {
        "x86_64" | "amd64" => STUB_X86_64,
        "aarch64" | "arm64" => STUB_AARCH64,
        _ => return StatusCode::NOT_FOUND.into_response(),
    };
    if bytes.is_empty() {
        return (StatusCode::SERVICE_UNAVAILABLE, "stub not embedded in this build").into_response();
    }
    (
        [
            (header::CONTENT_TYPE, "application/octet-stream"),
            (
                header::CONTENT_DISPOSITION,
                "attachment; filename=\"rsh\"",
            ),
        ],
        bytes,
    )
        .into_response()
}

fn is_valid_session(s: &str) -> bool {
    !s.is_empty()
        && s.len() <= 64
        && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-')
}

fn detect_scheme(headers: &HeaderMap) -> (&'static str, &'static str) {
    let proto = headers
        .get("x-forwarded-proto")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("http");
    if proto == "https" {
        ("https", "wss")
    } else {
        ("http", "ws")
    }
}