# Copyright 2021-2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
from __future__ import annotations

import dataclasses
from typing import Any
from typing_extensions import Self
from collections.abc import Sequence, Callable
import struct
import enum

from bumble import core
from bumble import device
from bumble import utils
from bumble import att
from bumble import gatt
from bumble import gatt_adapters
from bumble import gatt_client


# -----------------------------------------------------------------------------
class HeartRateService(gatt.TemplateService):
    UUID = gatt.GATT_HEART_RATE_SERVICE

    HEART_RATE_CONTROL_POINT_FORMAT = 'B'
    CONTROL_POINT_NOT_SUPPORTED = 0x80
    RESET_ENERGY_EXPENDED = 0x01

    heart_rate_measurement_characteristic: gatt.Characteristic[HeartRateMeasurement]
    body_sensor_location_characteristic: gatt.Characteristic[BodySensorLocation]
    heart_rate_control_point_characteristic: gatt.Characteristic[int]

    class BodySensorLocation(utils.OpenIntEnum):
        OTHER = 0
        CHEST = 1
        WRIST = 2
        FINGER = 3
        HAND = 4
        EAR_LOBE = 5
        FOOT = 6

    @dataclasses.dataclass
    class HeartRateMeasurement:
        heart_rate: int
        sensor_contact_detected: bool | None = None
        energy_expended: int | None = None
        rr_intervals: Sequence[float] | None = None

        class Flag(enum.IntFlag):
            INT16_HEART_RATE = 1 << 0
            SENSOR_CONTACT_DETECTED = 1 << 1
            SENSOR_CONTACT_SUPPORTED = 1 << 2
            ENERGY_EXPENDED_STATUS = 1 << 3
            RR_INTERVAL = 1 << 4

        def __post_init__(self) -> None:
            if self.heart_rate < 0 or self.heart_rate > 0xFFFF:
                raise core.InvalidArgumentError('heart_rate out of range')

            if self.energy_expended is not None and (
                self.energy_expended < 0 or self.energy_expended > 0xFFFF
            ):
                raise core.InvalidArgumentError('energy_expended out of range')

            if self.rr_intervals:
                for rr_interval in self.rr_intervals:
                    if rr_interval < 0 or rr_interval * 1024 > 0xFFFF:
                        raise core.InvalidArgumentError('rr_intervals out of range')

        @classmethod
        def from_bytes(cls, data: bytes) -> Self:
            flags = data[0]
            offset = 1

            if flags & cls.Flag.INT16_HEART_RATE:
                heart_rate = struct.unpack_from('<H', data, offset)[0]
                offset += 2
            else:
                heart_rate = struct.unpack_from('B', data, offset)[0]
                offset += 1

            if flags & cls.Flag.SENSOR_CONTACT_SUPPORTED:
                sensor_contact_detected = flags & cls.Flag.SENSOR_CONTACT_DETECTED != 0
            else:
                sensor_contact_detected = None

            if flags & cls.Flag.ENERGY_EXPENDED_STATUS:
                energy_expended = struct.unpack_from('<H', data, offset)[0]
                offset += 2
            else:
                energy_expended = None

            rr_intervals: Sequence[float] | None = None
            if flags & cls.Flag.RR_INTERVAL:
                rr_intervals = tuple(
                    struct.unpack_from('<H', data, i)[0] / 1024
                    for i in range(offset, len(data), 2)
                )

            return cls(
                heart_rate=heart_rate,
                sensor_contact_detected=sensor_contact_detected,
                energy_expended=energy_expended,
                rr_intervals=rr_intervals,
            )

        def __bytes__(self) -> bytes:
            flags = 0
            if self.heart_rate < 256:
                data = struct.pack('B', self.heart_rate)
            else:
                flags |= self.Flag.INT16_HEART_RATE
                data = struct.pack('<H', self.heart_rate)

            if self.sensor_contact_detected is not None:
                flags |= self.Flag.SENSOR_CONTACT_SUPPORTED
                if self.sensor_contact_detected:
                    flags |= self.Flag.SENSOR_CONTACT_DETECTED

            if self.energy_expended is not None:
                flags |= self.Flag.ENERGY_EXPENDED_STATUS
                data += struct.pack('<H', self.energy_expended)

            if self.rr_intervals is not None:
                flags |= self.Flag.RR_INTERVAL
                data += b''.join(
                    [
                        struct.pack('<H', int(rr_interval * 1024))
                        for rr_interval in self.rr_intervals
                    ]
                )

            return bytes([flags]) + data

    def __init__(
        self,
        read_heart_rate_measurement: Callable[
            [device.Connection], HeartRateMeasurement
        ],
        body_sensor_location: HeartRateService.BodySensorLocation | None = None,
        reset_energy_expended: Callable[[device.Connection], Any] | None = None,
    ):
        self.heart_rate_measurement_characteristic = (
            gatt_adapters.SerializableCharacteristicAdapter(
                gatt.Characteristic(
                    uuid=gatt.GATT_HEART_RATE_MEASUREMENT_CHARACTERISTIC,
                    properties=gatt.Characteristic.Properties.NOTIFY,
                    permissions=gatt.Characteristic.Permissions(0),
                    value=gatt.CharacteristicValue(read=read_heart_rate_measurement),
                ),
                HeartRateService.HeartRateMeasurement,
            )
        )
        characteristics: list[gatt.Characteristic] = [
            self.heart_rate_measurement_characteristic
        ]

        if body_sensor_location is not None:
            self.body_sensor_location_characteristic = (
                gatt_adapters.EnumCharacteristicAdapter(
                    gatt.Characteristic(
                        uuid=gatt.GATT_BODY_SENSOR_LOCATION_CHARACTERISTIC,
                        properties=gatt.Characteristic.Properties.READ,
                        permissions=gatt.Characteristic.READABLE,
                        value=body_sensor_location,
                    ),
                    cls=self.BodySensorLocation,
                    length=1,
                )
            )
            characteristics.append(self.body_sensor_location_characteristic)

        if reset_energy_expended:

            def write_heart_rate_control_point_value(
                connection: device.Connection, value: bytes
            ) -> None:
                if value == self.RESET_ENERGY_EXPENDED:
                    if reset_energy_expended is not None:
                        reset_energy_expended(connection)
                else:
                    raise att.ATT_Error(self.CONTROL_POINT_NOT_SUPPORTED)

            self.heart_rate_control_point_characteristic = (
                gatt_adapters.PackedCharacteristicAdapter(
                    gatt.Characteristic(
                        uuid=gatt.GATT_HEART_RATE_CONTROL_POINT_CHARACTERISTIC,
                        properties=gatt.Characteristic.Properties.WRITE,
                        permissions=gatt.Characteristic.WRITEABLE,
                        value=gatt.CharacteristicValue(
                            write=write_heart_rate_control_point_value
                        ),
                    ),
                    pack_format=HeartRateService.HEART_RATE_CONTROL_POINT_FORMAT,
                )
            )
            characteristics.append(self.heart_rate_control_point_characteristic)

        super().__init__(characteristics)


# -----------------------------------------------------------------------------
class HeartRateServiceProxy(gatt_client.ProfileServiceProxy):
    SERVICE_CLASS = HeartRateService

    heart_rate_measurement: gatt_client.CharacteristicProxy[
        HeartRateService.HeartRateMeasurement
    ]
    body_sensor_location: (
        gatt_client.CharacteristicProxy[HeartRateService.BodySensorLocation] | None
    )
    heart_rate_control_point: gatt_client.CharacteristicProxy[int] | None

    def __init__(self, service_proxy: gatt_client.ServiceProxy) -> None:
        self.service_proxy = service_proxy

        self.heart_rate_measurement = (
            gatt_adapters.SerializableCharacteristicProxyAdapter(
                service_proxy.get_required_characteristic_by_uuid(
                    gatt.GATT_HEART_RATE_MEASUREMENT_CHARACTERISTIC
                ),
                HeartRateService.HeartRateMeasurement,
            )
        )

        if characteristics := service_proxy.get_characteristics_by_uuid(
            gatt.GATT_BODY_SENSOR_LOCATION_CHARACTERISTIC
        ):
            self.body_sensor_location = gatt_adapters.EnumCharacteristicProxyAdapter(
                characteristics[0], cls=HeartRateService.BodySensorLocation, length=1
            )
        else:
            self.body_sensor_location = None

        if characteristics := service_proxy.get_characteristics_by_uuid(
            gatt.GATT_HEART_RATE_CONTROL_POINT_CHARACTERISTIC
        ):
            self.heart_rate_control_point = (
                gatt_adapters.PackedCharacteristicProxyAdapter(
                    characteristics[0],
                    pack_format=HeartRateService.HEART_RATE_CONTROL_POINT_FORMAT,
                )
            )
        else:
            self.heart_rate_control_point = None

    async def reset_energy_expended(self) -> None:
        if self.heart_rate_control_point is not None:
            return await self.heart_rate_control_point.write_value(
                HeartRateService.RESET_ENERGY_EXPENDED
            )