# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2024 Scipp contributors (https://github.com/scipp)

from __future__ import annotations

import dataclasses
import os
from datetime import UTC, datetime
from io import BytesIO
from os import PathLike
from pathlib import Path
from typing import Any, BinaryIO

import numpy as np
import numpy.typing as npt
import scipp as sc

from .._files import open_or_pass
from . import _ir as ir
from ._bytes import Byteorder
from ._low_level_io import LowLevelSqw
from ._models import (
    DataBlockName,
    SqwDataBlockDescriptor,
    SqwDataBlockType,
    SqwDndMetadata,
    SqwFileHeader,
    SqwFileType,
    SqwIXExperiment,
    SqwIXNullInstrument,
    SqwIXSample,
    SqwMainHeader,
    SqwMultiIXExperiment,
    SqwPixelMetadata,
    UniqueObjContainer,
    UniqueRefContainer,
)
from ._read_write import write_object_array

# Based on
# https://github.com/pace-neutrons/Horace/blob/master/documentation/add/05_file_formats.md
_DEFAULT_PIX_ROWS = (
    "u1",  # Coordinate 1 (h)
    "u2",  # Coordinate 2 (k)
    "u3",  # Coordinate 3 (l)
    "u4",  # Coordinate 1 (E)
    "irun",  # Run index in header block
    "idet",  # Detector group number
    "ien",  # Energy bin number
    "signal",  # Signal
    "error",  # Variance
)
_DEFAULT_PIX_ROW_UNITS = (
    '1/angstrom',
    '1/angstrom',
    '1/angstrom',
    'meV',
    None,
    None,
    None,
    'count',
    'count**2',
)


class SqwBuilder:
    def __init__(
        self,
        path: str | PathLike[str] | BinaryIO | BytesIO,
        title: str,
        *,
        byteorder: Byteorder,
    ) -> None:
        self._path = path
        self._stored_path = (
            None if isinstance(self._path, BinaryIO | BytesIO) else Path(self._path)
        )
        self._byteorder = byteorder
        self._n_dims = 0

        main_header = SqwMainHeader(
            full_filename=self._full_filename,
            title=title,
            # To be replaced when registering pixel data.
            nfiles=0,
            # To be replaced when writing the file.
            creation_date=datetime(1, 1, 1, tzinfo=UTC),
        )
        self._data_blocks: dict[DataBlockName, Any] = {  # TODO type
            ("", "main_header"): main_header,
        }

        self._dnd_data: _DndPlaceholder | _DndData | None = None
        self._pix_wrap: _PixWrap | None = None
        self._instrument: SqwIXNullInstrument | None = None
        self._sample: SqwIXSample | None = None

    def create(self, *, chunk_size: int = 8192) -> Path | None:
        with open_or_pass(self._path, "wb") as f:
            sqw_io = LowLevelSqw(
                f,
                path=self._stored_path,
                byteorder=self._byteorder,
            )

            _write_file_header(sqw_io, self._make_file_header())

            block_buffers, block_descriptors = self._serialize_data_blocks()
            bat_buffer, block_descriptors = self._serialize_block_allocation_table(
                block_descriptors=block_descriptors,
                bat_offset=sqw_io.position,
            )
            sqw_io.write_raw(bat_buffer)
            for name, buffer in block_buffers.items():
                descriptor = block_descriptors[name]
                match descriptor.block_type:
                    case SqwDataBlockType.regular:
                        # Type guaranteed by _serialize_data_blocks
                        sqw_io.write_raw(buffer)  # type: ignore[arg-type]
                    case SqwDataBlockType.pix:
                        # Type guaranteed by _serialize_data_blocks
                        self._pix_wrap.write(sqw_io, chunk_size=chunk_size)  # type: ignore[union-attr]
                    case SqwDataBlockType.dnd:
                        # Type guaranteed by _serialize_data_blocks
                        self._dnd_data.write(sqw_io)  # type: ignore[union-attr]
                    case _:
                        raise NotImplementedError(
                            f"Unsupported data block type: {descriptor.block_type}"
                        )

        return self._stored_path

    def add_pixel_data(
        self,
        data: npt.NDArray[np.float32],
        *,
        experiments: list[SqwIXExperiment],
        n_dims: int = 4,
        row_units: tuple[str | None, ...] = _DEFAULT_PIX_ROW_UNITS,
    ) -> SqwBuilder:
        self._n_dims = n_dims
        self._data_blocks[("experiment_info", "expdata")] = SqwMultiIXExperiment(
            experiments
        )

        self._pix_wrap = _PixWrap(row_data=data, row_units=row_units)
        metadata = self._make_pix_metadata(self._pix_wrap)
        self._data_blocks[("pix", "metadata")] = metadata
        self._data_blocks[("", "main_header")].nfiles = len(experiments)
        return self

    def add_empty_detector_params(self) -> SqwBuilder:
        self._data_blocks[("", "detpar")] = UniqueRefContainer(
            global_name="GLOBAL_NAME_DETECTORS_CONTAINER",
            objects=UniqueObjContainer(
                baseclass="IX_detector_array",
                objects=[],
                indices=[],
            ),
        )
        return self

    def _add_dnd_metadata(self, block: SqwDndMetadata) -> SqwBuilder:
        self._data_blocks[("data", "metadata")] = block
        return self

    def add_empty_dnd_data(self, metadata: SqwDndMetadata) -> SqwBuilder:
        # The file must always contain a DND block
        builder = self._add_dnd_metadata(metadata)
        builder._dnd_data = _DndPlaceholder(
            shape=tuple(map(int, metadata.axes.n_bins_all_dims))  # type: ignore[call-overload]
        )
        return builder

    def add_dnd_data(
        self, metadata: SqwDndMetadata, *, data: sc.Variable, counts: sc.Variable
    ) -> SqwBuilder:
        """Add a DND data block (image) to the file.

        Attention
        ---------
        ``data`` and ``counts`` must be given in row-major layout.
        They are converted to column-major (FORTRAN) layout internally.
        This means that the data must have shape ``['u1', 'u2', 'u3', 'u4']``.

        Parameters
        ----------
        metadata:
            Metadata describing the DND data.
        data:
            The values and variances of the image.
        counts:
            The number of observations ('pixels') per bin.

        Returns
        -------
        :
            Self.
        """
        builder = self._add_dnd_metadata(metadata)
        builder._dnd_data = _DndData(
            shape=tuple(map(int, metadata.axes.n_bins_all_dims)),  # type: ignore[call-overload]
            data=data,
            counts=counts,
        )
        return builder

    def add_default_instrument(self, instrument: SqwIXNullInstrument) -> SqwBuilder:
        self._instrument = instrument
        return self

    def add_default_sample(self, sample: SqwIXSample) -> SqwBuilder:
        self._sample = sample
        return self

    def _make_file_header(self) -> SqwFileHeader:
        return SqwFileHeader(
            prog_name="horace",
            prog_version=4.0,
            sqw_type=SqwFileType.SQW,
            n_dims=self._n_dims,
        )

    def _make_pix_metadata(self, pix_wrap: _PixWrap) -> SqwPixelMetadata:
        return SqwPixelMetadata(
            full_filename=self._full_filename,
            npix=pix_wrap.n_pixels(),
            data_range=np.vstack(
                [
                    (
                        pix_wrap.row_data[:, i].min().astype(np.float64),
                        pix_wrap.row_data[:, i].max().astype(np.float64),
                    )
                    for i, unit in enumerate(pix_wrap.row_units)
                ]
            ),
        )

    def _serialize_data_blocks(
        self,
    ) -> tuple[
        dict[DataBlockName, memoryview | None],
        dict[DataBlockName, SqwDataBlockDescriptor],
    ]:
        data_blocks = self._prepare_data_blocks()
        buffers: dict[DataBlockName, memoryview | None] = {}
        descriptors = {}
        for name, data_block in data_blocks.items():
            buffer = BytesIO()
            sqw_io = LowLevelSqw(
                buffer, path=self._stored_path, byteorder=self._byteorder
            )
            write_object_array(sqw_io, data_block.serialize_to_ir().to_object_array())

            buffer.seek(0)
            buf = buffer.getbuffer()
            buffers[name] = buf
            descriptors[name] = SqwDataBlockDescriptor(
                block_type=SqwDataBlockType.regular,
                name=name,
                position=0,
                size=len(buf),
            )

        self._serialize_dnd_data(buffers, descriptors)

        if self._pix_wrap is not None:
            buffers[("pix", "data_wrap")] = None
            descriptors[("pix", "data_wrap")] = SqwDataBlockDescriptor(
                block_type=SqwDataBlockType.pix,
                name=("pix", "data_wrap"),
                position=0,
                size=self._pix_wrap.size(),
            )

        return buffers, descriptors

    def _serialize_dnd_data(
        self,
        buffers: dict[DataBlockName, memoryview | None],
        descriptors: dict[tuple[str, str], SqwDataBlockDescriptor],
    ) -> None:
        match self._dnd_data:
            case _DndPlaceholder() as placeholder:
                buffers[("data", "nd_data")] = None
                descriptors[("data", "nd_data")] = SqwDataBlockDescriptor(
                    block_type=SqwDataBlockType.dnd,
                    name=("data", "nd_data"),
                    position=0,
                    size=placeholder.size(),
                )
            case _DndData() as data:
                buffers[("data", "nd_data")] = None
                descriptors[("data", "nd_data")] = SqwDataBlockDescriptor(
                    block_type=SqwDataBlockType.dnd,
                    name=("data", "nd_data"),
                    position=0,
                    size=data.size(),
                )

    def _prepare_data_blocks(self) -> dict[DataBlockName, Any]:
        filepath, filename = self._filepath_and_name
        blocks = {
            key: block.prepare_for_serialization(filepath=filepath, filename=filename)
            for key, block in self._data_blocks.items()
        }

        nfiles = blocks[("", "main_header")].nfiles
        if self._instrument is not None:
            blocks[("experiment_info", "instruments")] = _broadcast_unique_ref(
                self._instrument,
                n=nfiles,
                baseclass="IX_inst",
                global_name="GLOBAL_NAME_INSTRUMENTS_CONTAINER",
            )
        if self._sample is not None:
            blocks[("experiment_info", "samples")] = _broadcast_unique_ref(
                self._sample,
                n=nfiles,
                baseclass="IX_samp",
                global_name="GLOBAL_NAME_SAMPLES_CONTAINER",
            )

        return _to_canonical_block_order(blocks)

    def _serialize_block_allocation_table(
        self,
        block_descriptors: dict[DataBlockName, SqwDataBlockDescriptor],
        bat_offset: int,
    ) -> tuple[memoryview, dict[DataBlockName, SqwDataBlockDescriptor]]:
        # This function first writes the block allocation table (BAT) with placeholder
        # values in order to determine the size of the BAT.
        # Then, it computes the actual positions that data blocks will have in the file
        # and inserts those positions into the serialized BAT.
        # It returns a buffer of the BAT that can be inserted right after the file
        # header and an updated in-memory representation of the BAT.

        buffer = BytesIO()
        sqw_io = LowLevelSqw(buffer, path=self._stored_path, byteorder=self._byteorder)
        sqw_io.write_u32(0)  # Size of BAT in bytes, filled in below.
        bat_begin = sqw_io.position
        sqw_io.write_u32(len(block_descriptors))
        # Offsets are relative to the local sqw_io.
        position_offsets = {
            name: _write_data_block_descriptor(sqw_io, descriptor)
            for name, descriptor in block_descriptors.items()
        }
        bat_size = sqw_io.position - bat_begin

        block_position = bat_offset + sqw_io.position
        amended_descriptors = {}
        for name, descriptor in block_descriptors.items():
            amended_descriptors[name] = dataclasses.replace(
                descriptor, position=block_position
            )
            offset = position_offsets[name]
            sqw_io.seek(offset)
            sqw_io.write_u64(block_position)
            block_position += descriptor.size

        sqw_io.seek(0)
        sqw_io.write_u32(bat_size)
        return buffer.getbuffer(), amended_descriptors

    @property
    def _full_filename(self) -> str:
        return os.fspath(self._stored_path or "in_memory")

    @property
    def _filepath_and_name(self) -> tuple[str, str]:
        if self._stored_path is None:
            return "", ""
        return os.fspath(self._stored_path.parent), self._stored_path.name


def _write_file_header(sqw_io: LowLevelSqw, file_header: SqwFileHeader) -> None:
    sqw_io.write_char_array(file_header.prog_name)
    sqw_io.write_f64(file_header.prog_version)
    sqw_io.write_u32(file_header.sqw_type.value)
    sqw_io.write_u32(file_header.n_dims)


def _write_data_block_descriptor(
    sqw_io: LowLevelSqw, descriptor: SqwDataBlockDescriptor
) -> int:
    sqw_io.write_char_array(descriptor.block_type.value)
    sqw_io.write_char_array(descriptor.name[0])
    sqw_io.write_char_array(descriptor.name[1])
    pos = sqw_io.position
    sqw_io.write_u64(descriptor.position)
    sqw_io.write_u64(descriptor.size)
    return pos


def _broadcast_unique_ref(
    obj: ir.Serializable,
    n: int,
    baseclass: str,
    global_name: str,
) -> UniqueRefContainer:
    return UniqueRefContainer(
        global_name=global_name,
        objects=UniqueObjContainer(
            baseclass=baseclass,
            objects=[obj],
            indices=[0] * n,
        ),
    )


def _to_canonical_block_order(
    blocks: dict[DataBlockName, Any],
) -> dict[DataBlockName, Any]:
    order = (
        ("", "main_header"),
        ("", "detpar"),
        ("data", "metadata"),
        ("data", "nd_data"),
        ("experiment_info", "instruments"),
        ("experiment_info", "samples"),
        ("experiment_info", "expdata"),
        ("pix", "metadata"),
        ("pix", "data_wrap"),
    )
    blocks = dict(blocks)
    out = {
        name: block for name in order if (block := blocks.pop(name, None)) is not None
    }
    out.update(blocks)  # append remaining blocks if any
    return out


@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
class _DndPlaceholder:
    shape: tuple[int, ...]

    def size(self) -> int:
        n_elem = int(np.prod(self.shape))
        return 4 + 4 * len(self.shape) + 3 * 8 * n_elem

    def write(self, sqw_io: LowLevelSqw) -> None:
        sqw_io.write_u32(len(self.shape))
        for s in self.shape:
            sqw_io.write_u32(s)
        zero = np.zeros(self.shape, dtype="float64")
        sqw_io.write_array(zero)
        sqw_io.write_array(zero)
        sqw_io.write_array(np.zeros(self.shape, dtype="uint64"))


@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
class _DndData:
    data: sc.Variable
    counts: sc.Variable
    shape: tuple[int, ...]

    def size(self) -> int:
        n_elem = int(np.prod(self.shape))
        return 4 + 4 * len(self.shape) + 3 * 8 * n_elem

    def write(self, sqw_io: LowLevelSqw) -> None:
        sqw_io.write_u32(len(self.shape))
        for s in self.shape:
            sqw_io.write_u32(s)
        sqw_io.write_array(self.data.values.astype("float64", copy=False))
        sqw_io.write_array(np.sqrt(self.data.variances.astype("float64", copy=False)))
        # Strictly speaking uint64, but we don't support that in Scipp.
        # So we have to assume that all numbers are positive anyway.
        # Then we can avoid unnecessary type casts.
        sqw_io.write_array(self.counts.values.astype("int64", copy=False))


@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
class _PixWrap:
    row_data: npt.NDArray[np.float32]  # shape: (n_pixels, n_rows)
    row_units: tuple[str | None, ...]

    def size(self) -> int:
        # *4 for f32
        return 4 + 8 + self.n_rows() * 4 * self.n_pixels()

    def n_rows(self) -> int:
        return self.row_data.shape[1]

    def n_pixels(self) -> int:
        return self.row_data.shape[0]

    def write(self, sqw_io: LowLevelSqw, chunk_size: int) -> None:
        sqw_io.write_u32(self.n_rows())
        sqw_io.write_u64(self.n_pixels())
        sqw_io.write_array_fortran_layout(self.row_data)
