//! Keeps mirrors warm from `com.atproto.sync.subscribeRepos`: a commit that
//! touches a knotless collection marks its account's mirrors stale and
//! refreshes them, so the next clone needs no round trip to the PDS.

use crate::App;
use knotless::drisl::{link, record};
use knotless::jetstream::knotless;
pub use knotless::jetstream::{Event, IDLE, Op, Socket, frames};
use knotless::jetstream::{RETRY, RETRY_MAX};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio_tungstenite::tungstenite::Message;

/// Reads a frame. `None` when it is not one.
pub fn parse(frame: &[u8]) -> Option<Event> {
    let mut reader = frame;
    let header: ciborium::Value = ciborium::from_reader(&mut reader).ok()?;
    let body: ciborium::Value = ciborium::from_reader(&mut reader).ok()?;
    let field = |v: &ciborium::Value, name: &str| -> Option<ciborium::Value> {
        v.as_map()?
            .iter()
            .find(|(k, _)| k.as_text() == Some(name))
            .map(|(_, v)| v.clone())
    };
    let text = |v: &ciborium::Value, name: &str| field(v, name)?.as_text().map(str::to_string);
    let kind = text(&header, "t")?;
    let mut event = Event {
        seq: field(&body, "seq")
            .and_then(|s| s.as_integer())
            .and_then(|i| i64::try_from(i).ok()),
        ..Event::default()
    };
    match kind.as_str() {
        "#commit" => {
            let blocks = field(&body, "blocks");
            let blocks = blocks
                .as_ref()
                .and_then(|b| b.as_bytes())
                .map(Vec::as_slice)
                .map(car);
            for op in field(&body, "ops")?.as_array()? {
                let (Some(action), Some(path)) = (text(op, "action"), text(op, "path")) else {
                    continue;
                };
                let Some((collection, rkey)) = path.split_once('/') else {
                    continue;
                };
                if knotless(collection) {
                    let cid = field(op, "cid");
                    let block = cid.as_ref().and_then(link).zip(blocks.as_ref());
                    let record = block.and_then(|(cid, blocks)| record(blocks.get(cid)?));
                    event.ops.push(Op {
                        action,
                        collection: collection.into(),
                        rkey: rkey.into(),
                        record,
                    });
                }
            }
            if !event.ops.is_empty() {
                event.did = Some(text(&body, "repo")?);
                event.rev = text(&body, "rev");
            }
        }
        "#sync" | "#account" | "#identity" => {
            event.did = Some(text(&body, "did")?);
            event.identity = kind == "#identity";
            let active = field(&body, "active").and_then(|a| a.as_bool());
            let status = knotless::jetstream::status(active, text(&body, "status").as_deref());
            event.account = (kind == "#account").then_some(status);
        }
        _ => {}
    }
    Some(event)
}

/// The blocks of a CAR file, by their CIDs' bytes: a commit's carry the
/// records it wrote, which every follower of a shared PDS would otherwise
/// fetch, each of them, for every write.
fn car(bytes: &[u8]) -> HashMap<&[u8], &[u8]> {
    let mut blocks = HashMap::new();
    let mut rest = bytes;
    // The header first, then each block: its length, its CID, its data.
    let mut sections = std::iter::from_fn(|| {
        let len = usize::try_from(varint(&mut rest)?).ok()?;
        let section = rest.get(..len)?;
        rest = &rest[len..];
        Some(section)
    });
    sections.next();
    for section in sections {
        let mut cid = section;
        // A CIDv1: version, codec, then a multihash's code and length.
        let read = (0..4).try_fold(0, |_, _| varint(&mut cid));
        let at = read.and_then(|digest| {
            let digest = usize::try_from(digest).ok()?;
            Some(section.len() - cid.len() + digest).filter(|at| *at <= section.len())
        });
        if let Some(at) = at {
            blocks.insert(&section[..at], &section[at..]);
        }
    }
    blocks
}

/// An unsigned LEB128 integer off the front of `bytes`.
fn varint(bytes: &mut &[u8]) -> Option<u64> {
    let mut value = 0u64;
    for (i, &byte) in bytes.iter().enumerate().take(10) {
        value |= u64::from(byte & 0x7f) << (7 * i);
        if byte & 0x80 == 0 {
            *bytes = &bytes[i + 1..];
            return Some(value);
        }
    }
    None
}

/// The subscribeRepos endpoint for a relay or PDS base URL.
pub fn endpoint(base: &str, cursor: Option<i64>) -> String {
    let mut url = base.trim_end_matches('/').to_string();
    if !url.contains("/xrpc/") {
        url.push_str("/xrpc/com.atproto.sync.subscribeRepos");
    }
    if let Some(c) = cursor {
        url.push_str(&format!("?cursor={c}"));
    }
    url
}

/// Follows the firehose until the process ends, reconnecting from the last
/// sequence number. While disconnected, mirrors fall back to the TTL. A
/// connection that ends is made again after [`RETRY`], and one refused
/// waits twice as long as the last, up to [`RETRY_MAX`], so a relay that is
/// down is not asked twice a second.
pub async fn watch(app: Arc<App>, base: String) {
    let mut cursor = None;
    let mut wait = RETRY;
    app.follows("firehose", &base);
    loop {
        let connected = tokio_tungstenite::connect_async(endpoint(&base, cursor)).await;
        wait = next_wait(wait, connected.is_ok());
        if let Ok((mut ws, _)) = connected {
            app.connected("firehose", true);
            frames(&mut ws, IDLE, |msg| {
                if let Message::Binary(frame) = msg
                    && let Some(event) = parse(&frame)
                {
                    cursor = event.seq.or(cursor);
                    app.heard("firehose", event.seq);
                    app.observe(event);
                }
            })
            .await;
        }
        app.connected("firehose", false);
        tokio::time::sleep(wait).await;
    }
}

/// How long to wait before the next attempt: [`RETRY`] after a connection,
/// and after a refusal twice `wait`, up to [`RETRY_MAX`].
fn next_wait(wait: Duration, connected: bool) -> Duration {
    if connected {
        RETRY
    } else {
        (wait * 2).min(RETRY_MAX)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ciborium::Value as V;

    #[test]
    fn a_refused_relay_is_asked_less_often() {
        let mut wait = RETRY;
        for _ in 0..20 {
            wait = next_wait(wait, false);
        }
        assert_eq!(wait, RETRY_MAX);
        assert_eq!(next_wait(RETRY, false), RETRY * 2);
        assert_eq!(next_wait(RETRY_MAX, true), RETRY);
    }

    fn frame(header: V, body: V) -> Vec<u8> {
        let mut out = Vec::new();
        ciborium::into_writer(&header, &mut out).unwrap();
        ciborium::into_writer(&body, &mut out).unwrap();
        out
    }

    fn map(pairs: Vec<(&str, V)>) -> V {
        V::Map(
            pairs
                .into_iter()
                .map(|(k, v)| (V::Text(k.into()), v))
                .collect(),
        )
    }

    #[test]
    fn frames() {
        let commit = |ops: Vec<V>| {
            frame(
                map(vec![
                    ("op", V::Integer(1.into())),
                    ("t", V::Text("#commit".into())),
                ]),
                map(vec![
                    ("seq", V::Integer(7.into())),
                    ("repo", V::Text("did:plc:a".into())),
                    ("ops", V::Array(ops)),
                ]),
            )
        };
        let op = |action: &str, path: &str| {
            map(vec![
                ("action", V::Text(action.into())),
                ("path", V::Text(path.into())),
            ])
        };
        let event = parse(&commit(vec![
            op("create", "dev.knotless.ref/k.x"),
            op("delete", "dev.knotless.issue/3"),
            op("create", "app.bsky.feed.post/x"),
            op("create", "no-slash"),
            map(vec![("path", V::Text("dev.knotless.ref/y".into()))]),
        ]))
        .unwrap();
        assert_eq!(event.seq, Some(7));
        assert_eq!(event.did.as_deref(), Some("did:plc:a"));
        assert_eq!(
            event.ops,
            [
                Op {
                    action: "create".into(),
                    collection: "dev.knotless.ref".into(),
                    rkey: "k.x".into(),
                    record: None
                },
                Op {
                    action: "delete".into(),
                    collection: "dev.knotless.issue".into(),
                    rkey: "3".into(),
                    record: None
                },
            ]
        );
        let elsewhere = parse(&commit(vec![op("create", "app.bsky.feed.post/x")])).unwrap();
        assert_eq!((elsewhere.did, elsewhere.ops.len()), (None, 0));
        let identity = frame(
            map(vec![("t", V::Text("#identity".into()))]),
            map(vec![("did", V::Text("did:plc:b".into()))]),
        );
        let identity = parse(&identity).unwrap();
        assert_eq!(
            (identity.did.as_deref(), identity.identity),
            (Some("did:plc:b"), true)
        );
        let info = frame(map(vec![("t", V::Text("#info".into()))]), map(vec![]));
        assert_eq!(parse(&info), Some(Event::default()));
        assert_eq!(parse(b"\xff"), None);
        assert_eq!(parse(&frame(map(vec![]), map(vec![]))), None);
    }

    fn hex(s: &str) -> Vec<u8> {
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    /// A record and its CID as the reference implementation writes them,
    /// with a blob in it and bytes.
    const CID: &str = "01711220e257565846105b66bb046048fe489187b80bdd4ad18dc539082dcb8d5f670092";
    const RECORD: &str = "a7616e03626f6bf5636c6f67a463726566d82a58250001551220836ff184e7b41b1e13cb5fd89fa1de98dbbab99e9d2918913ff43b86a5c7c2136473697a650365247479706564626c6f62686d696d65547970656a746578742f706c61696e6372617743010203646e6f6e65f66474616773816161652474797065726465762e6b6e6f746c6573732e636865636b";

    #[test]
    fn records_come_with_their_commits() {
        let section = |bytes: &[u8]| {
            let len = bytes.len();
            let len = if len < 128 {
                vec![len as u8]
            } else {
                vec![len as u8 | 0x80, (len >> 7) as u8]
            };
            [len, bytes.to_vec()].concat()
        };
        let (cid, written) = (hex(CID), hex(RECORD));
        // A block cut short is left out, and a CAR cut short ends there.
        let short = section(&cid[..20]);
        let car = [
            section(b"\xa0"),
            short,
            section(&[cid.clone(), written].concat()),
            vec![60],
        ];
        let link = |cid: &[u8]| V::Tag(42, Box::new(V::Bytes([&[0][..], cid].concat())));
        let commit = |ops: Vec<V>| {
            let body = map(vec![
                ("repo", V::Text("did:plc:a".into())),
                ("blocks", V::Bytes(car.concat())),
                ("ops", V::Array(ops)),
            ]);
            parse(&frame(map(vec![("t", V::Text("#commit".into()))]), body)).unwrap()
        };
        let op = |path: &str, cid: V| {
            map(vec![
                ("action", V::Text("create".into())),
                ("path", V::Text(path.into())),
                ("cid", cid),
            ])
        };
        let event = commit(vec![
            op("dev.knotless.check/a", link(&cid)),
            op("dev.knotless.check/b", link(b"elsewhere")),
            op("dev.knotless.check/c", V::Null),
        ]);
        let records: Vec<_> = event.ops.into_iter().map(|o| o.record).collect();
        assert_eq!(
            records,
            [
                Some(serde_json::json!({
                    "$type": "dev.knotless.check",
                    "n": 3,
                    "ok": true,
                    "none": null,
                    "tags": ["a"],
                    "raw": {"$bytes": "AQID"},
                    "log": {
                        "$type": "blob",
                        "ref": {"$link": "bafkreiedn7yyjz5udmpbhs273cp2dxuy3o5lthu5femjcp7uhodklr6ccm"},
                        "mimeType": "text/plain",
                        "size": 3,
                    },
                })),
                None,
                None,
            ]
        );
        // What no record holds is fetched instead.
        let cbor = |v: V| {
            let mut out = Vec::new();
            ciborium::into_writer(&v, &mut out).unwrap();
            record(&out)
        };
        assert_eq!(cbor(V::Float(1.5)), None);
        assert_eq!(cbor(V::Integer(u64::MAX.into())), None);
        assert_eq!(cbor(V::Map(vec![(V::Integer(1.into()), V::Null)])), None);
        assert_eq!(cbor(V::Tag(42, Box::new(V::Text("x".into())))), None);
        assert_eq!(record(b"\xff"), None);
        assert_eq!(varint(&mut &[0xff; 11][..]), None);
    }

    #[test]
    fn endpoints() {
        assert_eq!(
            endpoint("wss://relay/", None),
            "wss://relay/xrpc/com.atproto.sync.subscribeRepos"
        );
        assert_eq!(
            endpoint("ws://h/xrpc/com.atproto.sync.subscribeRepos", Some(3)),
            "ws://h/xrpc/com.atproto.sync.subscribeRepos?cursor=3"
        );
    }
}