# -*- coding: utf-8 -*-
#
# Copyright (C) 2020-2026 Matthias Klumpp <matthias@tenstral.net>
#
# Licensed under the GNU Lesser General Public License Version 3
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the license, or
# (at your option) any later version.
#
# This software is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this software.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import annotations

import os
import json
import struct
import typing as T
import logging as log
from enum import IntEnum
from uuid import UUID
from pathlib import Path
from datetime import datetime, timezone

import pint
import numpy as np
from xxhash import xxh3_64

from .. import ureg

__all__ = ['TSyncFile', 'TSyncFileMode', 'TSyncTimeUnit']


class TSyncFileMode(IntEnum):
    """Time storage mode of a TSync file."""

    CONTINUOUS = 0  # Continous time-point mapping with no gaps
    SYNCPOINTS = 1  # Only synchronization points are saved


class TSyncTimeUnit(IntEnum):
    """Unit types for time representation in a TSync file."""

    INDEX = 0
    NANOSECONDS = 1
    MICROSECONDS = 2
    MILLISECONDS = 3
    SECONDS = 4


class TSyncDataType(IntEnum):
    """Data types use for storing time values in the data file."""

    INVALID = 0
    INT16 = 2
    INT32 = 3
    INT64 = 4

    UINT16 = 6
    UINT32 = 7
    UINT64 = 8


TSYNC_MAGIC = int('B28FE2434E53548A', 16)
TSYNC_VERSION_MAJOR = 1
TSYNC_VERSION_MINOR = 2

TSYNC_BLOCK_TERM = int('00000000009198E2', 16)


def tsync_dtype_to_pack_fmt_len(dtype: TSyncDataType) -> tuple[str, int]:
    """Convert tsync data type into Python unpack format string and length"""
    if dtype == TSyncDataType.INT16:
        return '<h', 2
    if dtype == TSyncDataType.UINT16:
        return '<H', 2
    if dtype == TSyncDataType.INT32:
        return '<i', 4
    if dtype == TSyncDataType.UINT32:
        return '<I', 4
    if dtype == TSyncDataType.INT64:
        return '<q', 8
    if dtype == TSyncDataType.UINT64:
        return '<Q', 8
    raise RuntimeError('No data defined for how to unpack type {}'.format(dtype))


def tsync_time_unit_to_punit(unit: TSyncTimeUnit) -> pint.Unit:
    if unit == TSyncTimeUnit.INDEX:
        return T.cast(pint.Unit, ureg.dimensionless)
    if unit == TSyncTimeUnit.NANOSECONDS:
        return T.cast(pint.Unit, ureg.nsec)
    if unit == TSyncTimeUnit.MICROSECONDS:
        return T.cast(pint.Unit, ureg.usec)
    if unit == TSyncTimeUnit.MILLISECONDS:
        return T.cast(pint.Unit, ureg.msec)
    if unit == TSyncTimeUnit.SECONDS:
        return T.cast(pint.Unit, ureg.sec)
    raise ValueError('Can not convert tsync time unit type "{}" to Pint unit type.'.format(unit))


class TSyncFile:
    """
    Read a TimeSync (.tsync) binary file as generated by the
    Syntalos DAQ system.
    """

    def __init__(self, fname: os.PathLike[str] | None = None):
        self._format_version = '1.0'
        self._time_created: datetime | None = None
        self._generator_name = ''
        self._collection_id = UUID(int=0x00)
        self._ts_mode = TSyncFileMode.CONTINUOUS
        self._block_size = 512
        self._custom: dict[str, T.Any] = {}
        self._time_labels = ('A', 'B')
        self._time_units = (
            tsync_time_unit_to_punit(TSyncTimeUnit.MICROSECONDS),
            tsync_time_unit_to_punit(TSyncTimeUnit.MICROSECONDS),
        )
        self._times = np.empty((0, 2))
        self._legacy_mode = False
        if fname:
            self.open(fname)

    @property
    def time_created(self) -> datetime | None:
        return self._time_created

    @property
    def tolerance(self) -> pint.Quantity[int]:
        """The tolerance range, as a time duration (usually usec)."""
        return int(self._custom.get('tolerance_us', 0)) * ureg.usec

    @tolerance.setter
    def tolerance(self, value: pint.Quantity[int] | int) -> None:
        # accept a dimensioned quantity, or a bare number interpreted as microseconds
        if isinstance(value, ureg.Quantity):
            value = value.to(ureg.usec).magnitude
        self._custom['tolerance_us'] = int(value)

    @property
    def generator_name(self) -> str:
        """Name of the module that generated this file."""
        return self._generator_name

    @generator_name.setter
    def generator_name(self, name: str) -> None:
        self._generator_name = name

    @property
    def collection_id(self) -> UUID:
        """Data collection ID this file belongs to."""
        return self._collection_id

    @collection_id.setter
    def collection_id(self, uuid: UUID) -> None:
        self._collection_id = uuid

    @property
    def sync_mode(self) -> TSyncFileMode:
        """Time data storage mode.."""
        return self._ts_mode

    @sync_mode.setter
    def sync_mode(self, mode: TSyncFileMode) -> None:
        self._ts_mode = mode

    @property
    def custom(self) -> dict[str, T.Any]:
        """User-defined custom properties of this file."""
        return self._custom

    @custom.setter
    def custom(self, v: dict[str, T.Any]) -> None:
        self._custom = v

    @property
    def time_labels(self) -> tuple[str, str]:
        """Labels of the two encoded times."""
        return self._time_labels

    @time_labels.setter
    def time_labels(self, v: tuple[str, str]) -> None:
        self._time_labels = v

    @property
    def time_units(self) -> tuple[T.Any, T.Any]:
        """Units of the two encoded times."""
        return self._time_units

    @time_units.setter
    def time_units(self, v: tuple[T.Any, T.Any]) -> None:
        self._time_units = v

    @property
    def times(self) -> np.ndarray:
        """The actual time values of the two clocks."""
        return self._times

    @times.setter
    def times(self, v: np.ndarray) -> None:
        self._times = v

    def _read_xxh_unpack(self, format: str, buffer: bytes) -> int:
        self._xxh.update(buffer)
        (v,) = struct.unpack(format, buffer)
        return int(v)

    def _read_utf8_xxh_from_file(self, f: T.BinaryIO) -> str:
        """Read UTF-8 encoded string from binary .tsync file"""

        raw_len = f.read(4)
        self._xxh.update(raw_len)
        (length,) = struct.unpack('<I', raw_len)
        if length == int('ffffffff', 16):
            return ''

        if length > (os.fstat(f.fileno()).st_size - f.tell() + 1):
            raise ValueError('String length in binary too long ({}).'.format(length))

        data = f.read(length)
        self._xxh.update(data)
        return str(data, 'utf-8')

    def open(self, fname: os.PathLike[str]) -> None:
        """Open a tsync file."""

        self._legacy_mode = False
        with open(fname, 'rb') as f:
            (magic_number,) = struct.unpack('<Q', f.read(8))
            if magic_number != TSYNC_MAGIC:
                if magic_number == int('F223434E5953548A', 16):
                    # we have an older file format from pre-3.0 Syntalos versions
                    self._legacy_mode = True
                else:
                    raise ValueError('Unrecognized file type: This file is no tsync file.')

            # read file header block
            self._xxh = xxh3_64()
            major_version = self._read_xxh_unpack('<H', f.read(2))
            minor_version = self._read_xxh_unpack('<H', f.read(2))
            self._format_version = '{}.{}'.format(major_version, minor_version)
            if major_version != TSYNC_VERSION_MAJOR or minor_version > TSYNC_VERSION_MINOR:
                raise ValueError(
                    'Can not read TSync format version {} (max {}.{})'.format(
                        self._format_version, TSYNC_VERSION_MAJOR, TSYNC_VERSION_MINOR
                    )
                )
            log.debug('Reading tsync {} file: {}'.format(self._format_version, fname))

            self._time_created = datetime.fromtimestamp(
                self._read_xxh_unpack('<q', f.read(8)), tz=timezone.utc
            )
            self._generator_name = self._read_utf8_xxh_from_file(f)
            self._collection_id = UUID(self._read_utf8_xxh_from_file(f))
            user_json_raw = self._read_utf8_xxh_from_file(f)
            self._custom = {}
            if user_json_raw:
                self._custom = json.loads(user_json_raw)

            self._ts_mode = TSyncFileMode(self._read_xxh_unpack('<H', f.read(2)))
            self._block_size = self._read_xxh_unpack('<i', f.read(4))

            time1Name = self._read_utf8_xxh_from_file(f)
            time1Unit = TSyncTimeUnit(self._read_xxh_unpack('<H', f.read(2)))
            time1DType = TSyncDataType(self._read_xxh_unpack('<H', f.read(2)))

            time2Name = self._read_utf8_xxh_from_file(f)
            time2Unit = TSyncTimeUnit(self._read_xxh_unpack('<H', f.read(2)))
            time2DType = TSyncDataType(self._read_xxh_unpack('<H', f.read(2)))

            self._time_labels = (time1Name, time2Name)
            self._time_units = (
                tsync_time_unit_to_punit(time1Unit),
                tsync_time_unit_to_punit(time2Unit),
            )

            # skip alignment padding
            padding = (f.tell() * -1) & (8 - 1)
            self._xxh.update(f.read(padding))

            # check header CRC
            term_bytecount = 16
            (block_term,) = struct.unpack('<Q', f.read(8))
            (expected_header_cs,) = struct.unpack('<Q', f.read(8))
            if block_term != TSYNC_BLOCK_TERM:
                raise ValueError(
                    'Header block terminator not found: The file is either '
                    'invalid or its header block was damaged.'
                )
            if expected_header_cs != self._xxh.intdigest():
                raise ValueError(
                    'Header checksum mismatch: The file is either invalid or '
                    'its header block was damaged.'
                )
            self._xxh.reset()

            tfmt1, tlen1 = tsync_dtype_to_pack_fmt_len(time1DType)
            tfmt2, tlen2 = tsync_dtype_to_pack_fmt_len(time2DType)

            self._times = np.empty((0, 2))
            bytes_per_entry = tlen1 + tlen2
            bytes_per_block = bytes_per_entry * self._block_size + term_bytecount

            bytes_remaining = os.fstat(f.fileno()).st_size - f.tell()
            if bytes_remaining <= 0:
                # no data is present
                return

            whole_block_count = bytes_remaining // bytes_per_block
            last_block_bytes_remaining = bytes_remaining - (whole_block_count * bytes_per_block)
            if last_block_bytes_remaining == 0:
                last_block_len = 0
            else:
                last_block_len_f = (last_block_bytes_remaining - term_bytecount) / bytes_per_entry
                if last_block_len_f.is_integer() and last_block_len_f > 0:
                    last_block_len = int(last_block_len_f)
                else:
                    raise ValueError(
                        'File "{}" may be corrupt: Suspicious size ({}) of '
                        'last data block.'.format(fname, last_block_len_f)
                    )
            entries_n = whole_block_count * self._block_size + last_block_len

            self._times = np.zeros((entries_n, 2), dtype=np.int64)
            self._block_crc = 0
            b_index = 0
            i = 0
            while True:
                if bytes_remaining == 0:
                    break

                time1 = self._read_xxh_unpack(tfmt1, f.read(tlen1))
                time2 = self._read_xxh_unpack(tfmt2, f.read(tlen2))
                bytes_remaining -= bytes_per_entry
                self._times[i] = np.array([time1, time2])

                i += 1
                b_index += 1
                if b_index == self._block_size or bytes_remaining == term_bytecount:
                    bytes_remaining -= term_bytecount

                    # check validity of the block we read last
                    (block_term,) = struct.unpack('<Q', f.read(8))
                    (expected_cs,) = struct.unpack('<Q', f.read(8))
                    if block_term != TSYNC_BLOCK_TERM:
                        raise ValueError(
                            'Block terminator not found: Some data is likely corrupted.'
                        )
                    if expected_cs != self._xxh.intdigest():
                        raise ValueError('Block checksum mismatch: Some data is likely corrupted.')
                    self._xxh.reset()
                    b_index = 0
            del self._xxh


class LegacyTSyncFile:
    """
    Read a legacy TimeSync (.tsync) binary file as generated by the
    Syntalos DAQ system.

    This class reads the Syntalos 2.x TSYNC variant, before 3.0 accidentally
    changed the format while porting away from Qt, by complying with the
    written spec and hashing length-prefixes, which was omitted by accident in
    the 2.x series. Great fun for the whole family, this class is the price we
    pay for not catching that issue before it reached production systems.
    """

    def __init__(self, fname: os.PathLike[str] | None = None):
        self._BLOCK_TERM = int('1126000000000000', 16)
        self._format_version = '1.0'
        self._time_created: datetime | None = None
        self._generator_name = ''
        self._collection_id = UUID(int=0x00)
        self._ts_mode = TSyncFileMode.CONTINUOUS
        self._block_size = 128
        self._custom: dict[str, T.Any] = {}
        self._time_labels = ('A', 'B')
        self._time_units = (
            tsync_time_unit_to_punit(TSyncTimeUnit.MICROSECONDS),
            tsync_time_unit_to_punit(TSyncTimeUnit.MICROSECONDS),
        )
        self._times = np.empty((0, 2))
        if fname:
            self.open(fname)

    @property
    def time_created(self) -> datetime | None:
        return self._time_created

    @property
    def tolerance(self) -> pint.Quantity[int]:
        """The tolerance range, as a time duration (usually usec)."""
        return int(self._custom.get('tolerance_us', 0)) * ureg.usec

    @tolerance.setter
    def tolerance(self, value: pint.Quantity[int] | int) -> None:
        # accept a dimensioned quantity, or a bare number interpreted as microseconds
        if isinstance(value, ureg.Quantity):
            value = value.to(ureg.usec).magnitude
        self._custom['tolerance_us'] = int(value)

    @property
    def generator_name(self) -> str:
        """Name of the module that generated this file."""
        return self._generator_name

    @generator_name.setter
    def generator_name(self, name: str) -> None:
        self._generator_name = name

    @property
    def collection_id(self) -> UUID:
        """Data collection ID this file belongs to."""
        return self._collection_id

    @collection_id.setter
    def collection_id(self, uuid: UUID) -> None:
        self._collection_id = uuid

    @property
    def sync_mode(self) -> TSyncFileMode:
        """Time data storage mode.."""
        return self._ts_mode

    @sync_mode.setter
    def sync_mode(self, mode: TSyncFileMode) -> None:
        self._ts_mode = mode

    @property
    def custom(self) -> dict[str, T.Any]:
        """User-defined custom properties of this file."""
        return self._custom

    @custom.setter
    def custom(self, v: dict[str, T.Any]) -> None:
        self._custom = v

    @property
    def time_labels(self) -> tuple[str, str]:
        """Labels of the two encoded times."""
        return self._time_labels

    @time_labels.setter
    def time_labels(self, v: tuple[str, str]) -> None:
        self._time_labels = v

    @property
    def time_units(self) -> tuple[T.Any, T.Any]:
        """Units of the two encoded times."""
        return self._time_units

    @time_units.setter
    def time_units(self, v: tuple[T.Any, T.Any]) -> None:
        self._time_units = v

    @property
    def times(self) -> np.ndarray:
        """The actual time values of the two clocks."""
        return self._times

    @times.setter
    def times(self, v: np.ndarray) -> None:
        self._times = v

    @staticmethod
    def is_legacy(fname: os.PathLike[str]) -> bool:
        with open(fname, 'rb') as f:
            (magic_number,) = struct.unpack('<Q', f.read(8))
            return bool(magic_number == int('F223434E5953548A', 16))

    def _read_xxh_unpack(self, format: str, buffer: bytes) -> int:
        self._xxh.update(buffer)
        if self._xxh_nolen is not None:
            self._xxh_nolen.update(buffer)
        (v,) = struct.unpack(format, buffer)
        return int(v)

    def _read_utf8_xxh_from_file(self, f: T.BinaryIO) -> str:
        """Read UTF-8 encoded string from binary .tsync file.

        While auto-detecting the variant, the length prefix is fed to ``_xxh``
        (spec-compliant 3.0 hashing) but withheld from ``_xxh_nolen``, which
        mirrors the 2.x writer that accidentally omitted it from the hash.
        """

        raw_len = f.read(4)
        self._xxh.update(raw_len)
        (length,) = struct.unpack('<I', raw_len)
        if length == int('ffffffff', 16):
            return ''

        if length > (os.fstat(f.fileno()).st_size - f.tell() + 1):
            raise ValueError('String length in binary too long ({}).'.format(length))

        data = f.read(length)
        self._xxh.update(data)
        if self._xxh_nolen is not None:
            self._xxh_nolen.update(data)
        return str(data, 'utf-8')

    def open(self, fname: os.PathLike[str]) -> None:
        """Open a tsync file."""

        with open(fname, 'rb') as f:
            (magic_number,) = struct.unpack('<Q', f.read(8))
            if magic_number != int('F223434E5953548A', 16):
                raise ValueError('Unrecognized file type: This file is no legacy tsync file.')

            # read file header block. We hash the header twice in parallel:
            # `_xxh` includes the string length prefixes (the spec-compliant
            # variant written by Syntalos 3.0), while `_xxh_nolen` omits them
            # (the 2.x variant). Both share the same magic and format version,
            # so we can only tell them apart by which checksum matches below.
            self._xxh = xxh3_64()
            self._xxh_nolen: T.Any = xxh3_64()
            major_version = self._read_xxh_unpack('<H', f.read(2))
            minor_version = self._read_xxh_unpack('<H', f.read(2))
            self._format_version = '{}.{}'.format(major_version, minor_version)
            if major_version != TSYNC_VERSION_MAJOR or minor_version > TSYNC_VERSION_MINOR:
                raise ValueError(
                    'Can not read TSync format version {} (max {}.{})'.format(
                        self._format_version, TSYNC_VERSION_MAJOR, TSYNC_VERSION_MINOR
                    )
                )
            log.debug('Reading tsync {} file: {}'.format(self._format_version, fname))
            check_xxh = major_version >= 1 and minor_version >= 2
            if not check_xxh:
                log.warning(
                    'Tsync file version ({}) is too old, header checksum validation for '
                    'file integrity will be skipped.'.format(self._format_version)
                )

            self._time_created = datetime.fromtimestamp(
                self._read_xxh_unpack('<q', f.read(8)), tz=timezone.utc
            )
            self._generator_name = self._read_utf8_xxh_from_file(f)
            self._collection_id = UUID(self._read_utf8_xxh_from_file(f))
            user_json_raw = self._read_utf8_xxh_from_file(f)
            self._custom = {}
            if user_json_raw:
                self._custom = json.loads(user_json_raw)

            self._ts_mode = TSyncFileMode(self._read_xxh_unpack('<H', f.read(2)))
            self._block_size = self._read_xxh_unpack('<i', f.read(4))

            time1Name = self._read_utf8_xxh_from_file(f)
            time1Unit = TSyncTimeUnit(self._read_xxh_unpack('<H', f.read(2)))
            time1DType = TSyncDataType(self._read_xxh_unpack('<H', f.read(2)))

            time2Name = self._read_utf8_xxh_from_file(f)
            time2Unit = TSyncTimeUnit(self._read_xxh_unpack('<H', f.read(2)))
            time2DType = TSyncDataType(self._read_xxh_unpack('<H', f.read(2)))

            self._time_labels = (time1Name, time2Name)
            self._time_units = (
                tsync_time_unit_to_punit(time1Unit),
                tsync_time_unit_to_punit(time2Unit),
            )

            # skip alignment padding (hashed by both variants)
            padding = (f.tell() * -1) & (8 - 1)
            padding_data = f.read(padding)
            self._xxh.update(padding_data)
            self._xxh_nolen.update(padding_data)

            # check header CRC
            if check_xxh:
                term_bytecount = 16
                (block_term,) = struct.unpack('<Q', f.read(8))
                (expected_header_cs,) = struct.unpack('<Q', f.read(8))
                if block_term != self._BLOCK_TERM:
                    raise ValueError(
                        'Header block terminator not found: The file is either '
                        'invalid or its header block was damaged.'
                    )
                # The header checksum disambiguates the two legacy variants:
                # whichever hashing matches tells us how this file was written.
                # Data blocks contain no length-prefixed fields, so from here on
                # a single hasher (the surviving one) is sufficient.
                if expected_header_cs == self._xxh.intdigest():
                    self._xxh_nolen = None
                elif expected_header_cs == self._xxh_nolen.intdigest():
                    self._xxh = self._xxh_nolen
                    self._xxh_nolen = None
                else:
                    raise ValueError(
                        'Header checksum mismatch: The file is either invalid or '
                        'its header block was damaged.'
                    )
            else:
                term_bytecount = 8
                self._xxh_nolen = None
                (block_term,) = struct.unpack('<I', f.read(4))
                f.read(4)
                if block_term != int('11260000', 16):
                    # check if we maybe had no padding due to an erroneous writer
                    f.seek((padding + 4 + 4) * -1, os.SEEK_CUR)
                    (block_term,) = struct.unpack('<I', f.read(4))
                    f.read(4)
                    if block_term != int('11260000', 16):
                        raise ValueError(
                            'Header block terminator not found: The file is either '
                            'invalid or its header block was damaged.'
                        )

            self._xxh.reset()

            tfmt1, tlen1 = tsync_dtype_to_pack_fmt_len(time1DType)
            tfmt2, tlen2 = tsync_dtype_to_pack_fmt_len(time2DType)

            self._times = np.empty((0, 2))
            bytes_per_entry = tlen1 + tlen2
            bytes_per_block = bytes_per_entry * self._block_size + term_bytecount

            bytes_remaining = os.fstat(f.fileno()).st_size - f.tell()
            if bytes_remaining <= 0:
                # no data is present
                return

            whole_block_count = bytes_remaining // bytes_per_block
            last_block_bytes_remaining = bytes_remaining - (whole_block_count * bytes_per_block)
            if last_block_bytes_remaining == 0:
                last_block_len = 0
            else:
                last_block_len_f = (last_block_bytes_remaining - term_bytecount) / bytes_per_entry
                if last_block_len_f.is_integer() and last_block_len_f > 0:
                    last_block_len = int(last_block_len_f)
                else:
                    raise ValueError(
                        'File "{}" may be corrupt: Suspicious size ({}) of '
                        'last data block.'.format(fname, last_block_len_f)
                    )
            entries_n = whole_block_count * self._block_size + last_block_len

            self._times = np.zeros((entries_n, 2), dtype=np.int64)
            self._block_crc = 0
            b_index = 0
            i = 0
            while True:
                if bytes_remaining == 0:
                    break

                time1 = self._read_xxh_unpack(tfmt1, f.read(tlen1))
                time2 = self._read_xxh_unpack(tfmt2, f.read(tlen2))
                bytes_remaining -= bytes_per_entry
                self._times[i] = np.array([time1, time2])

                i += 1
                b_index += 1
                if b_index == self._block_size or bytes_remaining == term_bytecount:
                    bytes_remaining -= term_bytecount

                    if not check_xxh:
                        (block_term,) = struct.unpack('<I', f.read(4))
                        f.read(4)
                        if block_term != int('11260000', 16):
                            raise ValueError(
                                'Block terminator not found: Some data may be corrupted.'
                            )
                        b_index = 0
                        continue

                    # check validity of the block we read last
                    (block_term,) = struct.unpack('<Q', f.read(8))
                    (expected_cs,) = struct.unpack('<Q', f.read(8))
                    if block_term != self._BLOCK_TERM:
                        raise ValueError(
                            'Block terminator not found: Some data is likely corrupted.'
                        )
                    if expected_cs != self._xxh.intdigest():
                        raise ValueError('Block checksum mismatch: Some data is likely corrupted.')
                    self._xxh.reset()
                    b_index = 0
            del self._xxh


def load_data(part_paths: T.Iterable[Path], aux_data_list: T.Any) -> T.Iterator[T.Any]:
    """Entry point for automatic dataset loading.

    This function is used internally to Syntalos' .tsync files
    as data or auxiliary data.
    """
    for fname in part_paths:
        tsync: LegacyTSyncFile | TSyncFile
        if LegacyTSyncFile.is_legacy(fname):
            tsync = LegacyTSyncFile(fname)
        else:
            tsync = TSyncFile(fname)
        yield tsync
