use anyhow::{Context, Result, bail, ensure};
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, SecondsFormat, Utc};
use cid::multibase::{self, Base};
use serde_json::{Map, Value, json};

const STRING: u8 = 0;
const INTERN: u8 = 1;
const KEY: u8 = 2;
const SIGNATURE: u8 = 3;
const CID: u8 = 4;
const ARRAY: u8 = 5;
const OBJECT: u8 = 6;
const BSKY_GENESIS: u8 = 7;
const TRUMP_GENESIS: u8 = 8;
const LITERAL_CBOR: u8 = 9;
const HANDLE: u8 = 10;
const BSKY_CREATE: u8 = 11;

const MAX_PROGRAM: usize = 1024 * 1024;
const MAX_DEPTH: usize = 64;
const MAX_STACK: usize = 16_384;
const MAX_WORK: usize = 16 * 1024 * 1024;

// These are format constants, not a description of Bluesky's current configuration.
pub(crate) const BSKY_ROTATIONS: [&str; 2] = [
    "did:key:zQ3shhCGUqDKjStzuDxPkTxN6ujddP4RkEKJJouJGRRkaLGbg",
    "did:key:zQ3shpKnbdPx3g3CmPf5cRVTPe1HtSwVn5ish3wSnDPQCbLJK",
];
const TRUMP_ALIASES: [&str; 7] = [
    "at://donald.trump.com",
    "at://jr.donald.trump.com",
    "at://eric.trump.com",
    "at://ivanka.trump.com",
    "at://tiffany.trump.com",
    "at://barron.trump.com",
    "at://melania.trump.com",
];

fn varint(mut n: u64, out: &mut Vec<u8>) {
    while n >= 128 {
        out.push(n as u8 | 128);
        n >>= 7;
    }
    out.push(n as u8);
}

fn literal(tag: u8, bytes: &[u8], out: &mut Vec<u8>) {
    out.push(tag);
    varint(bytes.len() as u64, out);
    out.extend_from_slice(bytes);
}

struct Reader<'a>(&'a [u8]);

impl<'a> Reader<'a> {
    fn take(&mut self, n: usize) -> Result<&'a [u8]> {
        ensure!(n <= self.0.len(), "truncated operand");
        let (head, tail) = self.0.split_at(n);
        self.0 = tail;
        Ok(head)
    }

    fn varint(&mut self) -> Result<u64> {
        let mut n = 0;
        for shift in (0..70).step_by(7) {
            let b = self.take(1)?[0];
            ensure!(shift != 63 || b <= 1, "varint overflow");
            n |= u64::from(b & 127) << shift;
            if b < 128 {
                return Ok(n);
            }
        }
        bail!("varint overflow")
    }

    fn bytes(&mut self) -> Result<&'a [u8]> {
        let n = usize::try_from(self.varint()?)?;
        self.take(n)
    }
}

fn genesis(endpoint: Value, aliases: Value, rotations: Value, key: Value, sig: Value) -> Value {
    json!({
        "type": "plc_operation",
        "prev": null,
        "rotationKeys": rotations,
        "alsoKnownAs": aliases,
        "services": {"atproto_pds": {"type": "AtprotoPersonalDataServer", "endpoint": endpoint}},
        "verificationMethods": {"atproto": key},
        "sig": sig,
    })
}

fn legacy_create(service: Value, handle: Value, key: Value, sig: Value) -> Value {
    json!({
        "type": "create",
        "prev": null,
        "recoveryKey": BSKY_ROTATIONS[0],
        "service": service,
        "handle": handle,
        "signingKey": key,
        "sig": sig,
    })
}

pub fn encode(value: &Value, intern: impl Fn(&str) -> Option<u64>) -> Result<Vec<u8>> {
    let depth = value_depth(value);
    ensure!(depth <= MAX_DEPTH, "operation too deep");
    // JSON text does not round-trip every f64 with serde_json's default parser.
    let fallback = serde_ipld_dagcbor::to_vec(value)?;
    ensure!(fallback.len() <= MAX_PROGRAM - 11, "operation too large");
    let mut out = Vec::new();
    // Wide/deep values use a single literal rather than exhausting the VM's work budget.
    if fallback.len() <= MAX_STACK / 2 && fallback.len() * 6 * (depth + 2) <= MAX_WORK {
        encode_value(value, &intern, &mut out, 0)?;
    }
    if out.is_empty() || out.len() > fallback.len() {
        out.clear();
        literal(LITERAL_CBOR, &fallback, &mut out);
    }
    Ok(out)
}

fn encode_value(
    value: &Value,
    intern: &impl Fn(&str) -> Option<u64>,
    out: &mut Vec<u8>,
    depth: usize,
) -> Result<()> {
    ensure!(depth <= MAX_DEPTH, "operation too deep");
    if value.get("type").and_then(Value::as_str) == Some("create") {
        let service = &value["service"];
        let handle = &value["handle"];
        let key = &value["signingKey"];
        let sig = &value["sig"];
        if [service, handle, key, sig].iter().all(|v| v.is_string())
            && value == &legacy_create(service.clone(), handle.clone(), key.clone(), sig.clone())
        {
            for operand in [service, handle, key, sig] {
                encode_value(operand, intern, out, depth + 1)?;
            }
            out.push(BSKY_CREATE);
            return Ok(());
        }
    }
    if value.get("type").and_then(Value::as_str) == Some("plc_operation") {
        let endpoint = &value["services"]["atproto_pds"]["endpoint"];
        let aliases = &value["alsoKnownAs"];
        let rotations = &value["rotationKeys"];
        let key = &value["verificationMethods"]["atproto"];
        let sig = &value["sig"];
        if endpoint.is_string() && key.is_string() && sig.is_string() {
            if rotations == &json!(BSKY_ROTATIONS)
                && aliases
                    .as_array()
                    .is_some_and(|a| a.len() == 1 && a[0].is_string())
                && value
                    == &genesis(
                        endpoint.clone(),
                        aliases.clone(),
                        rotations.clone(),
                        key.clone(),
                        sig.clone(),
                    )
            {
                for operand in [endpoint, &aliases[0], key, sig] {
                    encode_value(operand, intern, out, depth + 1)?;
                }
                out.push(BSKY_GENESIS);
                return Ok(());
            }
            if endpoint == "https://pds.trump.com"
                && aliases == &json!(TRUMP_ALIASES)
                && rotations
                    .as_array()
                    .is_some_and(|a| a.len() == 4 && a.iter().all(Value::is_string))
                && value
                    == &genesis(
                        endpoint.clone(),
                        aliases.clone(),
                        rotations.clone(),
                        key.clone(),
                        sig.clone(),
                    )
            {
                for operand in rotations.as_array().unwrap().iter().chain([key, sig]) {
                    encode_value(operand, intern, out, depth + 1)?;
                }
                out.push(TRUMP_GENESIS);
                return Ok(());
            }
        }
    }
    match value {
        Value::String(s) => {
            if let Some(id) = intern(s) {
                out.push(INTERN);
                varint(id, out);
            } else if let Some(handle) = s
                .strip_prefix("at://")
                .and_then(|s| s.strip_suffix(".bsky.social"))
            {
                literal(HANDLE, handle.as_bytes(), out);
            } else if let Some(key) = s.strip_prefix("did:key:z")
                && key.len() <= 128
                && let Ok(bytes) = Base::Base58Btc.decode(key)
                && Base::Base58Btc.encode(&bytes) == key
            {
                literal(KEY, &bytes, out);
            } else if s.len() == 86
                && let Ok(bytes) = URL_SAFE_NO_PAD.decode(s)
                && bytes.len() == 64
                && URL_SAFE_NO_PAD.encode(&bytes) == *s
            {
                out.push(SIGNATURE);
                out.extend_from_slice(&bytes);
            } else if let Ok(cid) = s.parse::<cid::Cid>()
                && cid.to_string() == *s
            {
                literal(CID, &cid.to_bytes(), out);
            } else {
                literal(STRING, s.as_bytes(), out);
            }
        }
        Value::Array(values) => {
            for value in values {
                encode_value(value, intern, out, depth + 1)?;
            }
            out.push(ARRAY);
            varint(values.len() as u64, out);
        }
        Value::Object(values) => {
            for (key, value) in values {
                encode_value(&Value::String(key.clone()), intern, out, depth + 1)?;
                encode_value(value, intern, out, depth + 1)?;
            }
            out.push(OBJECT);
            varint(values.len() as u64, out);
        }
        _ => literal(LITERAL_CBOR, &serde_ipld_dagcbor::to_vec(value)?, out),
    }
    Ok(())
}

pub fn decode<'a>(program: &[u8], intern: impl Fn(usize) -> Option<&'a str>) -> Result<Value> {
    ensure!(program.len() <= MAX_PROGRAM, "program too large");
    let mut input = Reader(program);
    let mut stack: Vec<(Value, usize)> = Vec::new();
    let mut work = 0usize;
    while !input.0.is_empty() {
        let opcode = input.take(1)?[0];
        let (value, depth) = match opcode {
            STRING | HANDLE => {
                let s = std::str::from_utf8(input.bytes()?)?;
                (
                    Value::String(if opcode == HANDLE {
                        format!("at://{s}.bsky.social")
                    } else {
                        s.to_owned()
                    }),
                    0,
                )
            }
            INTERN => {
                let id = usize::try_from(input.varint()?)?;
                let s = intern(id).context("missing interned string")?;
                ensure!(s.len() <= MAX_PROGRAM, "interned string too large");
                (Value::String(s.to_owned()), 0)
            }
            KEY => {
                let bytes = input.bytes()?;
                ensure!(bytes.len() <= 128, "key operand too large");
                (
                    Value::String(format!(
                        "did:key:{}",
                        multibase::encode(Base::Base58Btc, bytes)
                    )),
                    0,
                )
            }
            SIGNATURE => (Value::String(URL_SAFE_NO_PAD.encode(input.take(64)?)), 0),
            CID => (
                Value::String(cid::Cid::try_from(input.bytes()?)?.to_string()),
                0,
            ),
            LITERAL_CBOR => {
                let value: Value = serde_ipld_dagcbor::from_slice(input.bytes()?)?;
                let depth = value_depth(&value);
                (value, depth)
            }
            ARRAY | OBJECT | BSKY_GENESIS | TRUMP_GENESIS | BSKY_CREATE => {
                let n = match opcode {
                    ARRAY => usize::try_from(input.varint()?)?,
                    OBJECT => usize::try_from(input.varint()?)?
                        .checked_mul(2)
                        .context("object size overflow")?,
                    BSKY_GENESIS | BSKY_CREATE => 4,
                    _ => 6,
                };
                ensure!(n <= stack.len(), "stack underflow");
                let start = stack.len() - n;
                let depth = stack[start..].iter().map(|(_, d)| d + 1).max().unwrap_or(1);
                let mut values = stack.drain(start..).map(|(v, _)| v);
                let value = match opcode {
                    ARRAY => Value::Array(values.collect()),
                    OBJECT => {
                        let mut object = Map::new();
                        while let Some(key) = values.next() {
                            let Value::String(key) = key else {
                                bail!("object key is not a string")
                            };
                            ensure!(
                                object.insert(key, values.next().unwrap()).is_none(),
                                "duplicate object key"
                            );
                        }
                        Value::Object(object)
                    }
                    BSKY_GENESIS | BSKY_CREATE => {
                        let args: Vec<_> = values.collect();
                        ensure!(
                            args.iter().all(Value::is_string),
                            "invalid genesis operands"
                        );
                        let [endpoint, name, key, sig]: [Value; 4] = args.try_into().unwrap();
                        if opcode == BSKY_CREATE {
                            legacy_create(endpoint, name, key, sig)
                        } else {
                            genesis(endpoint, json!([name]), json!(BSKY_ROTATIONS), key, sig)
                        }
                    }
                    _ => {
                        let args: Vec<_> = values.collect();
                        ensure!(
                            args.iter().all(Value::is_string),
                            "invalid genesis operands"
                        );
                        let [a, b, c, d, key, sig]: [Value; 6] = args.try_into().unwrap();
                        genesis(
                            json!("https://pds.trump.com"),
                            json!(TRUMP_ALIASES),
                            json!([a, b, c, d]),
                            key,
                            sig,
                        )
                    }
                };
                (
                    value,
                    if opcode == BSKY_GENESIS || opcode == TRUMP_GENESIS {
                        depth + 2
                    } else {
                        depth
                    },
                )
            }
            _ => bail!("unknown opcode {opcode}"),
        };
        ensure!(depth <= MAX_DEPTH, "operation too deep");
        // Bound cumulative construction work, including dictionary and template expansion.
        work = work
            .checked_add(serde_json::to_vec(&value)?.len())
            .context("expansion overflow")?;
        ensure!(work <= MAX_WORK, "program expansion limit exceeded");
        ensure!(stack.len() < MAX_STACK, "stack limit exceeded");
        stack.push((value, depth));
    }
    ensure!(stack.len() == 1, "program must produce exactly one value");
    Ok(stack.pop().unwrap().0)
}

fn value_depth(value: &Value) -> usize {
    match value {
        Value::Array(a) => 1 + a.iter().map(value_depth).max().unwrap_or(0),
        Value::Object(o) => 1 + o.values().map(value_depth).max().unwrap_or(0),
        _ => 0,
    }
}

pub fn encode_did(did: &str) -> Vec<u8> {
    if let Some(suffix) = did.strip_prefix("did:plc:")
        && let Ok(bytes) = Base::Base32Lower.decode(suffix)
        && bytes.len() == 15
        && Base::Base32Lower.encode(&bytes) == suffix
    {
        let mut out = vec![1];
        out.extend(bytes);
        return out;
    }
    let mut out = vec![0];
    out.extend_from_slice(did.as_bytes());
    out
}

pub fn decode_did(bytes: &[u8]) -> Result<String> {
    match bytes.split_first() {
        Some((0, text)) => Ok(std::str::from_utf8(text)?.to_owned()),
        Some((1, bytes)) if bytes.len() == 15 => {
            Ok(format!("did:plc:{}", Base::Base32Lower.encode(bytes)))
        }
        _ => bail!("invalid DID encoding"),
    }
}

pub fn encode_timestamp(text: &str) -> Vec<u8> {
    if let Ok(time) = DateTime::parse_from_rfc3339(text) {
        let ms = time.timestamp_millis();
        if let Some(time) = DateTime::<Utc>::from_timestamp_millis(ms)
            && time.to_rfc3339_opts(SecondsFormat::Millis, true) == text
        {
            let mut out = vec![1];
            varint(((ms as u64) << 1) ^ ((ms >> 63) as u64), &mut out);
            return out;
        }
    }
    let mut out = vec![0];
    out.extend_from_slice(text.as_bytes());
    out
}

pub fn decode_timestamp(bytes: &[u8]) -> Result<String> {
    match bytes.split_first() {
        Some((0, text)) => Ok(std::str::from_utf8(text)?.to_owned()),
        Some((1, bytes)) => {
            let mut input = Reader(bytes);
            let n = input.varint()?;
            ensure!(input.0.is_empty(), "trailing timestamp bytes");
            let ms = (n >> 1) as i64 ^ -((n & 1) as i64);
            let time =
                DateTime::<Utc>::from_timestamp_millis(ms).context("timestamp out of range")?;
            Ok(time.to_rfc3339_opts(SecondsFormat::Millis, true))
        }
        _ => bail!("invalid timestamp encoding"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn legacy_bluesky_creates_are_compact_and_lossless() -> Result<()> {
        let original = json!({
            "handle": "paul.bsky.social",
            "prev": null,
            "recoveryKey": "did:key:zQ3shhCGUqDKjStzuDxPkTxN6ujddP4RkEKJJouJGRRkaLGbg",
            "service": "https://bsky.social",
            "sig": "DyaPWDItkJnVkN1izINSW-fdjUzP9BkIKlD7SnzD5axfK_870ZZ-1EYcrQLQtP9VkWcp2cdbyIHprjPfeUs8WQ",
            "signingKey": "did:key:zQ3shP5TBe1sQfSttXty15FAEHV1DZgcxRZNxvEWnPfLFwLxJ",
            "type": "create",
        });
        let strings = ["https://bsky.social"];
        let program = encode(&original, |s| {
            strings.iter().position(|v| *v == s).map(|n| n as u64)
        })?;
        assert!(program.len() <= 130, "{} bytes", program.len());
        println!("legacy create: {} bytes", program.len());
        assert_eq!(decode(&program, |id| strings.get(id).copied())?, original);

        let mut variants = vec![original.clone()];
        for field in original.as_object().unwrap().keys() {
            let mut absent = original.clone();
            absent.as_object_mut().unwrap().remove(field);
            variants.push(absent);
            let mut null = original.clone();
            null[field] = Value::Null;
            variants.push(null);
        }
        let mut custom = original.clone();
        custom["recoveryKey"] = original["signingKey"].clone();
        variants.push(custom);
        let mut extra = original;
        extra["future"] = json!({"nested": [false, null, 7]});
        variants.push(extra);
        for value in variants {
            let program = encode(&value, |_| None)?;
            let decoded = decode(&program, |_| None)?;
            assert_eq!(
                serde_ipld_dagcbor::to_vec(&decoded)?,
                serde_ipld_dagcbor::to_vec(&value)?
            );
        }
        Ok(())
    }

    #[test]
    fn values_and_canonical_bytes_survive() -> Result<()> {
        let values = [
            json!({"type":"create", "prev":null, "extra":[true, false, -7, 1.5, {}, []]}),
            json!({"type":"plc_tombstone", "prev":"bafyreih553tggj3y22ewyb5mlhrcszkjjles7fr57xixixf3sk52jeqoyy", "sig":"literal"}),
            json!({"alsoKnownAs":["at://alice.bsky.social", "at://☃.example"], "unknown":{"n":18446744073709551615u64}}),
            json!(-0.0),
        ];
        let strings = ["at://☃.example".to_owned()];
        for value in values {
            let program = encode(&value, |s| {
                strings.iter().position(|v| v == s).map(|n| n as u64)
            })?;
            let decoded = decode(&program, |id| strings.get(id).map(String::as_str))?;
            assert_eq!(decoded, value);
            assert_eq!(
                serde_ipld_dagcbor::to_vec(&decoded)?,
                serde_ipld_dagcbor::to_vec(&value)?
            );
        }
        Ok(())
    }

    #[test]
    fn metadata_is_exact() -> Result<()> {
        for did in [
            "did:plc:ragtjsm2j2vknwkz3zp4oxrd",
            "did:web:example.com",
            "did:plc:RAGT",
            "",
        ] {
            assert_eq!(decode_did(&encode_did(did))?, did);
        }
        for time in [
            "2026-09-12T12:16:35.070Z",
            "1960-01-01T00:00:00.000Z",
            "2026-09-12T12:16:35Z",
            "2026-09-12T12:16:35.070123Z",
            "2026-09-12T13:16:35.070+01:00",
            "unknown",
        ] {
            assert_eq!(decode_timestamp(&encode_timestamp(time))?, time);
        }
        Ok(())
    }

    #[test]
    fn malformed_programs_fail() {
        for bytes in [
            vec![],
            vec![255],
            vec![INTERN, 0],
            vec![STRING, 3, b'a'],
            vec![ARRAY, 1],
            vec![OBJECT, 1],
            vec![BSKY_GENESIS],
            vec![BSKY_CREATE],
            vec![
                LITERAL_CBOR,
                1,
                0xf6,
                STRING,
                0,
                STRING,
                0,
                STRING,
                0,
                BSKY_CREATE,
            ],
            vec![SIGNATURE, 0],
            vec![CID, 1, 0],
            vec![STRING, 1, 255],
            vec![INTERN, 255, 255, 255, 255, 255, 255, 255, 255, 255, 2],
            vec![STRING, 0, STRING, 0],
            vec![STRING, 0, STRING, 0, STRING, 0, STRING, 0, OBJECT, 2],
        ] {
            assert!(decode(&bytes, |_| None).is_err(), "{bytes:?}");
        }
        let mut nested = vec![STRING, 0];
        for _ in 0..100 {
            nested.extend([ARRAY, 1]);
        }
        assert!(decode(&nested, |_| None).is_err());
    }
}