// ingress: serve HTTP and XRPC from a wasm module at https://did-key-<multibase>.xrpc.fedproxy.com,
// through the fedproxy did:key ingress relay. A step protocol handler:
// the host calls `serve` once per event.
//
//   start            -> getServiceAuth x2 (host did:key, in parallel), lifetime timer
//   nonce token      -> getRegistrationNonce (relay, POST, did:key service-auth)
//   nonce            -> sign the nonce (host did:key)
//   signature + sub token -> subscribe to the relay (one WebSocket, kept open)
//   message #request -> respond on the same socket (send); describeRepo first reads the repo
//                       (many in flight at once), with a deadline timer that cancels the read
//   message #request for /xrpc/<nsid> in `routes` -> a call to that NSID (a composition edge serves it)
//                       with {method, params, body, authorization}; its output is the 200
//                       body, an error "<status> <Name>: <message>" that status (else 502)
//   forward          -> every request (keepalives aside) is one call to FORWARD with
//                       {method, path, params, headers, body, callerDid}; its output
//                       {status, body, contentType, headers} is the response (the host serves any path);
//                       headers: name -> string, or -> [string] for a repeated one (Set-Cookie); the
//                       relay drops the ones it owns (hop-by-hop, Content-Length, HSTS, Alt-Svc, ...)
//   message #subscribe (a caller's WebSocket on our host) -> ticks on a timer, then close
//   closed relay     -> re-register with backoff, 3 times
//   aliases          -> after every #registered, one com.fedproxy.temp.xrpc.alias frame {hosts}; the relay
//                       maps each host whose _xrpc.<host> TXT names our did:key and answers
//                       alias#result {accepted, rejected}, which the run result carries as `aliases`;
//                       while any host is rejected, the frame goes again every 60 s (the relay's TXT
//                       cache), so a TXT record added later is picked up on this connection
//   quietMs          -> fault injection: leave the relay's keepalive pings unanswered for the first
//                       quietMs of the first connection, so the relay evicts us (1011, after 90 s of
//                       silence) and the re-register path runs against the real relay
//   maxRequests served (and nothing in flight) -> done
//
// The did:key private key stays in the host: the module only ever sees tokens and signatures.

use std::collections::BTreeMap;

use abi::{Call, Event, Handle, Run, entries};

abi::call_ids! {
    #[derive(Clone, PartialEq)]
    enum Id {
        AuthNonce = "auth-nonce",
        AuthSub = "auth-sub",
        Nonce = "nonce",
        Sign = "sign",
        Relay = "relay",
        Lifetime = "lifetime",
        Reconnect = "reconnect",
        Req(String) = "req",
        Deadline(String) = "deadline",
        Tick(String) = "tick",
        AliasRetry = "alias-retry",
    }
}

type Out = abi::Out<Id>;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

const RELAY: &str = "https://xrpc.fedproxy.com";
const SUBSCRIBE: &str = "com.fedproxy.temp.xrpc.subscribe";
const NONCE: &str = "com.fedproxy.temp.xrpc.getRegistrationNonce";
const REGISTRATION: &str = "com.fedproxy.temp.xrpc.registration";
const ALIAS: &str = "com.fedproxy.temp.xrpc.alias";
const ALIAS_RESULT: &str = "com.fedproxy.temp.xrpc.alias#result";
const GET_SERVICE_AUTH: &str = "com.fedfork.exec.didkey.getServiceAuth";
const SIGN: &str = "com.fedfork.exec.didkey.sign";
const DESCRIBE: &str = "com.atproto.repo.describeRepo";
const KEEPALIVE: &str = "/xrpc/com.publicdomainrelay.temp.keepalive.ping";
const HELLO: &str = "/xrpc/com.example.ingress.hello";
const DESCRIBE_PATH: &str = "/xrpc/com.example.ingress.describeRepo";
const TICKS: &str = "com.example.ingress.ticks";
const FORWARD: &str = "com.fedfork.exec.ingress.request";
const MAX_RECONNECTS: u32 = 3;
const MAX_TTL_S: i64 = 3600;
const ALIAS_RETRY_MS: u64 = 60_000;

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct Args {
    now: String,
    #[serde(default)]
    relay: Option<String>,
    #[serde(default)]
    max_requests: Option<Value>,
    #[serde(default)]
    lifetime_ms: Option<Value>,
    #[serde(default)]
    deadline_ms: Option<Value>,
    // NSIDs to forward as calls: a list, or space/comma separated
    #[serde(default)]
    routes: Option<Value>,
    #[serde(default)]
    quiet_ms: Option<Value>,
    // every request to the host's FORWARD call: true, or "true" from a UI
    #[serde(default)]
    forward: Option<Value>,
    // custom host names to register with the relay: a list, or space/comma separated
    #[serde(default)]
    aliases: Option<Value>,
}

fn alias_hosts(run: &Args) -> Vec<String> {
    let list: Vec<&str> = match &run.aliases {
        Some(Value::Array(list)) => list.iter().filter_map(Value::as_str).collect(),
        Some(Value::String(list)) => list.split([' ', ',']).collect(),
        _ => vec![],
    };
    list.into_iter().map(|h| h.trim().to_ascii_lowercase()).filter(|h| !h.is_empty()).collect()
}

fn forwarding(run: &Args) -> bool {
    matches!(&run.forward, Some(Value::Bool(true))) || matches!(&run.forward, Some(Value::String(s)) if s.trim() == "true")
}

#[derive(Default, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
struct State {
    did: Option<String>,
    sub_token: Option<String>,
    nonce: Option<String>,
    signature: Option<String>,
    host: Option<String>,
    served: u64,
    reconnects: u32,
    // relay connections lost, over the whole run
    drops: u32,
    // unix seconds until which keepalives go unanswered (quietMs, first connection only)
    quiet_until: Option<i64>,
    closing: bool,
    // requestId of each proxied request waiting on a call
    pending: Vec<String>,
    // caller subscriptionId -> [ticks left, ms between ticks]
    ticks: BTreeMap<String, [u64; 2]>,
    // the relay's last alias#result {accepted, rejected} on this connection
    aliases: Option<Value>,
}

fn send_frame(out: &mut Out, kind: &str, mut fields: Value) {
    fields["$type"] = json!(format!("{SUBSCRIBE}#{kind}"));
    out.send(Id::Relay, fields);
}

fn respond(out: &mut Out, request_id: &str, status: u16, body: Value, content_type: &str) {
    send_frame(out, "response", json!({ "requestId": request_id, "status": status, "body": body, "contentType": content_type }));
}

// A forward output's headers, as the relay takes them: an object of string or [string] values; any
// other entry is dropped. None when nothing is left.
fn response_headers(headers: &Value) -> Option<Value> {
    let ok = |v: &Value| v.is_string() || v.as_array().is_some_and(|list| list.iter().all(Value::is_string));
    let kept: serde_json::Map<String, Value> = headers.as_object()?.iter().filter(|(name, v)| !name.is_empty() && ok(v)).map(|(k, v)| (k.clone(), v.clone())).collect();
    (!kept.is_empty()).then_some(Value::Object(kept))
}

// Fields from a UI arrive as strings, from a CLI as numbers.
fn number(value: &Option<Value>) -> Option<u64> {
    match value.as_ref()? {
        Value::Number(n) => n.as_u64(),
        Value::String(s) => s.trim().parse().ok(),
        _ => None,
    }
}

fn host_of(origin: &str) -> &str {
    let rest = origin.split_once("://").map_or(origin, |(_, rest)| rest);
    rest.split(['/', ':']).next().unwrap_or(rest)
}

// Unix seconds of an RFC 3339 UTC time (the host's `now`), for a token's exp.
fn unix(now: &str) -> Option<i64> {
    let n = |range: std::ops::Range<usize>| now.get(range)?.parse::<i64>().ok();
    let (y, m, d) = (n(0..4)?, n(5..7)?, n(8..10)?);
    let (hh, mm, ss) = (n(11..13)?, n(14..16)?, n(17..19)?);
    let y = if m <= 2 { y - 1 } else { y };
    let era = y.div_euclid(400);
    let yoe = y - era * 400;
    let doy = (153 * ((m + 9) % 12) + 2) / 5 + d - 1;
    let days = era * 146_097 + yoe * 365 + yoe / 4 - yoe / 100 + doy - 719_468;
    Some(days * 86_400 + hh * 3600 + mm * 60 + ss)
}

fn register(out: &mut Out, aud: &str, exp: i64) {
    out.call(Id::AuthNonce, Call::new("", GET_SERVICE_AUTH).params(json!({ "aud": aud, "lxm": NONCE })));
    out.call(Id::AuthSub, Call::new("", GET_SERVICE_AUTH).params(json!({ "aud": aud, "lxm": SUBSCRIBE, "exp": exp.to_string() })));
}

fn page(did: &str, served: u64) -> Value {
    json!(format!(
        "<!doctype html><meta charset=utf-8><title>atlibs ingress</title>\
         <h1>Served by a wasm module</h1><p>{did} answered this through xrpc.fedproxy.com \
         ({served} requests so far).</p><p>Try <a href=\"{HELLO}?name=you\">{HELLO}</a>.</p>"
    ))
}

fn routed(run: &Args, nsid: &str) -> bool {
    match &run.routes {
        Some(Value::Array(list)) => list.iter().any(|r| r == nsid),
        Some(Value::String(list)) => list.split([' ', ',']).any(|r| r == nsid),
        _ => false,
    }
}

// "<status> <Name>: <message>" from a route's error, else 502 UpstreamFailure.
fn route_error(error: &str) -> (u16, Value) {
    let parsed = error.split_once(' ').and_then(|(code, rest)| Some((code.parse::<u16>().ok().filter(|c| (400..600).contains(c))?, rest.split_once(": ")?)));
    match parsed {
        Some((status, (name, message))) if !name.contains(' ') => (status, json!({ "error": name, "message": message })),
        _ => (502, json!({ "error": "UpstreamFailure", "message": error })),
    }
}

fn on_request(run: &Args, state: &mut State, out: &mut Out, frame: &Value, now: i64) {
    let id = frame["requestId"].as_str().unwrap_or_default();
    let path = frame["path"].as_str().unwrap_or_default();
    let param = |key: &str| frame["params"][key].as_str().unwrap_or_default().trim().to_string();
    let did = state.did.clone().unwrap_or_default();
    if path == KEEPALIVE {
        if state.drops == 0 && state.quiet_until.is_some_and(|until| now < until) {
            return;
        }
        return respond(out, id, 200, json!({}), "application/json");
    }
    state.served += 1;
    if forwarding(run) {
        let input = json!({ "method": frame["method"], "path": path, "params": frame["params"], "headers": frame["headers"], "body": frame["body"], "callerDid": frame["callerDid"] });
        out.call(Id::Req(id.into()), Call::new("", FORWARD).input(input));
        out.timer(Id::Deadline(id.into()), number(&run.deadline_ms).unwrap_or(8000));
        state.pending.push(id.to_string());
        return;
    }
    if let Some(nsid) = path.strip_prefix("/xrpc/").filter(|nsid| routed(run, nsid)) {
        let input = json!({ "method": frame["method"], "params": frame["params"], "body": frame["body"], "authorization": frame["headers"]["authorization"] });
        out.call(Id::Req(id.into()), Call::new("", nsid).input(input));
        out.timer(Id::Deadline(id.into()), number(&run.deadline_ms).unwrap_or(8000));
        state.pending.push(id.to_string());
        return;
    }
    match path {
        "" | "/" => respond(out, id, 200, page(&did, state.served), "text/html; charset=utf-8"),
        HELLO => {
            let name = param("name");
            let name: String = if name.is_empty() { "world".into() } else { name.chars().take(64).collect() };
            respond(out, id, 200, json!({ "message": format!("hello {name}"), "did": did, "served": state.served, "now": run.now }), "application/json");
        }
        DESCRIBE_PATH => {
            let repo = param("repo");
            if !(repo.starts_with("did:plc:") || repo.starts_with("did:web:")) {
                return respond(out, id, 400, json!({ "error": "InvalidRequest", "message": "repo must be a did:plc or did:web" }), "application/json");
            }
            out.call(Id::Req(id.into()), Call::new("", DESCRIBE).params(json!({ "repo": repo })).service(repo));
            out.timer(Id::Deadline(id.into()), number(&run.deadline_ms).unwrap_or(8000));
            state.pending.push(id.to_string());
        }
        _ => respond(out, id, 404, json!({ "error": "MethodNotImplemented", "message": format!("{path} is not served here") }), "application/json"),
    }
}

fn on_subscribe(state: &mut State, out: &mut Out, frame: &Value) {
    let id = frame["subscriptionId"].as_str().unwrap_or_default().to_string();
    let nsid = frame["nsid"].as_str().unwrap_or_default();
    if nsid != TICKS {
        return send_frame(out, "subscriptionClose", json!({ "subscriptionId": id, "code": 1003, "reason": format!("{nsid} is not served here") }));
    }
    let param = |key: &str| frame["params"][key].as_str().and_then(|s| s.parse::<u64>().ok());
    let count = param("n").unwrap_or(3).clamp(1, 20);
    let ms = param("ms").unwrap_or(1000).clamp(200, 10_000);
    send_frame(out, "subscriptionOpen", json!({ "subscriptionId": id }));
    out.timer(Id::Tick(id.clone()), ms);
    state.ticks.insert(id, [count, ms]);
}

fn on_frame(run: &Args, state: &mut State, out: &mut Out, frame: &Value, now: i64) {
    if frame["$type"] == ALIAS_RESULT {
        state.aliases = Some(json!({ "accepted": frame["accepted"], "rejected": frame["rejected"] }));
        if frame["rejected"].as_object().is_some_and(|rejected| !rejected.is_empty()) {
            out.timer(Id::AliasRetry, ALIAS_RETRY_MS);
        }
        return;
    }
    let kind = frame["$type"].as_str().and_then(|t| t.strip_prefix(SUBSCRIBE)).and_then(|t| t.strip_prefix('#'));
    match kind {
        Some("registered") => {
            state.host = frame["subdomain"].as_str().map(|sub| format!("{sub}.{}", host_of(run.relay.as_deref().unwrap_or(RELAY))));
            state.reconnects = 0;
            if state.drops == 0 && state.quiet_until.is_none() {
                state.quiet_until = number(&run.quiet_ms).map(|ms| now + (ms / 1000) as i64);
            }
            let hosts = alias_hosts(run);
            if !hosts.is_empty() {
                out.send(Id::Relay, json!({ "$type": ALIAS, "hosts": hosts }));
            }
        }
        Some("request") => on_request(run, state, out, frame, now),
        Some("subscribe") => on_subscribe(state, out, frame),
        Some("subscriptionCancel" | "subscriptionClose") => {
            if let Some(id) = frame["subscriptionId"].as_str() {
                if state.ticks.remove(id).is_some() {
                    out.cancel(Id::Tick(id.into()));
                }
            }
        }
        _ => {}
    }
}

fn handle(Run { event, state, args: run }: Run<Args>) -> Result<Handle, String> {
    let mut state: State = abi::state(state)?;
    let mut out = Out::default();
    let mut done = false;
    let relay = run.relay.clone().unwrap_or_else(|| RELAY.into());
    let aud = format!("did:web:{}", host_of(&relay));
    let now = unix(&run.now).ok_or("now must be an RFC 3339 UTC time")?;
    let lifetime_ms = number(&run.lifetime_ms);
    let exp = now + lifetime_ms.map_or(MAX_TTL_S, |ms| (ms / 1000) as i64 + 60).min(MAX_TTL_S);
    let id_of = |id: &str| id.parse::<Id>().ok();

    match event {
        Event::Start => {
            register(&mut out, &aud, exp);
            if let Some(ms) = lifetime_ms {
                out.timer(Id::Lifetime, ms);
            }
        }
        Event::Result { id, nsid, result } => match id_of(&id) {
            // A send that could not be written; the relay times the request out.
            Some(Id::Relay) => {}
            Some(Id::Req(request_id)) => {
                if let Some(at) = state.pending.iter().position(|p| *p == request_id) {
                    state.pending.remove(at);
                    out.cancel(Id::Deadline(request_id.clone()));
                    match result.output() {
                        Ok(output) if nsid == FORWARD => {
                            let status = output["status"].as_u64().filter(|s| (100..600).contains(s)).unwrap_or(200) as u16;
                            let mut frame = json!({ "requestId": request_id, "status": status, "body": output["body"], "contentType": output["contentType"].as_str().unwrap_or("application/json") });
                            if let Some(headers) = response_headers(&output["headers"]) {
                                frame["headers"] = headers;
                            }
                            send_frame(&mut out, "response", frame);
                        }
                        Ok(output) if routed(&run, &nsid) => respond(&mut out, &request_id, 200, output.clone(), "application/json"),
                        Err(error) if nsid == FORWARD || routed(&run, &nsid) => {
                            let (status, body) = route_error(&error);
                            respond(&mut out, &request_id, status, body, "application/json");
                        }
                        Ok(output) => respond(&mut out, &request_id, 200, json!({ "did": output["did"], "handle": output["handle"], "collections": output["collections"], "servedBy": state.did }), "application/json"),
                        Err(error) => respond(&mut out, &request_id, 502, json!({ "error": "UpstreamFailure", "message": error }), "application/json"),
                    }
                }
            }
            call => {
                let output = result.output().map_err(|error| format!("{id}: {error}"))?;
                match call {
                    Some(Id::AuthNonce) => {
                        let did = output["did"].as_str().ok_or("getServiceAuth returned no did")?.to_string();
                        out.call(Id::Nonce, Call::new("", NONCE).service(&relay).input(json!({ "key": did, "signatures": [] })));
                        state.did = Some(did);
                    }
                    Some(Id::AuthSub) => state.sub_token = output["token"].as_str().map(str::to_string),
                    Some(Id::Nonce) => {
                        let nonce = output["nonce"].as_str().ok_or("relay returned no nonce")?.to_string();
                        out.call(Id::Sign, Call::new("", SIGN).input(json!({ "data": nonce })));
                        state.nonce = Some(nonce);
                    }
                    Some(Id::Sign) => state.signature = output["signature"].as_str().map(str::to_string),
                    _ => {}
                }
                // Both halves of the registration are in (they race): connect.
                if let (Some(did), Some(token), Some(nonce), Some(signature)) = (&state.did, &state.sub_token, &state.nonce, &state.signature) {
                    let registration = json!({ "$type": REGISTRATION, "key": did, "nonce": nonce, "signatures": [{ "key": did, "signature": signature }] });
                    out.call(Id::Relay, Call::new("", SUBSCRIBE).service(&relay).subscribe().close_on_overflow().params(json!({ "did": did, "registration": registration.to_string(), "service_auth": token })));
                    (state.nonce, state.signature, state.sub_token) = (None, None, None);
                }
            }
        },
        Event::Message { id, data, .. } if id_of(&id) == Some(Id::Relay) => on_frame(&run, &mut state, &mut out, &data, now),
        Event::Closed { id, error } if id_of(&id) == Some(Id::Relay) => {
            for request_id in state.pending.drain(..) {
                out.cancel(Id::Req(request_id.clone()));
                out.cancel(Id::Deadline(request_id));
            }
            for sub in std::mem::take(&mut state.ticks).into_keys() {
                out.cancel(Id::Tick(sub));
            }
            state.host = None;
            if state.aliases.take().is_some_and(|a| a["rejected"].as_object().is_some_and(|r| !r.is_empty())) {
                out.cancel(Id::AliasRetry);
            }
            state.drops += 1;
            if state.closing {
                done = true;
            } else if state.reconnects < MAX_RECONNECTS {
                state.reconnects += 1;
                out.timer(Id::Reconnect, 1000 << state.reconnects);
            } else {
                return Err(format!("relay closed {MAX_RECONNECTS} times: {}", error.as_deref().unwrap_or("closed")));
            }
        }
        Event::Timer { id } => match id_of(&id) {
            Some(Id::Lifetime) => state.closing = true,
            Some(Id::Reconnect) => register(&mut out, &aud, exp),
            Some(Id::AliasRetry) if state.host.is_some() && !state.closing => out.send(Id::Relay, json!({ "$type": ALIAS, "hosts": alias_hosts(&run) })),
            Some(Id::Deadline(request_id)) => {
                if let Some(at) = state.pending.iter().position(|p| *p == request_id) {
                    state.pending.remove(at);
                    out.cancel(Id::Req(request_id.clone()));
                    respond(&mut out, &request_id, 504, json!({ "error": "UpstreamTimeout", "message": "the upstream call took too long" }), "application/json");
                }
            }
            Some(Id::Tick(sub)) => {
                if let Some([left, ms]) = state.ticks.get_mut(&sub) {
                    *left -= 1;
                    let (left, ms) = (*left, *ms);
                    send_frame(&mut out, "subscriptionEvent", json!({ "subscriptionId": sub, "message": { "left": left, "time": run.now, "did": state.did } }));
                    if left == 0 {
                        state.ticks.remove(&sub);
                        send_frame(&mut out, "subscriptionClose", json!({ "subscriptionId": sub, "code": 1000 }));
                    } else {
                        out.timer(Id::Tick(sub), ms);
                    }
                }
            }
            _ => {}
        },
        _ => {}
    }

    if let Some(max) = number(&run.max_requests) {
        state.closing |= state.served >= max;
    }
    done |= state.closing && state.pending.is_empty() && state.ticks.is_empty();
    let url = state.host.as_ref().map(|host| format!("https://{host}/"));
    let mut result = json!({ "did": state.did, "url": url, "served": state.served, "drops": state.drops });
    if let Some(aliases) = &state.aliases {
        result["aliases"] = aliases.clone();
    }
    if done {
        out.finish(result);
    } else {
        out.result(result);
    }
    Ok(out.handle().state(serde_json::to_value(&state).map_err(|e| e.to_string())?))
}

entries! { serve => handle }