from __future__ import annotations
from typing import List, Union, TYPE_CHECKING
from collections.abc import Iterable
from struct import pack
import numpy as np

if TYPE_CHECKING:
    from .mul import Mul, MulImage


class Gwy:
    """Class for handeling the construction of bytes for a .gwy-file

    Args
        imgs: the Mul or MulImage instance that is converted to bytes
             according to the GWY file format

    """

    def __init__(self, imgs: Union[Mul, MulImage]) -> None:
        self.imgs: Union[Mul, List[MulImage]]

        if isinstance(imgs, Iterable):
            self.imgs = imgs
        else:
            self.imgs = [imgs]

        self.gwycontent = self._make_gwy()

    def _make_si_unit(self, dimension: str) -> bytes:
        """Constructs GwySIUnit

        Args:
            dimension: spacial dimension of the SI Unit (either xy or z)

        Returns:
            GwySIUnit bytes

        """

        si_name = (
            b"si_unit_" + bytes(dimension, "utf-8") + b"\0o" + b"GwySIUnit\0"
        )
        unit = b"unitstr\0s" + b"m\0"
        unit_size = pack("<i", len(unit))

        return si_name + unit_size + unit

    def _make_datafield(self, img: MulImage) -> bytes:
        """Constructs a GwyDatafield as defined by the .gwy file format

        Args:
            img: the image for which the datafield bytes are created

        Returns:
            GwyDataField bytes

        """

        datafield = b"".join(
            [
                b"xreal\0d" + pack("<d", float(img.xsize) * 1e-9),
                b"yreal\0d" + pack("<d", float(img.ysize) * 1e-9),
                b"xoff\0d" + pack("<d", 0),
                b"yoff\0d" + pack("<d", 0),
                b"xres\0i" + pack("<i", img.xres),
                b"yres\0i" + pack("<i", img.yres),
            ]
        )

        datafield += self._make_si_unit("xy")
        datafield += self._make_si_unit("z")

        data_arr = np.flip(img.img_data, axis=0).flatten() * 1e-9
        data_arr_size = len(data_arr)

        datafield += b"data\0D" + pack("<i", data_arr_size)
        datafield += data_arr.tobytes()

        datafield_size = pack("<i", len(datafield))

        return b"GwyDataField\0" + datafield_size + datafield

    def _make_meta(self, img: MulImage, channel_num: int) -> bytes:
        """Constructs a  GwyContainer containing the meta-data of an image

        Args:
            img: the image for which the datafield bytes are created
            channel_num: the number of the channel in the GWY file

        Return:
            GwyContainer containing meta-data in bytes

        """

        meta_name = bytes(f"/{channel_num}/meta\0o", "utf-8")
        meta_name += b"GwyContainer\0"

        meta_data = b"".join(
            [
                bytes(f"ID\0s{img.img_id}\0", "utf-8"),
                bytes(f"Bias\0s{img.bias * 1e-3:.5f} V\0", "utf-8"),
                bytes(f"Current\0s{img.current:.2f} nA\0", "utf-8"),
                bytes(f"Current factor\0s{img.currfac}\0", "utf-8"),
                bytes(f"Scan Size\0s{img.xsize} nm, {img.ysize} nm\0", "utf-8"),
                bytes(f"Resolution\0s{img.xres}, {img.yres}\0", "utf-8"),
                bytes(f"Datetime\0s{img.datetime}\0", "utf-8"),
                bytes(f"Gain\0s{img.gain}\0", "utf-8"),
                bytes(f"Mode\0s{img.mode}\0", "utf-8"),
                bytes(f"Postprocessing\0s{img.postpr}\0", "utf-8"),
                bytes(f"Scan Duration\0s{img.speed:.2f} s\0", "utf-8"),
                bytes(f"Line Time\0s{img.line_time:.2f} ms\0", "utf-8"),
                bytes(f"Tilt\0s{img.tilt} deg\0", "utf-8"),
                bytes(f"X-Offset\0s{img.xoffset:.2f} nm\0", "utf-8"),
                bytes(f"Y-Offset\0s{img.yoffset:.2f} nm\0", "utf-8"),
                bytes(f"Z-Scale\0s{img.zscale}\0", "utf-8"),
            ]
        )

        meta_data_size = pack("<i", len(meta_data))

        return meta_name + meta_data_size + meta_data

    def _make_data_container(self) -> bytes:
        """Constructs a top-level GwyContainer containing image data as
        GwyDataFields and image meta-data as nested GwyContainers

        Returns:
            Top-level GwyConainer with DataFields and meta-data in bytes

        """

        container = b""
        for channel_num, img in enumerate(self.imgs):
            channel_title = bytes(f"{img.img_id}\0", "utf-8")
            channel_key = bytes(f"/{channel_num}/data\0o", "utf-8")
            title_key = bytes(f"/{channel_num}/data/title\0s", "utf-8")

            container += (
                channel_key
                + self._make_datafield(img)
                + title_key
                + channel_title
                + self._make_meta(img, channel_num)
            )

        container_size = pack("<i", len(container))

        return b"GwyContainer\0" + container_size + container

    def _make_gwy(self) -> bytes:
        """Constructs the content of a .gwy-file as defined by the GWY file format

        Returns:
            Bytes of the complete GWY file

        """

        GWY_HEADER = b"GWYP"
        content = GWY_HEADER + self._make_data_container()
        return content

    def save_gwyfile(self, output_name: str) -> None:
        """Writes the content of a .gwy file into a file

        Args:
            output_name: name of the .gwy file to create

        """

        with open(output_name, "wb") as f:
            f.write(self.gwycontent)
