//! The payload of `wai.generate.spectral`: the instruction object, form v0 (SPEC §5,
//! *Payload of `wai.generate.spectral`*). One canonical JSON object names the model by
//! its SHA-256 and gives the prompt, the seed, the steps and every size a sink may
//! make. `wai wrap` checks a `wai.generate.spectral` payload against it before it
//! writes an envelope, `wai inspect` and `wai unwrap` report it, and `wai generate
//! --instr` makes a listed size from one.
//!
//! Reading the object needs no feature: it is JSON and a comparison of bytes. The
//! object is read by its members (a member given twice is refused, since readers differ
//! on which value they keep) and then written again in the canonical form; it is
//! accepted only if that equals the bytes given, so another order, whitespace, an
//! escape or a trailing line feed is refused.
//!
//! Every refusal carries a reason code ([`Refusal::code`]), the first rule the object
//! breaks in the order [`Instruction::parse`] checks them. The conformance vector
//! (`tests/vectors/generate/instr-v0.json`) states the code for each object it lists as
//! refused, so a test notices when a rule stops refusing and another rule happens to
//! refuse the same object. The codes are this tool's; a sink MUST refuse each object,
//! by any reason.

use super::strict_json;
use serde_json::Value;

/// The members of a form-v0 object, in their canonical (ascending byte) order.
pub const MEMBERS: [&str; 5] = ["model", "prompt", "seed", "sizes", "steps"];
/// The longest prompt, in bytes.
pub const MAX_PROMPT: usize = 256;
/// The largest `steps`.
pub const MAX_STEPS: u64 = 64;
/// The most sizes one object lists.
pub const MAX_SIZES: usize = 16;
/// The longest side of a listed size, in pixels.
pub const MAX_SIDE: u64 = 4096;
/// The most pixels at a listed size (`w·h`), 2^23.
pub const MAX_PIXELS: u64 = 1 << 23;

/// Why bytes are not a form-v0 instruction object.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Refusal {
    /// The rule broken: `not_json`, `member_repeated`, `not_object`, `member_unknown`,
    /// `member_missing`, `member_type`, `not_integer`, `model_digest`, `prompt_length`,
    /// `prompt_byte`, `seed_range`, `steps_range`, `sizes_count`, `size_not_pair`,
    /// `size_side`, `size_pixels`, `sizes_order` or `not_canonical`.
    pub code: &'static str,
    /// The same, in words.
    pub detail: String,
}

impl std::fmt::Display for Refusal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}: {}", self.code, self.detail)
    }
}

fn no(code: &'static str, detail: impl Into<String>) -> Refusal {
    Refusal { code, detail: detail.into() }
}

/// A form-v0 instruction object.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Instruction {
    /// The SHA-256 of the model file, as 64 lowercase hexadecimal digits.
    pub model_sha256: String,
    pub prompt: String,
    pub seed: u32,
    pub steps: u32,
    /// `(w, h)`, in strictly ascending order of `w` and then `h`.
    pub sizes: Vec<(u32, u32)>,
}

fn is_hex64(s: &str) -> bool {
    s.len() == 64 && s.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
}

/// Whether `prompt` is a form-v0 prompt: at most [`MAX_PROMPT`] bytes, each a
/// printable ASCII character other than `"` and `\`, so its JSON string needs no
/// escape and its bytes are the same in every JSON writer.
pub fn prompt_ok(prompt: &str) -> Result<(), Refusal> {
    if prompt.len() > MAX_PROMPT {
        return Err(no("prompt_length", format!("the prompt is {} bytes; form v0 allows at most {MAX_PROMPT}", prompt.len())));
    }
    if let Some(b) = prompt.bytes().find(|&b| !(0x20..=0x7e).contains(&b) || b == b'"' || b == b'\\') {
        return Err(no("prompt_byte", format!("the prompt holds byte 0x{b:02x}; form v0 allows 0x20-0x7e other than '\"' and '\\'")));
    }
    Ok(())
}

/// One size within form v0's bounds.
pub fn size_ok(w: u64, h: u64) -> Result<(), Refusal> {
    if !(1..=MAX_SIDE).contains(&w) || !(1..=MAX_SIDE).contains(&h) {
        return Err(no("size_side", format!("size {w}x{h}: each side is 1 to {MAX_SIDE} pixels in form v0")));
    }
    if w * h > MAX_PIXELS {
        return Err(no("size_pixels", format!("size {w}x{h} is {} pixels; form v0 allows at most {MAX_PIXELS}", w * h)));
    }
    Ok(())
}

/// A member's value as an integer written in decimal: `member_type` when it is not a
/// number, `not_integer` when it is one with a fraction or an exponent (which the
/// canonical form would also refuse, but for a reason that says less).
fn integer(v: &Value, name: &str) -> Result<i128, Refusal> {
    match v {
        Value::Number(n) => n.as_u64().map(i128::from).or_else(|| n.as_i64().map(i128::from))
            .ok_or_else(|| no("not_integer", format!("`{name}` is {n}, not an integer in decimal"))),
        _ => Err(no("member_type", format!("`{name}` is not a number"))),
    }
}

impl Instruction {
    /// The canonical serialization: the members in ascending byte order, no
    /// whitespace, integers in decimal. The prompt is written as it is, which
    /// [`prompt_ok`] makes a valid JSON string.
    pub fn canonical(&self) -> Vec<u8> {
        let sizes: Vec<String> = self.sizes.iter().map(|(w, h)| format!("[{w},{h}]")).collect();
        format!(
            "{{\"model\":\"sha256:{}\",\"prompt\":\"{}\",\"seed\":{},\"sizes\":[{}],\"steps\":{}}}",
            self.model_sha256, self.prompt, self.seed, sizes.join(","), self.steps
        )
        .into_bytes()
    }

    /// Every rule of form v0 that is not about the bytes: the model digest, the
    /// prompt, the ranges and the order of the sizes.
    pub fn check(&self) -> Result<(), Refusal> {
        if !is_hex64(&self.model_sha256) {
            return Err(no("model_digest", "`model` is not `sha256:` and 64 lowercase hexadecimal digits"));
        }
        prompt_ok(&self.prompt)?;
        if !(1..=MAX_STEPS).contains(&(self.steps as u64)) {
            return Err(no("steps_range", format!("`steps` is {}; form v0 allows 1 to {MAX_STEPS}", self.steps)));
        }
        if self.sizes.is_empty() || self.sizes.len() > MAX_SIZES {
            return Err(no("sizes_count", format!("`sizes` lists {} sizes; form v0 allows 1 to {MAX_SIZES}", self.sizes.len())));
        }
        for &(w, h) in &self.sizes {
            size_ok(w as u64, h as u64)?;
        }
        if self.sizes.windows(2).any(|p| p[0] >= p[1]) {
            return Err(no("sizes_order", "`sizes` is not in strictly ascending order of w and then h"));
        }
        Ok(())
    }

    /// Read a form-v0 object from the payload bytes. Refuses anything that is not
    /// exactly the canonical serialization of an object that passes [`Self::check`],
    /// with the code of the first rule it breaks, in this order: the JSON (a member
    /// given twice among it), an object, its members, `model`, `prompt`, `seed`,
    /// `steps`, `sizes` (the count, each pair, the order), and last the bytes.
    pub fn parse(bytes: &[u8]) -> Result<Instruction, Refusal> {
        let v = strict_json::from_slice(bytes).map_err(|e| match e {
            strict_json::Refused::Repeated(m) => no("member_repeated", m),
            strict_json::Refused::NotJson(m) => no("not_json", m),
        })?;
        let o = v.as_object().ok_or_else(|| no("not_object", "the payload is JSON but not an object"))?;
        if let Some(k) = o.keys().find(|k| !MEMBERS.contains(&k.as_str())) {
            return Err(no("member_unknown", format!("member `{k}` is not one of form v0's")));
        }
        if let Some(k) = MEMBERS.iter().find(|k| !o.contains_key(**k)) {
            return Err(no("member_missing", format!("member `{k}` is missing")));
        }
        let model = o["model"].as_str().ok_or_else(|| no("member_type", "`model` is not a string"))?;
        let model_sha256 = model.strip_prefix("sha256:").filter(|h| is_hex64(h))
            .ok_or_else(|| no("model_digest", "`model` is not `sha256:` and 64 lowercase hexadecimal digits"))?
            .to_string();
        let prompt = o["prompt"].as_str().ok_or_else(|| no("member_type", "`prompt` is not a string"))?.to_string();
        prompt_ok(&prompt)?;
        let seed = integer(&o["seed"], "seed")?;
        let seed = u32::try_from(seed).map_err(|_| no("seed_range", format!("`seed` is {seed}; form v0 allows 0 to {}", u32::MAX)))?;
        let steps = integer(&o["steps"], "steps")?;
        if !(1..=MAX_STEPS as i128).contains(&steps) {
            return Err(no("steps_range", format!("`steps` is {steps}; form v0 allows 1 to {MAX_STEPS}")));
        }
        let list = o["sizes"].as_array().ok_or_else(|| no("member_type", "`sizes` is not an array"))?;
        if list.is_empty() || list.len() > MAX_SIZES {
            return Err(no("sizes_count", format!("`sizes` lists {} sizes; form v0 allows 1 to {MAX_SIZES}", list.len())));
        }
        let mut sizes = Vec::with_capacity(list.len());
        for s in list {
            let pair = s.as_array().filter(|a| a.len() == 2).ok_or_else(|| no("size_not_pair", "a size is not a pair [w, h]"))?;
            let (w, h) = (integer(&pair[0], "a size's w")?, integer(&pair[1], "a size's h")?);
            if !(1..=MAX_SIDE as i128).contains(&w) || !(1..=MAX_SIDE as i128).contains(&h) {
                return Err(no("size_side", format!("size {w}x{h}: each side is 1 to {MAX_SIDE} pixels in form v0")));
            }
            size_ok(w as u64, h as u64)?;
            sizes.push((w as u32, h as u32));
        }
        let i = Instruction { model_sha256, prompt, seed, steps: steps as u32, sizes };
        i.check()?;
        if i.canonical() != bytes {
            return Err(no("not_canonical", "not in canonical form: members in ascending byte order, no whitespace, \
                                            no escape, nothing after the object"));
        }
        Ok(i)
    }

    /// Whether `(w, h)` is a listed size.
    pub fn lists(&self, w: u32, h: u32) -> bool {
        self.sizes.contains(&(w, h))
    }
}

/// `WxH` → `(w, h)`, within form v0's bounds.
pub fn parse_size(s: &str) -> Result<(u32, u32), String> {
    let (w, h) = s.split_once(['x', 'X']).ok_or_else(|| format!("size `{s}` is not WxH"))?;
    let w: u64 = w.parse().map_err(|_| format!("size `{s}`: width is not an integer"))?;
    let h: u64 = h.parse().map_err(|_| format!("size `{s}`: height is not an integer"))?;
    size_ok(w, h).map_err(|e| e.to_string())?;
    Ok((w as u32, h as u32))
}
