// A step protocol toy: one handler, every event kind.
// start: describe the user's repo and subscribe to a repo event stream, with a stop timer, at once.
// result: remember the handle. message: count; at `limit`, cancel the stream and finish.
// input: echo the data up the stream. timer or closed: finish.
use abi::records::exec::iroh;
use abi::{Call, Did, Event, Handle, Run, entries};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

const FEED: &str = "feed";
const STOP: &str = "stop";

#[derive(Deserialize)]
struct Args {
    did: Did,
    service: String,
    #[serde(default)]
    limit: Option<Value>,
    #[serde(default)]
    seconds: Option<Value>,
}

#[derive(Default, Deserialize, Serialize)]
struct State {
    handle: Option<String>,
    count: u64,
    last: Option<Value>,
}

// Form fields arrive as strings; tests and hosts may pass numbers.
fn number(value: Option<&Value>, default: u64, max: u64) -> Result<u64, String> {
    let n = match value {
        None | Some(Value::Null) => return Ok(default),
        Some(Value::Number(n)) => n.as_u64(),
        Some(Value::String(s)) => s.trim().parse().ok(),
        _ => None,
    };
    n.filter(|n| (1..=max).contains(n)).ok_or_else(|| format!("must be 1 to {max}"))
}

fn finish(state: &State, why: &str) -> Handle {
    Handle::new().result(json!({ "handle": state.handle, "count": state.count, "last": state.last, "stopped": why })).done()
}

fn watch_handler(run: Run<Args>) -> Result<Handle, String> {
    let args = run.args;
    let limit = number(args.limit.as_ref(), 10, 1000).map_err(|e| format!("limit {e}"))?;
    let mut state: State = abi::state(run.state)?;
    Ok(match run.event {
        Event::Start => {
            if !args.service.starts_with("https://") {
                return Err("service must be an https origin".into());
            }
            let seconds = number(args.seconds.as_ref(), 10, 50).map_err(|e| format!("seconds {e}"))?;
            Handle::new()
                .call(Call::new("who", "com.atproto.repo.describeRepo").params(json!({ "repo": args.did })))
                .call(Call::new(FEED, "com.atproto.sync.subscribeRepos").service(args.service).subscribe())
                .timer(STOP, seconds * 1000)
                .state(json!(state))
        }
        Event::Result { id, result, .. } if id == "who" => {
            state.handle = result.output().ok().and_then(|o| o["handle"].as_str()).map(String::from);
            Handle::new().state(json!(state))
        }
        Event::Message { data, .. } => {
            state.count += 1;
            state.last = Some(json!({ "$type": data["$type"], "seq": data["seq"] }));
            if state.count >= limit { finish(&state, "limit").cancel(FEED) } else { Handle::new().state(json!(state)) }
        }
        Event::Input { data } => Handle::new().send(FEED, data),
        Event::Timer { .. } => finish(&state, "time"),
        Event::Closed { error, .. } => finish(&state, error.as_deref().unwrap_or("closed")),
        _ => Handle::new(),
    })
}

// session: a com.fedfork.exec.iroh.connect session: stream a guest's
// output (ssh `command`, or a raw `alpn` stream), write `send` after the first stdout, echo inputs
// up the session, and finish with everything it printed and its exit code when it closes.
const SESSION: &str = "session";

#[derive(Deserialize)]
struct SessionArgs {
    ticket: String,
    #[serde(default)]
    alpn: Option<String>,
    #[serde(default, rename = "hostKeys")]
    host_keys: Option<Value>,
    #[serde(default)]
    command: Option<String>,
    #[serde(default)]
    send: Option<String>,
}

#[derive(Default, Deserialize, Serialize)]
struct SessionState {
    stdout: String,
    stderr: String,
    #[serde(rename = "exitCode")]
    exit_code: Value,
    sent: bool,
}

fn session_handler(run: Run<SessionArgs>) -> Result<Handle, String> {
    let args = run.args;
    let mut state: SessionState = abi::state(run.state)?;
    Ok(match run.event {
        Event::Start => {
            let input = iroh::Connect { ticket: args.ticket, alpn: args.alpn, host_keys: args.host_keys, command: args.command, stdin: None, send: None };
            Handle::new().call(Call::new(SESSION, "com.fedfork.exec.iroh.connect").input(json!(input)).subscribe().close_on_overflow()).state(json!(state))
        }
        Event::Message { data, .. } => {
            let mut out = Handle::new();
            if let Some(text) = data["stdout"].as_str() {
                state.stdout.push_str(text);
                if let (false, Some(line)) = (state.sent, args.send) {
                    state.sent = true;
                    out = out.send(SESSION, json!(line));
                }
            }
            if let Some(text) = data["stderr"].as_str() {
                state.stderr.push_str(text);
            }
            if let Some(code) = data.get("exitCode") {
                state.exit_code = code.clone();
            }
            out.state(json!(state))
        }
        Event::Input { data } => Handle::new().send(SESSION, data),
        Event::Closed { error, .. } => Handle::new()
            .result(json!({ "stdout": state.stdout, "stderr": state.stderr, "exitCode": state.exit_code, "error": error }))
            .done(),
        _ => Handle::new(),
    })
}

// connect: the one-shot com.fedfork.exec.iroh.connect (ssh `command` with `stdin`, or `alpn` with
// `send`), the input passed through as is; finishes with the call's output ({output, exitCode?}).
fn connect_handler(run: Run<Value>) -> Result<Handle, String> {
    Ok(match run.event {
        Event::Start => Handle::new().call(Call::new("connect", "com.fedfork.exec.iroh.connect").input(run.args)),
        Event::Result { result, .. } => Handle::new().result(result.output()?.clone()).done(),
        _ => Handle::new(),
    })
}

entries! { watch => watch_handler, session => session_handler, connect => connect_handler }