// What the bidder reports about the VM for a receipt (iroh ticket, webrtc fingerprints, host keys,
// signals, termination), and a command on the guest over a host transport. wasm holds no socket:
// the module names the guest and the host connects (com.fedfork.exec.webrtc.connect / iroh.connect).
use abi::records::exec::{iroh, webrtc};
use abi::{AtUri, Call, Fan, Handle, On, Run};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

use crate::calls::{hex16, read, receipt_uri};
use crate::ids::CallId;
use crate::input::{Args, Polls, Step, at_uri, did, stdin, strong, tag, tags};
use crate::market::Receipt;
use crate::nsid::{COMPUTE_EVENTS, EVENT, IROH, WEBRTC};

// The guest's own event streams per tag (collection, kind), and the reads a guest step makes of each.
const STREAMS: [(&str, &str); 2] = [("guest.signalWebRTC", "answer"), ("guest.publicKey", "age-x25519")];
pub const STREAM_MAX: usize = 32;

#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Via {
    Iroh,
    Webrtc,
}

// ---- report records -----------------------------------------------------------------------------
// The bidder (socialweb-computer-kcp internal/market/bidder adapters.go recordKeys, syncs.go
// reportSyncKey) writes each report's payload at a key derived from the receipt:
//   <prefix>-payload-<hex(sha256(key))[0..32]>, key = onNetwork:[<tag>:]<receipt> | hostKeys:<receipt> | terminated:<receipt>
// so those are read directly. A guest's own events are streams, key
// guest:<nsid>:<tag>:<kind>:<n>:<receipt>, read send/recv style: n = 0, 1, ... until one is missing.

fn report(bidder: &str, collection: &str, prefix: &str, key: &str) -> AtUri {
    AtUri::new(bidder, format!("{COMPUTE_EVENTS}{collection}"), format!("{prefix}-payload-{}", hex16(key)))
}

pub fn on_network(bidder: &str, receipt: &str, tag: Option<&str>) -> AtUri {
    let key = match tag {
        Some(tag) => format!("onNetwork:{tag}:{receipt}"),
        None => format!("onNetwork:{receipt}"),
    };
    report(bidder, "vm.onNetwork", "onnetwork", &key)
}

pub fn host_keys(bidder: &str, receipt: &str) -> AtUri {
    report(bidder, "vm.hostKeys", "hostkeys", &format!("hostKeys:{receipt}"))
}

pub fn terminated(bidder: &str, receipt: &str) -> AtUri {
    report(bidder, "vm.terminated", "termination", &format!("terminated:{receipt}"))
}

pub fn report_refs(bidder: &str, receipt: &str, tags: &[String]) -> Vec<AtUri> {
    let mut refs = vec![on_network(bidder, receipt, None), host_keys(bidder, receipt), terminated(bidder, receipt), by_nsid(bidder, receipt, "vm.registerIdentity")];
    refs.extend(tags.iter().map(|tag| on_network(bidder, receipt, Some(tag))));
    refs
}

// A report keyed <nsid>:<receipt>: its prefix is the nsid's letters and digits, lowercased.
fn by_nsid(bidder: &str, receipt: &str, collection: &str) -> AtUri {
    let nsid = format!("{COMPUTE_EVENTS}{collection}");
    let prefix: String = nsid.to_lowercase().chars().filter(char::is_ascii_alphanumeric).collect();
    report(bidder, collection, &prefix, &format!("{nsid}:{receipt}"))
}

/// One guest event stream of a receipt and the index read next: the guest's (the bidder writes
/// guest-payload-<hex16("guest:<key>")>) or the requester's own signals (signal-payload-, signal-event-).
#[derive(Clone, Default, Serialize, Deserialize)]
pub struct Stream {
    collection: String,
    tag: String,
    kind: String,
    pub next: usize,
    pub done: bool,
}

// The requester's own signal streams per tag (signal entry: webrtc offers, age keys).
const OWN: [(&str, &str); 2] = [("guest.signalWebRTC", "offer"), ("guest.publicKey", "age-x25519")];

impl Stream {
    pub fn new(collection: &str, tag: &str, kind: &str) -> Self {
        Stream { collection: collection.into(), tag: tag.into(), kind: kind.into(), ..Stream::default() }
    }

    pub fn of(tags: &[String]) -> Vec<Stream> {
        tags.iter().flat_map(|tag| STREAMS.map(|(collection, kind)| Stream::new(collection, tag, kind))).collect()
    }

    pub fn own(tags: &[String]) -> Vec<Stream> {
        tags.iter().flat_map(|tag| OWN.map(|(collection, kind)| Stream::new(collection, tag, kind))).collect()
    }

    fn key(&self, receipt: &str) -> String {
        format!("{COMPUTE_EVENTS}{}:{}:{}:{}:{receipt}", self.collection, self.tag, self.kind, self.next)
    }

    pub fn uri(&self, bidder: &str, receipt: &str) -> AtUri {
        report(bidder, &self.collection, "guest", &format!("guest:{}", self.key(receipt)))
    }

    /// The payload and market.event of the requester's own signal at `next`.
    pub fn signal(&self, me: &str, receipt: &str) -> (AtUri, AtUri) {
        let key = format!("signal:{}", self.key(receipt));
        (report(me, &self.collection, "signal", &key), AtUri::new(me, EVENT, format!("signal-event-{}", hex16(&format!("event:{key}")))))
    }
}

/// A guest record read: its kind (the collection under compute.events), and its value or the error.
#[derive(Clone, Serialize, Deserialize)]
pub struct GuestEvent {
    kind: String,
    uri: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    value: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<String>,
}

impl GuestEvent {
    pub fn read(uri: &str, result: Result<Value, String>) -> Self {
        let kind = uri.parse::<AtUri>().map(|u| u.collection.strip_prefix(COMPUTE_EVENTS).unwrap_or(&u.collection).to_string()).unwrap_or_default();
        match result {
            Ok(mut output) => GuestEvent { kind, uri: uri.into(), value: Some(output["value"].take()), error: None },
            Err(error) => GuestEvent { kind, uri: uri.into(), value: None, error: Some(error) },
        }
    }

    /// A report the bidder has written: a read that found a record.
    pub fn found(uri: &str, result: &Result<Value, String>) -> Option<Self> {
        let output = result.as_ref().ok().filter(|o| o["value"].is_object())?;
        Some(Self::read(uri, Ok(output.clone())))
    }
}

#[derive(Clone, Default, Serialize, Deserialize)]
pub struct Webrtc {
    fingerprint: Value,
    pub tag: Value,
}

/// What the reports say. vm.onNetwork: iroh tickets have no kind (iroh:// prefix); kind "webrtc"
/// carries the DTLS fingerprint.
#[derive(Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Reports {
    pub iroh: Option<String>,
    pub webrtc: Vec<Webrtc>,
    pub host_keys: Option<Value>,
    pub node_id: Option<String>,
    pub signals: Vec<Value>,
    pub terminated: bool,
}

impl Reports {
    pub fn of(events: &[GuestEvent]) -> Self {
        let of = |kind: &'static str| events.iter().filter(move |e| e.kind == kind).filter_map(|e| e.value.as_ref());
        Reports {
            iroh: of("vm.onNetwork").find_map(|v| v["address"].as_str().filter(|a| a.starts_with("iroh://")).map(str::to_string)),
            webrtc: of("vm.onNetwork").filter(|v| v["kind"] == "webrtc").map(|v| Webrtc { fingerprint: v["address"].clone(), tag: v["tag"].clone() }).collect(),
            host_keys: of("vm.hostKeys").last().map(|v| v["keys"].clone()),
            node_id: of("vm.registerIdentity").find_map(|v| v["computeIdentity"]["nodeId"].as_str().map(str::to_string)),
            signals: of("guest.signalWebRTC").cloned().collect(),
            terminated: of("vm.terminated").next().is_some(),
        }
    }

    pub fn pinned(&self) -> bool {
        self.host_keys.as_ref().and_then(Value::as_array).is_some_and(|k| !k.is_empty())
    }

    /// via iroh needs a ticket and host keys; via webrtc a fingerprint; neither, host keys.
    /// The guest has one DTLS certificate for all its tags and the bidder can drop one tag's
    /// vm.onNetwork, so any tag's fingerprint will do, as atlibs-webrtc's guestOf falls back to.
    pub fn ready(&self, via: Option<Via>) -> bool {
        match via {
            Some(Via::Iroh) => self.pinned() && self.iroh.is_some(),
            Some(Via::Webrtc) => !self.webrtc.is_empty(),
            None => self.pinned(),
        }
    }

    fn answered(&self, tag: &str) -> bool {
        self.signals.iter().any(|v| v["kind"] == "answer" && v["tag"] == tag)
    }
}

// ---- receipt ------------------------------------------------------------------------------------

/// The winner's receipt for an accept {uri, cid}, read at the key the winner gives it.
pub fn receipt(run: Run<Args>) -> Result<Handle, String> {
    let (mut step, on) = Step::<()>::new(run)?;
    match on {
        On::Start => {
            let bidder = did(step.args.arg("bidder"), "bidder")?;
            let accept = strong(step.args.arg("accept"), "accept")?;
            step.out.call(CallId::Receipt, read(&receipt_uri(&bidder, &accept)));
        }
        On::Result(CallId::Receipt, result) => {
            let receipt = result.ok().filter(|o| o["value"].is_object()).map(|o| Receipt::of(&o));
            step.out.finish(json!({ "receipt": receipt }));
        }
        _ => {}
    }
    step.handle()
}

// ---- guest, and shell -----------------------------------------------------------------------------

#[derive(Serialize, Deserialize)]
pub struct Command {
    command: String,
    stdin: Option<String>,
}

/// The guest's reports: the receipt-derived ones, and each tag's guest event streams (answers,
/// age keys) read in rounds until every stream reaches a missing record (at most STREAM_MAX each).
#[derive(Default, Serialize, Deserialize)]
#[serde(default, rename_all = "camelCase")]
pub struct Guest {
    bidder: String,
    receipt: String,
    refs: Vec<AtUri>,
    streams: Vec<Stream>,
    round: Vec<usize>,
    reads: Fan,
    events: Vec<GuestEvent>,
    shell: Option<Command>,
    via: Option<Via>,
}

pub fn guest(run: Run<Args>) -> Result<Handle, String> {
    let (mut step, on) = Step::<Guest>::new(run)?;
    match on {
        On::Start => {
            let bidder = did(step.args.arg("bidder"), "bidder")?;
            let receipt = at_uri(step.args.arg("receipt"), "receipt")?.to_string();
            let tags = tags(step.args.arg("tags"))?;
            read_reports(&mut step, bidder, receipt, &tags);
        }
        On::Result(id, result) => guest_result(&mut step, id, result)?,
        _ => {}
    }
    step.handle()
}

/// webrtc: the host reads the guest, signals and opens the datachannel itself; without a command
/// the channel is only opened (the page's terminal takes it). iroh: the ticket and pinned host
/// keys come from the guest's derived reports (no scan).
pub fn shell(run: Run<Args>) -> Result<Handle, String> {
    let (mut step, on) = Step::<Guest>::new(run)?;
    step.args.signed_in()?;
    match on {
        On::Start => {
            let args = &step.args;
            let receipt = strong(args.arg("receipt"), "receipt")?;
            let command = match args.arg("command") {
                Value::Null => None,
                Value::String(c) if c.is_empty() => None,
                Value::String(c) if c.len() <= 4096 => Some(c.clone()),
                _ => return Err("command must be a string of at most 4096 bytes".into()),
            };
            let stdin = stdin(args.arg("stdin"))?;
            match args.arg("via").as_str().filter(|v| !v.is_empty()).unwrap_or("webrtc") {
                "webrtc" => {
                    let tag = match args.arg("tag") {
                        Value::Null => "ssh".into(),
                        Value::String(t) if t.is_empty() => "ssh".into(),
                        value => tag(value)?,
                    };
                    let stdin = command.as_ref().and(stdin);
                    step.out.call(CallId::Shell, Call::new("", WEBRTC).input(json!(webrtc::Connect { receipt, tag, command, stdin })));
                    step.state.via = Some(Via::Webrtc);
                }
                "iroh" => {
                    let command = command.ok_or("command is required over iroh")?;
                    let bidder = receipt.uri.parse::<AtUri>().map(|u| u.repo).unwrap_or_default();
                    step.state.shell = Some(Command { command, stdin });
                    read_reports(&mut step, bidder, receipt.uri.clone(), &[]);
                }
                _ => return Err("via must be webrtc or iroh".into()),
            }
        }
        On::Result(id, result) => guest_result(&mut step, id, result)?,
        _ => {}
    }
    step.handle()
}

// The derived reports (a missing record is just not reported yet) and the first read of each stream.
fn read_reports(step: &mut Step<Guest>, bidder: String, receipt: String, tags: &[String]) {
    let s = &mut step.state;
    s.refs = report_refs(&bidder, &receipt, tags);
    s.streams = Stream::of(tags);
    let mut reads: Vec<_> = s.refs.iter().enumerate().map(|(k, uri)| (CallId::Report(k), read(uri))).collect();
    s.round = (0..s.streams.len()).collect();
    reads.extend(s.round.iter().map(|&i| (CallId::GuestEvent(i), read(&s.streams[i].uri(&bidder, &receipt)))));
    step.out.fan_out(&mut s.reads, reads);
    (s.bidder, s.receipt) = (bidder, receipt);
}

fn guest_result(step: &mut Step<Guest>, id: CallId, result: Result<Value, String>) -> Result<(), String> {
    let s = &mut step.state;
    match id {
        CallId::Report(_) | CallId::GuestEvent(_) if s.reads.expects(&id) => {
            let Some(results) = s.reads.land(&id, result) else { return Ok(()) };
            let (reports, streamed) = results.split_at(results.len() - s.round.len());
            s.events.extend(s.refs.iter().zip(reports).filter_map(|(uri, result)| GuestEvent::found(&uri.to_string(), result)));
            for (&i, result) in s.round.iter().zip(streamed) {
                let uri = s.streams[i].uri(&s.bidder, &s.receipt).to_string();
                let stream = &mut s.streams[i];
                match GuestEvent::found(&uri, result) {
                    Some(event) => {
                        s.events.push(event);
                        stream.next += 1;
                        stream.done = stream.next >= STREAM_MAX;
                    }
                    None => stream.done = true,
                }
            }
            s.round = (0..s.streams.len()).filter(|&i| !s.streams[i].done).collect();
            if s.round.is_empty() {
                return answer(step);
            }
            let reads: Vec<_> = s.round.iter().map(|&i| (CallId::GuestEvent(i), read(&s.streams[i].uri(&s.bidder, &s.receipt)))).collect();
            step.out.fan_out(&mut s.reads, reads);
        }
        CallId::Shell => step.out.finish(json!({ "via": s.via, "result": result? })),
        _ => {}
    }
    Ok(())
}

fn answer(step: &mut Step<Guest>) -> Result<(), String> {
    let s = &mut step.state;
    let reports = Reports::of(&s.events);
    let Some(shell) = &s.shell else {
        let mut result = json!(reports);
        result["receipt"] = json!(s.receipt);
        result["events"] = json!(s.events);
        step.out.finish(result);
        return Ok(());
    };
    let ticket = reports.iroh.clone().ok_or("the bidder has not reported an iroh ticket (vm.onNetwork iroh://) for this receipt")?;
    if !reports.pinned() {
        return Err("the bidder has not reported vm.hostKeys: refusing an unpinned guest".into());
    }
    let input = iroh::Connect { ticket, alpn: None, host_keys: reports.host_keys, command: Some(shell.command.clone()), stdin: shell.stdin.clone(), send: None };
    step.out.call(CallId::Shell, Call::new("", IROH).input(json!(input)));
    s.via = Some(Via::Iroh);
    Ok(())
}

// ---- watch --------------------------------------------------------------------------------------

/// The receipt-derived reports, re-read on a timer until what `via` (or `until`: terminated,
/// answer) needs is there, or the poll budget runs out. Each poll also reads every guest event
/// stream at its next index, and on through the ones it finds (signalWebRTC answers, publicKey).
#[derive(Default, Serialize, Deserialize)]
#[serde(default, rename_all = "camelCase")]
pub struct Watch {
    bidder: String,
    receipt: String,
    refs: Vec<AtUri>,
    streams: Vec<Stream>,
    via: Option<Via>,
    tag: String,
    poll: Fan,
    fresh: Vec<GuestEvent>,
    events: Vec<GuestEvent>,
    polls: Polls,
}

pub fn watch(run: Run<Args>) -> Result<Handle, String> {
    let (mut step, on) = Step::<Watch>::new(run)?;
    step.args.signed_in()?;
    match on {
        On::Start => {
            let args = &step.args;
            let bidder = did(args.arg("bidder"), "bidder")?;
            let receipt = at_uri(args.arg("receipt"), "receipt")?.to_string();
            let tag: String = args.arg("tag").as_str().unwrap_or("ssh").into();
            let mut tags = tags(args.arg("tags"))?;
            if !tags.contains(&tag) {
                tags.push(tag.clone());
            }
            step.state = Watch {
                refs: report_refs(&bidder, &receipt, &tags),
                streams: Stream::of(&tags),
                via: serde_json::from_value(args.arg("via").clone()).ok(),
                tag,
                bidder,
                receipt,
                ..Watch::default()
            };
            poll(&mut step);
        }
        On::Timer(_) => poll(&mut step),
        On::Result(id, result) => watch_result(&mut step, id, result),
        On::Other => {}
    }
    step.handle()
}

fn poll(step: &mut Step<Watch>) {
    let s = &mut step.state;
    s.fresh.clear();
    let mut reads: Vec<_> = s.refs.iter().enumerate().map(|(k, uri)| (CallId::Report(k), read(uri))).collect();
    reads.extend(s.streams.iter().enumerate().map(|(i, stream)| (CallId::GuestEvent(i), read(&stream.uri(&s.bidder, &s.receipt)))));
    step.out.fan_out(&mut s.poll, reads);
}

fn watch_result(step: &mut Step<Watch>, id: CallId, result: Result<Value, String>) {
    let s = &mut step.state;
    match id {
        CallId::Report(k) => s.fresh.extend(s.refs.get(k).and_then(|uri| GuestEvent::found(&uri.to_string(), &result))),
        CallId::GuestEvent(i) if i < s.streams.len() => {
            let uri = s.streams[i].uri(&s.bidder, &s.receipt).to_string();
            if let Some(event) = GuestEvent::found(&uri, &result) {
                s.events.push(event);
                s.streams[i].next += 1;
                if s.streams[i].next % STREAM_MAX != 0 {
                    step.out.fan_out(&mut s.poll, [(CallId::GuestEvent(i), read(&s.streams[i].uri(&s.bidder, &s.receipt)))]);
                }
            }
        }
        _ => return,
    }
    if s.poll.land(&id, result).is_none() {
        return;
    }
    // A poll is complete: derived reports are re-read each poll; streamed guest events accumulate.
    let all: Vec<GuestEvent> = s.fresh.drain(..).chain(s.events.iter().cloned()).collect();
    let reports = Reports::of(&all);
    let ready = match step.args.arg("until").as_str().unwrap_or("") {
        "terminated" => reports.terminated,
        "answer" => reports.answered(&s.tag),
        _ => reports.ready(s.via),
    };
    if !ready && s.polls.again(&step.args, &mut step.out, CallId::Watch) {
        return;
    }
    let mut answer = json!(reports);
    answer["receipt"] = json!(s.receipt);
    answer["ready"] = json!(ready);
    answer["events"] = json!(all);
    step.out.finish(answer);
}