//! Where a neural decode reads its parameters from (SPEC §3.1, *Decoding against
//! a verified set*).
//!
//! An unpinned decode reads the registered artifact and the files beside it, as a
//! sink always has. A pinned decode reads only the bytes the pin verified
//! ([`crate::pin::Verified`]) — the very bytes that were hashed, never a second read
//! of the same path — and a file the verified set does not contain is refused with
//! `set_incomplete` instead of being read from disk unverified. So does an unpinned
//! decode at a sink that holds a key log: it reads only the bytes check 5 ran over,
//! read once ([`ParamSource::Checked`], [`ParamSource::UnpinnedRead`]).
//!
//! Sessions are built for determinism (`crate::ort_session`): one intra-op and one
//! inter-op thread and a fixed graph-optimization level, because multi-threaded
//! reductions sum partial results in a nondeterministic order. That makes a float
//! decode reproducible on one build; it does not make it reproducible across
//! hardware, which only an integer decode is (SPEC §7).

use std::borrow::Cow;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use ort::session::Session;

use crate::container::PinRefusal;
use crate::pin::Verified;

use super::DecodeError;

/// The parameters one decode reads.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ParamSource {
    /// No pin applied: the registered path, and the files beside it (the path
    /// itself when it is a directory).
    Unpinned(PathBuf),
    /// A pin verified this set; the decode reads nothing else.
    Pinned(Verified),
    /// No pin applied, and the sink holds a key log (SPEC §4, *Module authorisation*):
    /// the set the decode reads, read once when the decode started, over whose every
    /// identity check 5 ran ([`super::ModelRegistry`]) — the registered file of a
    /// one-file set, or every file below the directory the decoder reads from; where
    /// the key log holds the capability's sets to jwp-receipts §7.9, the set §7.9
    /// authorised. `model` is the registered file's name within a set, where a file
    /// was registered. The decode reads nothing else, as a pinned decode does: a file
    /// the set does not contain is `set_incomplete`.
    Checked { set: Verified, model: Option<String> },
    /// No pin applied, and the set an integer capability's decode reads was read into
    /// memory once, when the decode started: its members (`int_payload::set_members`)
    /// by the names a set lists them under, a member the sink could not read being
    /// absent. A key log the registry holds ran check 5 over exactly these bytes, and
    /// the decode reads nothing else ([`super::ModelRegistry`]).
    UnpinnedRead(BTreeMap<String, Vec<u8>>),
}

fn incomplete(name: &str) -> DecodeError {
    DecodeError::PriorMismatch(PinRefusal::SetIncomplete { path: name.to_string() })
}

impl ParamSource {
    /// The directory an unpinned decode reads siblings from.
    pub(crate) fn dir(path: &Path) -> PathBuf {
        if path.is_dir() {
            path.to_path_buf()
        } else {
            path.parent().map(Path::to_path_buf).unwrap_or_default()
        }
    }

    /// Whether a pin was verified for this decode.
    pub fn is_pinned(&self) -> bool {
        matches!(self, Self::Pinned(_))
    }

    /// The artifact the capability registers, named `set_name` within a set: for an
    /// unpinned decode the registered file itself; for one verified file, that file;
    /// for a set read where a file was registered, that file; for another set, the
    /// file it lists as `set_name`.
    pub(crate) fn model(&self, set_name: &str) -> Result<Cow<'_, [u8]>, DecodeError> {
        fn from<'a>(v: &'a Verified, name: &str) -> Result<Cow<'a, [u8]>, DecodeError> {
            match v {
                Verified::File(b) => Ok(Cow::Borrowed(b.as_slice())),
                v => v.file(name).map(Cow::Borrowed).map_err(DecodeError::PriorMismatch),
            }
        }
        match self {
            Self::Unpinned(p) => std::fs::read(p).map(Cow::Owned)
                .map_err(|e| DecodeError::InvalidPayload(format!("{}: {e}", p.display()))),
            Self::UnpinnedRead(files) => files.get(set_name).map(|b| Cow::Borrowed(b.as_slice())).ok_or_else(|| incomplete(set_name)),
            Self::Pinned(v) => from(v, set_name),
            Self::Checked { set, model } => from(set, model.as_deref().unwrap_or(set_name)),
        }
    }

    /// The file `name` beside the registered artifact (unpinned), or in the verified
    /// set. One verified file (a form-less pin, or a one-file set §7.9 authorised)
    /// holds no sibling.
    pub(crate) fn sibling(&self, name: &str) -> Result<Cow<'_, [u8]>, DecodeError> {
        match self {
            Self::Unpinned(p) => {
                let f = Self::dir(p).join(name);
                std::fs::read(&f).map(Cow::Owned)
                    .map_err(|e| DecodeError::InvalidPayload(format!("{}: {e}", f.display())))
            }
            Self::UnpinnedRead(files) => files.get(name).map(|b| Cow::Borrowed(b.as_slice())).ok_or_else(|| incomplete(name)),
            Self::Pinned(v) | Self::Checked { set: v, .. } => v.file(name).map(Cow::Borrowed).map_err(DecodeError::PriorMismatch),
        }
    }

    /// [`Self::sibling`], `None` when the file is absent (an unreadable unpinned
    /// file counts as absent).
    pub(crate) fn try_sibling(&self, name: &str) -> Option<Cow<'_, [u8]>> {
        match self {
            Self::Unpinned(p) => std::fs::read(Self::dir(p).join(name)).ok().map(Cow::Owned),
            Self::UnpinnedRead(files) => files.get(name).map(|b| Cow::Borrowed(b.as_slice())),
            Self::Pinned(v) | Self::Checked { set: v, .. } => v.file(name).ok().map(Cow::Borrowed),
        }
    }

    /// The JSON file `<stem>.json`, parsed.
    pub(crate) fn json<T: for<'de> serde::Deserialize<'de>>(&self, stem: &str) -> Result<T, DecodeError> {
        let name = format!("{stem}.json");
        let raw = self.sibling(&name)?;
        serde_json::from_slice(&raw).map_err(|e| DecodeError::InvalidPayload(format!("{name} parse: {e}")))
    }

    /// A deterministic session over the registered artifact ([`Self::model`]). An
    /// unpinned decode loads it from its path, as before (so a model that keeps
    /// tensors in external files still loads); a pinned one from the verified bytes
    /// alone ([`session_from_verified`]).
    pub(crate) fn session_model(&self, set_name: &str) -> Result<Session, DecodeError> {
        match self {
            Self::Unpinned(p) => crate::ort_session::load_deterministic(p).map_err(DecodeError::Ort),
            Self::Pinned(_) | Self::Checked { .. } | Self::UnpinnedRead(_) => session_from_verified(&self.model(set_name)?),
        }
    }

    /// A deterministic session over the sibling `name` ([`Self::sibling`]).
    pub(crate) fn session_sibling(&self, name: &str) -> Result<Session, DecodeError> {
        match self {
            Self::Unpinned(p) => crate::ort_session::load_deterministic(&Self::dir(p).join(name)).map_err(DecodeError::Ort),
            Self::Pinned(_) | Self::Checked { .. } | Self::UnpinnedRead(_) => session_from_verified(&self.sibling(name)?),
        }
    }
}

/// A deterministic session over verified model bytes. A tensor whose data the model
/// keeps in another file would be read by the runtime from outside the verified
/// set, so such a model is refused with `set_incomplete`, naming that file
/// ([`super::onnx_external`]); bytes that do not parse as a protobuf message are
/// refused as an invalid model.
fn session_from_verified(bytes: &[u8]) -> Result<Session, DecodeError> {
    match super::onnx_external::external_tensor(bytes) {
        Ok(None) => crate::ort_session::load_deterministic_from_memory(bytes).map_err(DecodeError::Ort),
        Ok(Some(location)) => Err(incomplete(&location)),
        Err(e) => Err(DecodeError::Ort(format!("verified model is not an ONNX protobuf: {e}"))),
    }
}

#[cfg(test)]
mod tests {
    use super::super::onnx_external::tests::{external_tensor_proto, inline_tensor, model};
    use super::*;

    /// A pinned decode never builds a session over a model whose tensors live in
    /// another file: it is refused with `set_incomplete` naming that file, before
    /// the runtime sees it. The same graph with its tensor inside loads.
    #[test]
    fn a_pinned_model_with_external_tensors_is_set_incomplete() {
        let ext = ParamSource::Pinned(Verified::File(model(&external_tensor_proto("w", &[1], "weights.bin"))));
        match ext.session_model("model.onnx") {
            Err(DecodeError::PriorMismatch(r)) => {
                assert_eq!(r.code(), "set_incomplete");
                assert_eq!(r, PinRefusal::SetIncomplete { path: "weights.bin".into() });
            }
            Err(e) => panic!("expected set_incomplete, got {e:?}"),
            Ok(_) => panic!("expected set_incomplete, got a session"),
        }
        let inline = ParamSource::Pinned(Verified::File(model(&inline_tensor("w", &[1]))));
        assert!(inline.session_model("model.onnx").is_ok(), "a self-contained model loads");
    }
}
