//! The file model the tools share: an open file, its objects, their links
//! and the messages that describe a dataset, all read through
//! `clawhdf5-format` (and the `clawhdf5` facade for dataset values).
//!
//! Nothing here trusts the file: every size is checked before it is used,
//! and every failure is an [`Error`] carrying the address it happened at.

use std::cell::RefCell;
use std::collections::HashMap;
use std::path::{Path, PathBuf};

use clawhdf5::File;
use clawhdf5_format::attribute::{AttributeMessage, extract_attributes_tolerant};
use clawhdf5_format::attribute_info::AttributeInfoMessage;
use clawhdf5_format::btree_v2::{BTreeV2Header, collect_btree_v2_records};
use clawhdf5_format::data_layout::DataLayout;
use clawhdf5_format::dataspace::{Dataspace, DataspaceType};
use clawhdf5_format::datatype::Datatype;
use clawhdf5_format::error::FormatError;
use clawhdf5_format::filter_pipeline::FilterPipeline;
use clawhdf5_format::fractal_heap::FractalHeapHeader;
use clawhdf5_format::group_v1;
use clawhdf5_format::link_info::LinkInfoMessage;
use clawhdf5_format::link_message::{LinkMessage, LinkTarget};
use clawhdf5_format::message_type::MessageType;
use clawhdf5_format::object_header::ObjectHeader;
use clawhdf5_format::superblock::Superblock;
use clawhdf5_format::symbol_table::SymbolTableMessage;

/// Largest dataset or attribute (in bytes) the tools will decode by default.
/// A corrupt dataspace can claim far more elements than the file holds; the
/// limit keeps such a file from exhausting memory. `--max-bytes` changes it.
pub const DEFAULT_MAX_BYTES: u64 = 1 << 30;

/// Upper bound on the objects one traversal visits, so a crafted file with
/// millions of links cannot make a walk run for ever.
pub const MAX_OBJECTS: usize = 1_000_000;

/// A tool error: what went wrong and, when known, the file address (relative
/// to the superblock, as every HDF5 address is) of the structure involved.
#[derive(Debug, Clone)]
pub struct Error {
    pub addr: Option<u64>,
    pub msg: String,
    pub kind: ErrorKind,
}

/// Whether an error says something is wrong with the file.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorKind {
    /// The file is damaged or not valid HDF5 (as far as clawhdf5 knows).
    Corrupt,
    /// Valid HDF5 that clawhdf5 cannot decode (a filter it does not
    /// implement, external raw data files, ...).
    Unsupported,
    /// Refused by a limit of the tool (`--max-bytes`).
    Limit,
}

impl Error {
    pub fn new(msg: impl Into<String>) -> Self {
        Self {
            addr: None,
            msg: msg.into(),
            kind: ErrorKind::Corrupt,
        }
    }

    pub fn at(addr: u64, msg: impl Into<String>) -> Self {
        Self {
            addr: Some(addr),
            msg: msg.into(),
            kind: ErrorKind::Corrupt,
        }
    }

    pub fn with_kind(mut self, kind: ErrorKind) -> Self {
        self.kind = kind;
        self
    }

    /// The same error with `prefix: ` in front of its message.
    pub fn context(mut self, prefix: &str) -> Self {
        self.msg = format!("{prefix}: {}", self.msg);
        self
    }
}

fn format_error_kind(e: &FormatError) -> ErrorKind {
    match e {
        FormatError::UnsupportedFilter(_)
        | FormatError::ExternalDataFilesUnsupported
        | FormatError::ExternalLinkUnsupported { .. } => ErrorKind::Unsupported,
        _ => ErrorKind::Corrupt,
    }
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.addr {
            Some(a) => write!(f, "{} (at address {a:#x})", self.msg),
            None => write!(f, "{}", self.msg),
        }
    }
}

impl From<FormatError> for Error {
    fn from(e: FormatError) -> Self {
        Error::new(e.to_string()).with_kind(format_error_kind(&e))
    }
}

impl From<clawhdf5::Error> for Error {
    fn from(e: clawhdf5::Error) -> Self {
        let kind = match &e {
            clawhdf5::Error::Format(f) => format_error_kind(f),
            _ => ErrorKind::Corrupt,
        };
        Error::new(e.to_string()).with_kind(kind)
    }
}

pub type Result<T> = std::result::Result<T, Error>;

/// Attach an address to a format error.
pub fn fe_at(addr: u64) -> impl Fn(FormatError) -> Error {
    move |e| Error::at(addr, e.to_string())
}

/// What an object header describes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Kind {
    Group,
    Dataset,
    Datatype,
    Unknown,
}

impl Kind {
    pub fn of(h: &ObjectHeader) -> Kind {
        let has = |t: MessageType| h.messages.iter().any(|m| m.msg_type == t);
        if has(MessageType::DataLayout) {
            Kind::Dataset
        } else if has(MessageType::SymbolTable)
            || has(MessageType::LinkInfo)
            || has(MessageType::Link)
            || has(MessageType::GroupInfo)
        {
            Kind::Group
        } else if has(MessageType::Datatype) {
            Kind::Datatype
        } else {
            Kind::Unknown
        }
    }

    pub fn name(self) -> &'static str {
        match self {
            Kind::Group => "Group",
            Kind::Dataset => "Dataset",
            Kind::Datatype => "Type",
            Kind::Unknown => "Unknown",
        }
    }
}

/// Where a link points.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LinkKind {
    Hard(u64),
    Soft(String),
    External {
        file: String,
        path: String,
    },
    /// A user-defined link class (type 65-255): its target means something
    /// only to the application that registered the class.
    UserDefined(u8),
}

#[derive(Debug, Clone)]
pub struct Link {
    pub name: String,
    pub kind: LinkKind,
}

/// An open file.
pub struct H5 {
    pub path: PathBuf,
    pub file: File,
    pub max_bytes: u64,
    /// Fractal heaps whose blocks were verified: `None` = sound.
    verified_heaps: RefCell<HashMap<u64, Option<Error>>>,
}

impl H5 {
    pub fn open(path: &Path) -> Result<H5> {
        if !path.exists() {
            return Err(Error::new(format!("{}: no such file", path.display())));
        }
        let file = File::open(path).map_err(|e| {
            Error::new(format!(
                "{}: not an HDF5 file this tool can open: {e}",
                path.display()
            ))
        })?;
        Ok(H5 {
            path: path.to_path_buf(),
            file,
            max_bytes: DEFAULT_MAX_BYTES,
            verified_heaps: RefCell::new(HashMap::new()),
        })
    }

    /// The file's bytes from the superblock on: what every address indexes.
    pub fn data(&self) -> &[u8] {
        self.file.as_bytes()
    }

    pub fn sb(&self) -> &Superblock {
        self.file.superblock()
    }

    pub fn os(&self) -> u8 {
        self.sb().offset_size
    }

    pub fn ls(&self) -> u8 {
        self.sb().length_size
    }

    pub fn root(&self) -> u64 {
        self.sb().root_group_address
    }

    pub fn header(&self, addr: u64) -> Result<ObjectHeader> {
        let off = usize::try_from(addr).map_err(|_| Error::at(addr, "address out of range"))?;
        ObjectHeader::parse(self.data(), off, self.os(), self.ls())
            .map_err(|e| Error::at(addr, format!("object header: {e}")))
    }

    /// The payload of the first message of type `t`, resolving a shared
    /// message to the message it points at.
    pub fn payload(&self, h: &ObjectHeader, t: MessageType) -> Result<Option<Vec<u8>>> {
        match h.messages.iter().find(|m| m.msg_type == t) {
            None => Ok(None),
            Some(m) => {
                clawhdf5_format::shared_message::message_data(self.data(), m, self.os(), self.ls())
                    .map(|c| Some(c.into_owned()))
                    .map_err(|e| Error::new(format!("{t:?} message: {e}")))
            }
        }
    }

    pub fn datatype(&self, h: &ObjectHeader) -> Result<Datatype> {
        let b = self
            .payload(h, MessageType::Datatype)?
            .ok_or_else(|| Error::new("no datatype message"))?;
        Datatype::parse(&b)
            .map(|(d, _)| d)
            .map_err(|e| Error::new(format!("datatype message: {e}")))
    }

    pub fn dataspace(&self, h: &ObjectHeader) -> Result<Dataspace> {
        let b = self
            .payload(h, MessageType::Dataspace)?
            .ok_or_else(|| Error::new("no dataspace message"))?;
        Dataspace::parse(&b, self.ls()).map_err(|e| Error::new(format!("dataspace message: {e}")))
    }

    pub fn layout(&self, h: &ObjectHeader) -> Result<DataLayout> {
        let m = h
            .messages
            .iter()
            .find(|m| m.msg_type == MessageType::DataLayout)
            .ok_or_else(|| Error::new("no layout message"))?;
        DataLayout::parse(&m.data, self.os(), self.ls())
            .map_err(|e| Error::new(format!("layout message: {e}")))
    }

    pub fn filters(&self, h: &ObjectHeader) -> Result<Option<FilterPipeline>> {
        match self.payload(h, MessageType::FilterPipeline)? {
            None => Ok(None),
            Some(b) => FilterPipeline::parse(&b)
                .map(Some)
                .map_err(|e| Error::new(format!("filter pipeline message: {e}"))),
        }
    }

    /// Every attribute that can be read, plus one error per attribute that
    /// cannot.
    pub fn attributes(&self, h: &ObjectHeader) -> Result<(Vec<AttributeMessage>, Vec<String>)> {
        if let Some(m) = h
            .messages
            .iter()
            .find(|m| m.msg_type == MessageType::AttributeInfo)
            && let Ok(ai) = AttributeInfoMessage::parse(&m.data, self.os())
            && let Some(fh) = ai.fractal_heap_address
        {
            self.verified_heap(fh)
                .map_err(|e| e.context("dense attribute storage"))?;
        }
        let (mut attrs, errs) = extract_attributes_tolerant(self.data(), h, self.os(), self.ls())
            .map_err(|e| Error::new(format!("attributes: {e}")))?;
        attrs.sort_by(|a, b| a.name.cmp(&b.name));
        Ok((attrs, errs.iter().map(|e| e.to_string()).collect()))
    }

    /// Every link of the group whose header is `h`, sorted by name. An
    /// object that is not a group has none.
    pub fn links(&self, h: &ObjectHeader) -> Result<Vec<Link>> {
        let os = self.os();
        let ls = self.ls();
        let data = self.data();
        let mut out = Vec::new();
        if let Some(m) = h
            .messages
            .iter()
            .find(|m| m.msg_type == MessageType::SymbolTable)
        {
            let stm = SymbolTableMessage::parse(&m.data, os)
                .map_err(|e| Error::new(format!("symbol table message: {e}")))?;
            let entries = group_v1::resolve_v1_group_entries(data, &stm, os, ls)
                .map_err(|e| Error::at(stm.btree_address, format!("symbol table: {e}")))?;
            let has_soft = entries.iter().any(group_v1::is_v1_soft_link);
            for e in entries {
                if !group_v1::is_v1_soft_link(&e) {
                    out.push(Link {
                        name: e.name,
                        kind: LinkKind::Hard(e.object_header_address),
                    });
                }
            }
            if has_soft {
                let soft = group_v1::v1_soft_links(data, &stm, os, ls)
                    .map_err(|e| Error::at(stm.btree_address, format!("soft links: {e}")))?;
                for (name, target) in soft {
                    out.push(Link {
                        name,
                        kind: LinkKind::Soft(target),
                    });
                }
            }
        } else {
            let li = match h
                .messages
                .iter()
                .find(|m| m.msg_type == MessageType::LinkInfo)
            {
                Some(m) => Some(
                    LinkInfoMessage::parse(&m.data, os)
                        .map_err(|e| Error::new(format!("link info message: {e}")))?,
                ),
                None => None,
            };
            for m in h
                .messages
                .iter()
                .filter(|m| m.msg_type == MessageType::Link)
            {
                out.push(link_from_message(&m.data, os)?);
            }
            if let Some(li) = li
                && let Some(fh) = li.fractal_heap_address
            {
                for bytes in self.dense_heap_objects(fh, li.btree_name_index_address, 5)? {
                    out.push(link_from_message(&bytes, os)?);
                }
            }
        }
        out.sort_by(|a, b| a.name.cmp(&b.name));
        Ok(out)
    }

    /// The objects a dense-storage B-tree (`btree`, of record type
    /// `name_type`: 5 for links, 8 for attributes) indexes in the fractal
    /// heap at `heap`.
    pub fn dense_heap_objects(
        &self,
        heap: u64,
        btree: Option<u64>,
        name_type: u8,
    ) -> Result<Vec<Vec<u8>>> {
        self.verified_heap(heap)?;
        let data = self.data();
        let os = self.os();
        let ls = self.ls();
        let fh = FractalHeapHeader::parse(data, to_usize(heap)?, os, ls)
            .map_err(|e| Error::at(heap, format!("fractal heap header: {e}")))?;
        let bt = btree.ok_or_else(|| Error::at(heap, "dense storage without a name index"))?;
        let hdr = BTreeV2Header::parse(data, to_usize(bt)?, os, ls)
            .map_err(|e| Error::at(bt, format!("v2 B-tree header: {e}")))?;
        let recs = collect_btree_v2_records(data, &hdr, os, ls)
            .map_err(|e| Error::at(bt, format!("v2 B-tree: {e}")))?;
        // Name-index records: hash(4) + heap ID; creation-order ones: order(8) + heap ID.
        let skip = if hdr.tree_type == name_type { 4 } else { 8 };
        let idlen = usize::from(fh.heap_id_length);
        let mut out = Vec::with_capacity(recs.len());
        for r in &recs {
            let id = r
                .data
                .get(skip..skip + idlen)
                .ok_or_else(|| Error::at(bt, "v2 B-tree record shorter than a heap ID"))?;
            let obj = fh
                .read_managed_object(data, id, os)
                .map_err(|e| Error::at(heap, format!("fractal heap object: {e}")))?;
            out.push(obj);
        }
        Ok(out)
    }

    /// Verify every block of the fractal heap at `addr` (once per heap):
    /// the library reads heap objects without checking the blocks'
    /// checksums, which libhdf5 does, so a damaged block would otherwise be
    /// read as if it were sound.
    pub fn verified_heap(&self, addr: u64) -> Result<()> {
        if let Some(r) = self.verified_heaps.borrow().get(&addr) {
            return r.clone().map_or(Ok(()), Err);
        }
        let report = crate::heap_blocks::verify(self, addr);
        let n = report.problems.len();
        let r = report.problems.into_iter().next().map(|mut e| {
            if n > 1 {
                e.msg = format!("{} (and {} more problems in this heap)", e.msg, n - 1);
            }
            e
        });
        self.verified_heaps.borrow_mut().insert(addr, r.clone());
        r.map_or(Ok(()), Err)
    }

    /// The dataspace of the dataset at `path` with a virtual dataset's
    /// extent resolved from its sources (as libhdf5 reports it) instead of
    /// the stored one.
    pub fn resolved_dataspace(&self, path: &str, h: &ObjectHeader) -> Result<Dataspace> {
        let mut ds = self.dataspace(h)?;
        if matches!(self.layout(h), Ok(DataLayout::Virtual { .. })) {
            let d = self.file.dataset(path)?;
            ds.dimensions = d.shape()?;
            if ds.space_type == DataspaceType::Simple
                && let Some(m) = &ds.max_dimensions
                && m.len() != ds.dimensions.len()
            {
                ds.max_dimensions = None;
            }
        }
        Ok(ds)
    }

    /// Read a dataset's values (raw, in file byte order) through the
    /// `clawhdf5` facade, refusing one larger than `max_bytes`. `ds` must be
    /// the [resolved](Self::resolved_dataspace) dataspace.
    pub fn read_dataset(&self, path: &str, dt: &Datatype, ds: &Dataspace) -> Result<Vec<u8>> {
        let need = byte_len(ds, dt)?;
        if need > self.max_bytes {
            return Err(Error::new(format!(
                "dataset is {need} bytes, over the {} byte limit (--max-bytes)",
                self.max_bytes
            ))
            .with_kind(ErrorKind::Limit));
        }
        let d = self.file.dataset(path)?;
        let raw = d.read_selection(&clawhdf5::Selection::All)?;
        if raw.len() as u64 != need {
            return Err(Error::new(format!(
                "read {} bytes, expected {need}",
                raw.len()
            )));
        }
        Ok(raw)
    }

    /// Walk every object reachable by hard links from the root, depth first
    /// in name order, calling `visit` once per link (the root is visited
    /// first with an empty link name). Each object is described once; later
    /// hard links to it are reported with `first_path` set.
    pub fn walk(&self, mut visit: impl FnMut(&WalkItem<'_>)) -> Result<()> {
        self.walk_from(self.root(), "/", &mut visit)
    }

    pub fn walk_from(
        &self,
        start: u64,
        start_path: &str,
        visit: &mut dyn FnMut(&WalkItem<'_>),
    ) -> Result<()> {
        let mut seen: HashMap<u64, String> = HashMap::new();
        let mut stack: Vec<(u64, String, Option<Link>, usize)> =
            vec![(start, start_path.to_string(), None, 0)];
        let mut count = 0usize;
        while let Some((addr, path, link, depth)) = stack.pop() {
            count += 1;
            if count > MAX_OBJECTS {
                return Err(Error::new(format!(
                    "more than {MAX_OBJECTS} links; stopped walking"
                )));
            }
            // Soft, external and user-defined links have no address.
            if let Some(Link { kind, .. }) = &link
                && !matches!(kind, LinkKind::Hard(_))
            {
                visit(&WalkItem {
                    path: &path,
                    link: link.as_ref(),
                    addr: None,
                    first_path: None,
                    header: None,
                    depth,
                });
                continue;
            }
            if let Some(first) = seen.get(&addr) {
                visit(&WalkItem {
                    path: &path,
                    link: link.as_ref(),
                    addr: Some(addr),
                    first_path: Some(first),
                    header: None,
                    depth,
                });
                continue;
            }
            seen.insert(addr, path.clone());
            let header = self.header(addr);
            visit(&WalkItem {
                path: &path,
                link: link.as_ref(),
                addr: Some(addr),
                first_path: None,
                header: Some(&header),
                depth,
            });
            let Ok(h) = &header else { continue };
            if Kind::of(h) != Kind::Group {
                continue;
            }
            // Link errors are the visitor's to report (it sees the header).
            let Ok(links) = self.links(h) else { continue };
            let base = if path == "/" { "" } else { path.as_str() };
            for l in links.into_iter().rev() {
                let child = format!("{base}/{}", l.name);
                let a = match l.kind {
                    LinkKind::Hard(a) => a,
                    _ => 0,
                };
                stack.push((a, child, Some(l), depth + 1));
            }
        }
        Ok(())
    }

    /// Resolve an absolute object path to its header address, following
    /// soft links like the library does.
    pub fn resolve(&self, path: &str) -> Result<u64> {
        let p = path.trim_matches('/');
        if p.is_empty() {
            return Ok(self.root());
        }
        clawhdf5_format::group_v2::resolve_path_any(self.data(), self.sb(), p)
            .map_err(|e| Error::new(format!("{path}: {e}")))
    }
}

/// An owned copy of a [`WalkItem`], from [`H5::walk_collect`].
pub struct Visited {
    pub path: String,
    pub addr: Option<u64>,
    pub first_path: Option<String>,
    pub header: Option<Result<ObjectHeader>>,
}

impl H5 {
    /// Every step of [`H5::walk`], collected, and the walk's own error.
    pub fn walk_collect(&self) -> (Vec<Visited>, Result<()>) {
        let mut items = Vec::new();
        let r = self.walk(|it| {
            items.push(Visited {
                path: it.path.to_string(),
                addr: it.addr,
                first_path: it.first_path.map(str::to_string),
                header: it.header.cloned(),
            });
        });
        (items, r)
    }
}

/// One step of [`H5::walk`].
pub struct WalkItem<'a> {
    pub path: &'a str,
    /// The link that led here (`None` for the start object).
    pub link: Option<&'a Link>,
    /// Header address (`None` for a soft/external/user-defined link).
    pub addr: Option<u64>,
    /// Set when this object was already visited under another path.
    pub first_path: Option<&'a str>,
    /// The parsed header, for an object seen for the first time.
    pub header: Option<&'a Result<ObjectHeader>>,
    pub depth: usize,
}

fn link_from_message(data: &[u8], os: u8) -> Result<Link> {
    match LinkMessage::parse(data, os) {
        Ok(l) => Ok(Link {
            name: l.name,
            kind: match l.link_target {
                LinkTarget::Hard {
                    object_header_address,
                } => LinkKind::Hard(object_header_address),
                LinkTarget::Soft { target_path } => LinkKind::Soft(target_path),
                LinkTarget::External {
                    filename,
                    object_path,
                } => LinkKind::External {
                    file: filename,
                    path: object_path,
                },
            },
        }),
        Err(FormatError::InvalidLinkType(t)) if t >= 65 => Ok(Link {
            name: user_defined_link_name(data).unwrap_or_else(|| "?".into()),
            kind: LinkKind::UserDefined(t),
        }),
        Err(e) => Err(Error::new(format!("link message: {e}"))),
    }
}

/// The name of a user-defined link, which `LinkMessage::parse` refuses.
fn user_defined_link_name(d: &[u8]) -> Option<String> {
    // version(1) flags(1) [type(1)] [corder(8)] [cset(1)] len(1|2|4|8) name
    let flags = *d.get(1)?;
    let mut p = 2usize;
    if flags & 0x08 != 0 {
        p += 1;
    }
    if flags & 0x04 != 0 {
        p += 8;
    }
    if flags & 0x10 != 0 {
        p += 1;
    }
    let w = 1usize << (flags & 0x03);
    let mut n = 0usize;
    for i in 0..w {
        n |= usize::from(*d.get(p + i)?) << (8 * i);
    }
    p += w;
    let name = d.get(p..p.checked_add(n)?)?;
    Some(String::from_utf8_lossy(name).into_owned())
}

pub fn to_usize(a: u64) -> Result<usize> {
    usize::try_from(a).map_err(|_| Error::at(a, "address out of range"))
}

/// Number of elements a dataspace holds (0 for a null dataspace).
pub fn num_elements(ds: &Dataspace) -> Result<u64> {
    match ds.space_type {
        DataspaceType::Null => Ok(0),
        DataspaceType::Scalar => Ok(1),
        DataspaceType::Simple => ds
            .dimensions
            .iter()
            .try_fold(1u64, |a, &d| a.checked_mul(d))
            .ok_or_else(|| Error::new("dataspace element count overflows")),
    }
}

/// Bytes needed for all of a dataspace's elements of type `dt`.
pub fn byte_len(ds: &Dataspace, dt: &Datatype) -> Result<u64> {
    num_elements(ds)?
        .checked_mul(u64::from(dt.type_size()))
        .ok_or_else(|| Error::new("dataset byte size overflows"))
}

/// Split a `FILE[/object/path]` argument the way h5ls does: the longest
/// prefix that is an existing file is the file.
pub fn split_file_arg(arg: &str) -> (String, Option<String>) {
    if Path::new(arg).is_file() {
        return (arg.to_string(), None);
    }
    let mut idx: Vec<usize> = arg.match_indices('/').map(|(i, _)| i).collect();
    idx.reverse();
    for i in idx {
        let (f, rest) = arg.split_at(i);
        if !f.is_empty() && Path::new(f).is_file() {
            return (f.to_string(), Some(rest.to_string()));
        }
    }
    (arg.to_string(), None)
}