"""
XML-based snapshot reader for AMUN simulations.

This module implements the AmunXML class, a concrete Amun snapshot reader
for simulation outputs stored in XML metadata files with associated
binary payloads.

Responsibilities include:
- parsing snapshot metadata from XML files
- reading binary dataset payloads
- handling optional compression and byte-level data encoders
- verifying data integrity using embedded digests

The implementation provides the reference behavior against which other
snapshot formats (e.g. AmunH5) are validated.
"""
import os
import numpy as np
import xml.etree.ElementTree as ET

from .amun import Amun
from .compression import decompress
from .data_filters import filter_decode, DataFilterError
from .digests import verify_digest

class AmunXML(Amun):
    """
    Snapshot reader for AMUN XML-format outputs.

    AmunXML reads snapshot metadata from XML files and loads datasets from
    external binary payload files referenced in the XML. It supports:
    - per-payload compression (via `compression_format`)
    - byte-level decoding filters (via `data_encoder`)
    - integrity checks via embedded digests (`digest`, `compressed_digest`)
    """
    digest_policy = "raise"


    def __init_snapshot__(self):
        """
        Initialize snapshot format metadata and validate the XML header.

        - Resolves the metadata filename (defaults to 'metadata.xml' for directory snapshots).
        - Parses the XML root element and confirms it is an AMUNFile.
        - Sets `self.dataformat = 'AmunXML'`.
        - Optionally reads the snapshot version from the XML if not already known.
        """
        if not self.path_is_file:
            self.filename = "metadata.xml"

        mfile = os.path.join(self.dirname, self.filename)
        if os.path.exists(mfile):
            tree = ET.parse(mfile)
            root = tree.getroot()
            if self.version < 0:
                if 'version' in root.attrib:
                    self.version = int(float(root.attrib['version']))
            if root.tag == 'AMUNFile':
                self.dataformat = 'AmunXML'
            else:
                raise Exception("{} does not seem to be an AmunXML snapshot!".format(mfile))
        else:
            raise Exception("{} does not exist!".format(mfile))


    def __fill_attributes__(self):
        """
        Initialize snapshot format metadata and validate the XML header.

        - Resolves the metadata filename (defaults to 'metadata.xml' for directory snapshots).
        - Parses the XML root element and confirms it is an AMUNFile.
        - Sets `self.dataformat = 'AmunXML'`.
        - Optionally reads the snapshot version from the XML if not already known.
        """
        exclude_list = ['nproc']

        tree = ET.parse(os.path.join(self.dirname, self.filename))
        root = tree.getroot()
        for child in root:
            if not child.tag == 'BinaryFiles':
                for item in child:
                    if item.tag == 'Attribute':
                        if not item.attrib['name'] in exclude_list:
                            if item.attrib['type'] == 'double':
                                self.attributes[item.attrib['name']] = float(item.text)
                            elif item.attrib['type'] == 'integer':
                                self.attributes[item.attrib['name']] = int(item.text)
                            else:
                                self.attributes[item.attrib['name']] = item.text

        if not 'nchunks' in self.attributes and 'nprocs' in self.attributes:
            self.attributes['nchunks'] = self.attributes['nprocs']
            del self.attributes['nprocs']

        if not 'zmin' in self.attributes:
            self.attributes['zmin'] = 0
        if not 'zmax' in self.attributes:
            self.attributes['zmax'] = 1
        if not 'zblocks' in self.attributes:
            self.attributes['zblocks'] = 1


    def __fill_variables__(self):
        """
        Populate `self.variables` from the 'variables' attribute.

        The 'variables' attribute is expected to be a whitespace-separated list
        of dataset names available in the snapshot.
        """
        if 'variables' in self.attributes:
            for v in self.attributes['variables'].split():
                self.variables[v] = v
        else:
            raise Exception("No attribute 'variables' in {}!".format(self.filename))


    def __fill_chunks__(self):
        """
        Parse per-chunk metadata and derive per-block geometry information.

        This method:
        1) Reads the <BinaryFiles> section of metadata.xml into `self.binaries`.
        2) For each chunk (0..nchunks-1):
        - parses datablocks_XXXXXX.xml
        - records dataset payload descriptors under `self.chunks[n][dataset_name]`
        - reads the chunk's 'dblocks' count
        3) Reads global metadata payloads:
        - 'fields' to map block IDs -> (level, coordinates)
        - 'bounds' to map block IDs -> physical bounds
        4) For each chunk, loads the chunk-local 'ids' dataset and constructs:
        - `levels`, `coords`, `bounds` arrays for that chunk's blocks

        All binary payload reads go through the shared decode/verify pipeline
        implemented by `__read_binary_payload`.
        """
        self.binaries = dict()

        fname = os.path.join(self.dirname, self.filename)
        tree  = ET.parse(fname)
        root  = tree.getroot()
        for child in root:
            if child.tag == 'BinaryFiles':
                for item in child:
                    self.binaries[item.attrib['name']]         = item.attrib
                    self.binaries[item.attrib['name']]['file'] = item.text

        self.chunkname = 'datablocks_{:06d}.xml'
        for n in range(self.attributes['nchunks']):
            self.chunks[n] = dict()
            self.chunks[n]['filename'] = self.chunkname.format(n)
            fname = os.path.join(self.dirname, self.chunks[n]['filename'])
            if os.path.exists(fname):
                tree = ET.parse(fname)
                root = tree.getroot()
                for item in root.findall('./BinaryFiles/Attribute'):
                    self.chunks[n][item.attrib['name']]         = item.attrib
                    self.chunks[n][item.attrib['name']]['file'] = item.text
                for item in root.findall('./DataBlocks/Attribute'):
                    if item.attrib['name'] == 'dblocks':
                        self.chunks[n]['dblocks'] = int(item.text)
            else:
                raise Exception("Snapshot's chunk '{}' does not exist!".format(fname))

        mset = self.__read_binary_meta('fields')

        nleafs = int(self.attributes['nleafs'])
        if nleafs <= 0:
            raise ValueError(f"Invalid nleafs={nleafs} in snapshot metadata")
        if mset.size % nleafs != 0:
            raise ValueError(
                f"Corrupt 'fields' payload: size={mset.size} not divisible by nleafs={nleafs}"
            )
        n = mset.size // nleafs
        if self.version > 0:
            mset  = np.reshape(mset, [self.attributes['nleafs'],n])
        else:
            mset  = self.__swap__(np.reshape(mset, [n, self.attributes['nleafs']]))

        index = dict()
        level = dict()
        coord = dict()
        for n in range(self.attributes['nleafs']):
            index[mset[n,0]] = n
            level[mset[n,0]] = mset[n,  1]
            coord[mset[n,0]] = mset[n,2:5]

        bounds = self.__read_binary_meta('bounds', dtype='float64')
        if bounds.size != nleafs * 2 * 3:
            raise ValueError(
                f"Corrupt 'bounds' payload: size={bounds.size}, expected {nleafs*2*3}"
            )
        if self.version > 0:
            bounds = np.reshape(bounds, [self.attributes['nleafs'], 2, 3])
        else:
            bounds = self.__swap__(np.reshape(bounds, [2, 3, self.attributes['nleafs']]))

        for n in range(self.attributes['nchunks']):
            nblk = self.chunks[n]['dblocks']
            if nblk > 0:
                ids = self.__read_binary_data('ids', n, dtype='int32')

                self.chunks[n]['levels'] = np.array([level[ids[p]] for p in range(nblk)])
                self.chunks[n]['coords'] = np.array([coord[ids[p]] for p in range(nblk)])

                ii  = [ index[ids[p]] for p in range(nblk) ]

                self.chunks[n]['bounds'] = np.array([bounds[ii[p],:,:] for p in range(nblk)])
            else:
                # Metadata-only chunk (valid when nprocs > number of blocks).
                # Keep base-class invariants: arrays exist with correct zero-length shapes.
                self.chunks[n]['levels'] = np.empty((0,), dtype=np.int32)
                # In this codebase, coords/bounds are stored with 3 columns even for 2D.
                self.chunks[n]['coords']  = np.empty((0, 3), dtype=np.float64)
                self.chunks[n]['bounds']  = np.empty((0, 2, 3), dtype=np.float64)


    def __read_binary_data__(self, dataset_name, chunk_number):
        """
        Return a chunk dataset as a reshaped NumPy array.

        This method:
        - reads the dataset payload bytes for (dataset_name, chunk_number)
        - reshapes the flat array into the chunk dimensions
        - applies version-dependent axis swapping for legacy snapshots

        Notes:
        - The target shape is taken from `self.chunks[chunk_number]['dims']`.
        - For snapshots with version < 1, the dimension order is adjusted
        before reshaping.
        """
        dims = self.chunks[chunk_number]['dims']
        if self.version < 1:
            dims = np.roll(self.chunks[chunk_number]['dims'], -1)
        dset = np.reshape(self.__read_binary_data(dataset_name, chunk_number), dims)
        return self.__swap__(dset)


    def __read_binary_payload(self, fname, meta, dtype):
        """
        Read and decode a binary payload referenced by XML metadata.

        Pipeline:
        - read raw bytes from `fname`
        - verify `compressed_digest` (if present) on the raw bytes
        - decompress using `meta['compression_format']` (if present)
        - apply byte-level decoding filter(s) via `meta['data_encoder']` (if present)
        - verify `digest` (if present) on the decoded bytes
        - return `np.frombuffer(decoded, dtype=dtype)`

        Parameters
        ----------
        fname : str
            Path to the binary payload file.
        meta : dict
            Metadata attributes describing encoding/compression/digests.
        dtype : str or np.dtype
            Target NumPy dtype of the decoded payload.

        Returns
        -------
        np.ndarray
            Flat 1D array view of the decoded payload.
        """
        if os.path.exists(fname):
            with open(fname, mode='rb') as f:
                stream = f.read()

            if 'compressed_digest' in meta:
                verify_digest(
                    data=stream,
                    expected=meta["compressed_digest"],
                    hash_type=meta.get("digest_type", "xxh64"),
                    filename=fname,
                    label="compressed_digest",
                    policy=self.digest_policy,
                )

            data = decompress(stream, meta.get("compression_format"))

            data_encoder = meta.get("data_encoder")
            try:
                data = filter_decode(data, data_encoder=data_encoder, dtype=dtype)
            except DataFilterError as e:
                raise DataFilterError(f"{fname}: {e}") from e

            if "digest" in meta:
                verify_digest(
                    data=data,
                    expected=meta["digest"],
                    hash_type=meta.get("digest_type", "xxh64"),
                    filename=fname,
                    label="digest",
                    policy=self.digest_policy,
                )

            return np.frombuffer(data, dtype=dtype)
        else:
            raise Exception("Binary file '{}' does not exist!".format(fname))


    def __read_binary_meta(self, dataset, dtype='int32'):
        """
        Read a global (non-chunk) metadata payload.

        Parameters
        ----------
        dataset : str
            Name of the global metadata payload (e.g. 'fields', 'bounds').
        dtype : str or np.dtype
            Target dtype used to decode the payload.

        Returns
        -------
        np.ndarray
            Flat 1D array of decoded metadata values.
        """
        meta = self.binaries[dataset]
        fname = os.path.join(self.path, meta['file'])
        return self.__read_binary_payload(fname, meta, dtype=dtype)


    def __read_binary_data(self, dataset_name, chunk_number, dtype='float64'):
        """
        Read a raw (flat) dataset payload from a given chunk.

        This returns a 1D array. Higher-level assembly (reshape + axis swap)
        is handled by `__read_binary_data__()`.

        Parameters
        ----------
        dataset_name : str
            Dataset name as stored in the chunk XML.
        chunk_number : int
            Chunk index.
        dtype : str or np.dtype
            Target dtype used to decode the payload.

        Returns
        -------
        np.ndarray
            Flat 1D array view of the decoded dataset.
        """
        meta = self.chunks[chunk_number][dataset_name]
        fname = os.path.join(self.path, meta['file'])
        return self.__read_binary_payload(fname, meta, dtype=dtype)
