# -*- 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/>.

import os
import json
import struct
import typing as T
import logging as log
from enum import IntEnum
from uuid import UUID
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('F223434E5953548A', 16)

TSYNC_VERSION_MAJOR = 1
TSYNC_VERSION_MINOR = 2

TSYNC_BLOCK_TERM = int('1126000000000000', 16)
TSYNC_BLOCK_TERM_32 = int('11260000', 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))


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

    (length,) = struct.unpack('<I', f.read(4))
    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)
    xxh.update(data)
    return str(data, 'utf-8')


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

    def __init__(self, fname: 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 = 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) -> int:
        """The tolerance range value, in microseconds"""
        return int(self._custom.get('tolerance_us', 0))

    @tolerance.setter
    def tolerance(self, usec: int) -> None:
        self._custom['tolerance_us'] = usec

    @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:
        return read_utf8_xxh_from_file(f, self._xxh)

    def open(self, fname: str) -> None:
        with open(fname, 'rb') as f:
            (magic_number,) = struct.unpack('<Q', f.read(8))
            if magic_number != TSYNC_MAGIC:
                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))
            check_xxh = major_version >= 1 and minor_version >= 2
            if not check_xxh:
                log.warning(
                    'Tsync file version ({}) is too old, checksum validation for '
                    'integrity checks 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
            padding = (f.tell() * -1) & (8 - 1)
            self._xxh.update(f.read(padding))

            # 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 != 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.'
                    )
            else:
                term_bytecount = 8
                (block_term,) = struct.unpack('<I', f.read(4))
                f.read(4)
                if block_term != TSYNC_BLOCK_TERM_32:
                    # 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 != TSYNC_BLOCK_TERM_32:
                        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 != TSYNC_BLOCK_TERM_32:
                            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 != 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 (legacy variant for an older, experimental
    version of this file format that was briefly in use).
    """

    def __init__(self, fname: str | None = None):
        self._format_version = 1
        self._time_created: datetime | None = None
        self._tolerance_us = 0
        self._generator_name = ''
        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) -> int:
        """The tolerance range value, in microseconds"""
        return self._tolerance_us

    @tolerance.setter
    def tolerance(self, usec: int) -> None:
        self._tolerance_us = usec

    @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 sync_mode(self) -> TSyncFileMode:
        """Time data storage mode.."""
        return TSyncFileMode.SYNCPOINTS

    @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: str) -> bool:
        with open(fname, 'rb') as f:
            (magic_number,) = struct.unpack('<I', f.read(4))
            return bool(magic_number == int('C6BBDFBC', 16))

    def open(self, fname: str) -> None:
        with open(fname, 'rb') as f:
            (magic_number,) = struct.unpack('<I', f.read(4))
            if magic_number != int('C6BBDFBC', 16):
                raise ValueError('Unrecognized file type.')

            (self._format_version,) = struct.unpack('<I', f.read(4))
            if self._format_version != 1:
                raise ValueError(
                    'Can not read TSync format version {}'.format(self._format_version)
                )
            log.debug('Reading legacy tsync file: {}'.format(fname))

            (ts,) = struct.unpack('<q', f.read(8))
            self._time_created = datetime.fromtimestamp(ts, tz=timezone.utc)
            (self._tolerance_us,) = struct.unpack('<I', f.read(4))

            xxh = xxh3_64()
            try:
                self._generator_name = read_utf8_xxh_from_file(f, xxh)
            except UnicodeDecodeError as uni_e:
                raise ValueError(
                    'This legacy tsync file is damaged and can not be read.'
                ) from uni_e

            json_raw = read_utf8_xxh_from_file(f, xxh)
            self._custom = {}
            if json_raw:
                self._custom = json.loads(json_raw)

            tlabel1 = read_utf8_xxh_from_file(f, xxh)
            tlabel2 = read_utf8_xxh_from_file(f, xxh)
            self._time_labels = (tlabel1, tlabel2)

            (tuv1,) = struct.unpack('<H', f.read(2))
            (tuv2,) = struct.unpack('<H', f.read(2))
            self._time_units = (
                tsync_time_unit_to_punit(TSyncTimeUnit(tuv1)),
                tsync_time_unit_to_punit(TSyncTimeUnit(tuv2)),
            )

            self._times = np.empty((0, 2))
            bytes_per_block = 4 + 8 + 8
            bytes_remaining = os.fstat(f.fileno()).st_size - f.tell()
            if bytes_remaining <= 0:
                # no data is present
                return

            if bytes_remaining % bytes_per_block != 0:
                raise ValueError('File may be corrupt: Not a whole number of data blocks found!')

            num_data_blocks = int(bytes_remaining / bytes_per_block)
            indices_continuous = True
            self._times = np.zeros((num_data_blocks, 2), dtype=np.int64)
            for i in range(num_data_blocks):
                (index,) = struct.unpack('<I', f.read(4))
                (time1,) = struct.unpack('<q', f.read(8))
                (time2,) = struct.unpack('<q', f.read(8))
                if index != i:
                    indices_continuous = False
                self._times[i] = np.array([time1, time2])

            if not indices_continuous:
                print('WARNING: Indices in time sync file were not continuous.')


def load_data(part_paths: T.Iterable[str], 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
