# Copyright 2024 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 pprint
import asyncio
import contextlib
import logging
import wave
import itertools
import struct
from typing import cast, Any, AsyncGenerator, Coroutine, Dict, Optional, Tuple, List
import itertools


try:
    import lc3  # type: ignore  # pylint: disable=E0401
except ImportError as e:
    raise ImportError("Try `python -m pip install \".[lc3]\"`.") from e

from bumble.colors import color
from bumble import company_ids
from bumble import core
from bumble import gatt
from bumble import hci
from bumble.profiles import bap
from bumble.profiles import le_audio
from bumble.profiles import pbp
from bumble.profiles import bass
import bumble.device
import bumble.transport
import bumble.utils
from bumble.device import Host, BIGInfoAdvertisement, AdvertisingChannelMap
from bumble.audio import io as audio_io

from auracast import auracast_config


# modified from bumble
class ModWaveAudioInput(audio_io.ThreadedAudioInput):
    """Audio input that reads PCM samples from a .wav file."""

    def __init__(self, filename: str) -> None:
        super().__init__()
        self._filename = filename
        self._wav: wave.Wave_read | None = None
        self._bytes_read = 0
        self.rewind=True

    def _open(self) -> audio_io.PcmFormat:
        self._wav = wave.open(self._filename, 'rb')
        if self._wav.getsampwidth() != 2:
            raise ValueError('sample width not supported')
        return audio_io.PcmFormat(
            audio_io.PcmFormat.Endianness.LITTLE,
            audio_io.PcmFormat.SampleType.INT16,
            self._wav.getframerate(),
            self._wav.getnchannels(),
        )

    def _read(self, frame_size: int) -> bytes:
        if not self._wav:
            return b''

        pcm_samples = self._wav.readframes(frame_size)
        if not pcm_samples and self._bytes_read:
            if not self.rewind:
                return None
            # Loop around.
            self._wav.rewind()
            self._bytes_read = 0
            pcm_samples = self._wav.readframes(frame_size)

        self._bytes_read += len(pcm_samples)
        return pcm_samples

    def _close(self) -> None:
        if self._wav:
            self._wav.close()

audio_io.WaveAudioInput = ModWaveAudioInput

# -----------------------------------------------------------------------------
# Logging
# -----------------------------------------------------------------------------
logger = logging.getLogger(__name__)

@contextlib.asynccontextmanager
async def create_device(config: auracast_config.AuracastGlobalConfig) -> AsyncGenerator[bumble.device.Device, Any]:
    async with await bumble.transport.open_transport(config.transport) as (
        hci_source,
        hci_sink,
    ):
        device_config = bumble.device.DeviceConfiguration(
            name=config.device_name,
            address=config.auracast_device_address,
            keystore='JsonKeyStore',
            #le_simultaneous_enabled=True #TODO: What is this doing ?
        )

        device = bumble.device.Device.from_config_with_hci(
            device_config,
            hci_source,
            hci_sink,
        )
        await device.power_on()

        yield device


def run_async(async_command: Coroutine) -> None:
    try:
        asyncio.run(async_command)
    except core.ProtocolError as error:
        if error.error_namespace == 'att' and error.error_code in list(
            bass.ApplicationError
        ):
            message = bass.ApplicationError(error.error_code).name
        else:
            message = str(error)

        print(
            color('!!! An error occurred while executing the command:', 'red'), message
        )

async def init_broadcast(
        device,
        global_config : auracast_config.AuracastGlobalConfig,
        big_config: List[auracast_config.AuracastBigConfig]
        ) -> dict:

    bap_sampling_freq = getattr(bap.SamplingFrequency, f"FREQ_{global_config.auracast_sampling_rate_hz}")
    bigs = {}
    for i, conf in enumerate(big_config):
        bigs[f'big{i}'] = {}
        # Config advertising set
        bigs[f'big{i}']['basic_audio_announcement'] = bap.BasicAudioAnnouncement(
            presentation_delay=global_config.presentation_delay_us,
            subgroups=[
                bap.BasicAudioAnnouncement.Subgroup(
                    codec_id=hci.CodingFormat(codec_id=hci.CodecID.LC3),
                    codec_specific_configuration=bap.CodecSpecificConfiguration(
                        sampling_frequency=bap_sampling_freq,
                        frame_duration=bap.FrameDuration.DURATION_10000_US,
                        octets_per_codec_frame=global_config.octets_per_frame,
                    ),
                    metadata=le_audio.Metadata(
                        [
                            le_audio.Metadata.Entry(
                                tag=le_audio.Metadata.Tag.LANGUAGE, data=conf.language.encode()
                            ),
                            le_audio.Metadata.Entry(
                                tag=le_audio.Metadata.Tag.PROGRAM_INFO, data=conf.program_info.encode()
                            ),
                            le_audio.Metadata.Entry(
                            tag=le_audio.Metadata.Tag.BROADCAST_NAME, data=conf.name.encode()
                            ),
                        ]
                    ),
                    bis=[
                        bap.BasicAudioAnnouncement.BIS(
                            index=1,
                            codec_specific_configuration=bap.CodecSpecificConfiguration(
                                audio_channel_allocation=bap.AudioLocation.FRONT_LEFT
                            ),
                        ),
                    ],
                )
            ],
        )
        logger.info('Setup Advertising')
        advertising_manufacturer_data = (
        b''
        if global_config.manufacturer_data is None
        else bytes(
            core.AdvertisingData(
                [
                    (
                        core.AdvertisingData.MANUFACTURER_SPECIFIC_DATA,
                        struct.pack('<H', global_config.manufacturer_data[0])
                        + global_config.manufacturer_data[1],
                    )
                ]
            )
        )
    )
        bigs[f'big{i}']['broadcast_audio_announcement'] = bap.BroadcastAudioAnnouncement(conf.id)
        advertising_set = await device.create_advertising_set(
            random_address=conf.random_address,
            advertising_parameters=bumble.device.AdvertisingParameters(
                advertising_event_properties=bumble.device.AdvertisingEventProperties(
                    is_connectable=False
                ),
                primary_advertising_interval_min=100,
                primary_advertising_interval_max=200,
                advertising_sid=i,
                primary_advertising_phy=hci.Phy.LE_1M, # 2m phy config throws error - because for primary advertising channels, 1mbit is only supported
                secondary_advertising_phy=hci.Phy.LE_2M, # this is the secondary advertising beeing send on non advertising channels (extendend advertising)
                #advertising_tx_power= # tx power in dbm (max 20)
                #secondary_advertising_max_skip=10,
            ),
            advertising_data=(
                bigs[f'big{i}']['broadcast_audio_announcement'].get_advertising_data()
                + bytes(
                    core.AdvertisingData(
                        [(core.AdvertisingData.BROADCAST_NAME, conf.name.encode())]
                    )
                )
                + advertising_manufacturer_data
            ),
            periodic_advertising_parameters=bumble.device.PeriodicAdvertisingParameters(
                periodic_advertising_interval_min=80,
                periodic_advertising_interval_max=160,
            ),
            periodic_advertising_data=bigs[f'big{i}']['basic_audio_announcement'].get_advertising_data(),
            auto_restart=True,
            auto_start=True,
        )
        bigs[f'big{i}']['advertising_set'] = advertising_set

        logging.info('Start Periodic Advertising')
        await advertising_set.start_periodic()

        logging.info('Setup BIG')
        if global_config.qos_config.iso_int_multiple_10ms == 1:
            frame_enable = 0
        else:
            frame_enable = 1

        big = await device.create_big(
            bigs[f'big{i}']['advertising_set'],
            parameters=bumble.device.BigParameters(
                num_bis=1,
                sdu_interval=global_config.qos_config.iso_int_multiple_10ms*10000, # Is the same as iso interval
                max_sdu=global_config.octets_per_frame,
                max_transport_latency=global_config.qos_config.max_transport_latency_ms,
                rtn=global_config.qos_config.number_of_retransmissions,
                broadcast_code=(
                    bytes.fromhex(conf.code) if conf.code else None
                ),
                framing=frame_enable # needed if iso interval is not frame interval of codedc
            ),
        )
        bigs[f'big{i}']['big'] = big

        for bis_link in big.bis_links:
            await bis_link.setup_data_path(
                direction=bis_link.Direction.HOST_TO_CONTROLLER
            )

        iso_queue = bumble.device.IsoPacketStream(big.bis_links[0], conf.iso_que_len)

        logging.info('Setup ISO Data Path')

        bigs[f'big{i}']['iso_queue'] = iso_queue

        logging.debug(f'big{i} parameters are:')
        logging.debug('%s', pprint.pformat(vars(big)))
        logging.debug(f'Finished setup of big{i}.')

    await asyncio.sleep(i+1) # Wait for advertising to set up

    def on_flow():
        data_packet_queue = iso_queue.data_packet_queue
        print(
            f'\rPACKETS: pending={data_packet_queue.pending}, '
            f'queued={data_packet_queue.queued}, '
            f'completed={data_packet_queue.completed}',
            end='',
        )

    if global_config.debug:
        bigs[f'big{0}']['iso_queue'].data_packet_queue.on('flow', on_flow)

    return bigs

async def init_audio(
        bigs,
        global_config : auracast_config.AuracastGlobalConfig,
        big_config: List[auracast_config.AuracastBigConfig]
    ):
    for i, big in enumerate(bigs.values()):
        audio_source = big_config[i].audio_source
        input_format = big_config[i].input_format
        audio_input = await audio_io.create_audio_input(audio_source, input_format)
        audio_input.rewind = big_config[i].loop_wav
        pcm_format = await audio_input.open()

        #try:
        if pcm_format.channels != 1:
            print("Only 1 channels PCM configurations are supported")
            return
        if pcm_format.sample_type == audio_io.PcmFormat.SampleType.INT16:
            pcm_bit_depth = 16
        elif pcm_format.sample_type == audio_io.PcmFormat.SampleType.FLOAT32:
            pcm_bit_depth = None
        else:
            print("Only INT16 and FLOAT32 sample types are supported")
            return
        encoder = lc3.Encoder(
            frame_duration_us=global_config.frame_duration_us,
            sample_rate_hz=global_config.auracast_sampling_rate_hz,
            num_channels=1,
            input_sample_rate_hz=pcm_format.sample_rate,
        )
        lc3_frame_samples = encoder.get_frame_samples() # number of the pcm samples per lc3 frame
        lc3_frame_size = global_config.octets_per_frame #encoder.get_frame_bytes(bitrate)
        lc3_bytes_per_frame = lc3_frame_size #* 2 #multiplied by number of channels

        big['pcm_bit_depth'] = pcm_bit_depth
        big['lc3_bytes_per_frame'] = lc3_bytes_per_frame
        big['lc3_frame_samples'] = lc3_frame_samples
        big['audio_input'] = audio_input
        big['encoder'] = encoder

class Streamer():
    def __init__(self, bigs):
        self.task = None
        self.is_streaming = False
        self.bigs = bigs

    def start_streaming(self):
        if not self.is_streaming:
            self.task = asyncio.create_task(self.stream())
        else:
            logging.warning('Streamer is already running')

    def stop_streaming(self):
        """Stops the background task if running."""
        if self.is_streaming:
            self.is_streaming = False
            if self.task:
                self.task.cancel()  # Cancel the task safely
                self.task = None

    async def stream(self):
        # TODO: do some pre buffering so the stream is stable from the beginning. One half iso queue would be appropriate
        logging.info("Streaming audio...")
        bigs = self.bigs
        self.is_streaming = True
        while self.is_streaming:
            stream_finished = [False for _ in range(len(bigs))]
            for i, big in enumerate(bigs.values()):
                pcm_frame = await anext(big['audio_input'].frames(big['lc3_frame_samples']), None)
                if pcm_frame is None: # Not all streams may stop at the same time
                    stream_finished[i] = True
                    continue

                lc3_frame = big['encoder'].encode(
                    pcm_frame, num_bytes=big['lc3_bytes_per_frame'], bit_depth=big['pcm_bit_depth']
                )
                await big['iso_queue'].write(lc3_frame)

            if all(stream_finished): # Take into account that multiple files have different lengths
                logging.info('All streams finished, stopping streamer')
                self.is_streaming = False
                break


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------

async def broadcast(global_conf: auracast_config.AuracastGlobalConfig, big_conf: List[auracast_config.AuracastBigConfig]):
    """Start a broadcast."""
    async with create_device(global_conf) as device:
        if not device.supports_le_periodic_advertising:
            logger.error(color('Periodic advertising not supported', 'red'))
            return

        bigs = await init_broadcast( # the bigs dictionary contains all the global configurations
            device,
            global_conf,
            big_conf
        )
        await init_audio(
            bigs,
            global_conf,
            big_conf
        )
        streamer = Streamer(bigs)
        streamer.start_streaming()

        await asyncio.wait([streamer.task])


# -----------------------------------------------------------------------------
if __name__ == "__main__":
    logging.basicConfig(
        level=logging.DEBUG,
        format='%(module)s.py:%(lineno)d %(levelname)s: %(message)s'
    )

    global_conf = auracast_config.global_base_config

    #global_conf.transport='serial:/dev/serial/by-id/usb-ZEPHYR_Zephyr_HCI_UART_sample_81BD14B8D71B5662-if00,1000000,rtscts' # transport for nrf52 dongle

    #global_conf.transport='serial:/dev/serial/by-id/usb-SEGGER_J-Link_001050076061-if02,1000000,rtscts' # transport for nrf53dk

    #global_conf.transport='serial:/dev/serial/by-id/usb-SEGGER_J-Link_001057705357-if02,1000000,rtscts' # transport for nrf54l15dk

    global_conf.transport='serial:/dev/serial/by-id/usb-ZEPHYR_Zephyr_HCI_UART_sample_81BD14B8D71B5662-if00,115200,rtscts' #nrf52dongle hci_uart usb cdc

    # global_conf.transport='usb:2fe3:000b' #nrf52dongle hci_usb # TODO: iso packet over usb not supported


    # TODO: How can we use other iso interval than 10ms ?(medium or low rel) ? - nrf53audio receiver repports I2S tx underrun
    #global_conf.qos_config = auracast_config.qos_config_mono_medium_rel
    global_conf.qos_config = auracast_config.qos_config_mono_high_rel

    bigs = [
        auracast_config.broadcast_de,
        auracast_config.broadcast_en,
        auracast_config.broadcast_fr,
        #auracast_config.broadcast_es,
        #auracast_config.broadcast_it,
        ]
    #for big in bigs: # TODO. investigate this further
    #    big.code = 'ff'*16 # returns hci/HCI_ENCRYPTION_MODE_NOT_ACCEPTABLE_ERROR
        #big.code = '78 e5 dc f1 34 ab 42 bf c1 92 ef dd 3a fd 67 ae'
    

    # 16kHz works reliably with 3 streams
    # 24kHz is only working with 2 streams - probably airtime constraint
    # TODO: with more than three broadcasters (16kHz) no advertising (no primary channels is present anymore)
    # TODO: find the bottleneck - probably airtime
    # TODO: test encrypted streams

    global_conf.auracast_sampling_rate_hz = 16000
    global_conf.octets_per_frame = 40 # 32kbps@16kHz
    #global_conf.debug = True

    run_async(
        broadcast(
            global_conf,
            bigs
        )
    )

    # TODO: possible inputs:
    # wav file locally
    # precoded lc3 file locally
    # realtime audio locally
    # realtime audio network lc3 coded
    # (realtime audio network uncoded)