// Copyright 2026 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use futures::stream::{BoxStream, StreamExt};
use std::collections::HashMap;

use bt_common::Uuid;
use bt_gatt::client::{CharacteristicNotification, PeerService, ServiceCharacteristic};
use bt_gatt::types::Handle;

use crate::types::*;

/// 16-bit UUID value for the characteristics offered by the Audio
/// Stream Control Service.
pub const ASE_CONTROL_POINT_UUID: Uuid = Uuid::from_u16(0x2BC6);
pub const SINK_ASE_UUID: Uuid = Uuid::from_u16(0x2BC4);
pub const SOURCE_ASE_UUID: Uuid = Uuid::from_u16(0x2BC5);

#[derive(Debug, thiserror::Error, PartialEq, Clone)]
pub enum ClientError {
    #[error("Remote server is missing control point characteristic")]
    MissingControlPointCharacteristic,
    #[error("Remote server should only have 1 control point characteristic")]
    ExtraControlPointCharacteristic,
    #[error("Remote server doesn't have any audio stream endpoints")]
    MissingAudioStreamEndpoints,
}

/// Represents a single source/sink ASE state characteristic.
pub struct AudioStreamEndpointHandle {
    pub endpoint: AudioStreamEndpoint,
    pub notification_stream:
        BoxStream<'static, Result<CharacteristicNotification, bt_gatt::types::Error>>,
}

// See ASCS v1.0.1 4.2 for details.
pub struct AseControlPoint {
    pub handle: Handle,
    pub notification_stream:
        BoxStream<'static, Result<CharacteristicNotification, bt_gatt::types::Error>>,
}

/// Creates Audio Stream Control Service (ASCS) client instance
pub struct AudioStreamControlServiceClient<T: bt_gatt::GattTypes> {
    pub gatt_client: T::PeerService,
    pub control_point: AseControlPoint,
    pub sink_endpoints: HashMap<Handle, AudioStreamEndpointHandle>,
    pub source_endpoints: HashMap<Handle, AudioStreamEndpointHandle>,
}

pub type SinkAndSourceEndpoints =
    (HashMap<Handle, AudioStreamEndpointHandle>, HashMap<Handle, AudioStreamEndpointHandle>);

impl<T: bt_gatt::GattTypes> AudioStreamControlServiceClient<T> {
    pub async fn create(gatt_client: T::PeerService) -> Result<Self, Error>
    where
        <T as bt_gatt::GattTypes>::NotificationStream: std::marker::Send,
    {
        let control_point = Self::discover_control_point(&gatt_client).await?;
        let (sink_endpoints, source_endpoints) = Self::discover_all_endpoints(&gatt_client).await?;

        Ok(Self { gatt_client, control_point, sink_endpoints, source_endpoints })
    }

    async fn discover_control_point(gatt_client: &T::PeerService) -> Result<AseControlPoint, Error>
    where
        <T as bt_gatt::GattTypes>::NotificationStream: std::marker::Send,
    {
        let cp_chars = ServiceCharacteristic::<T>::find(gatt_client, ASE_CONTROL_POINT_UUID)
            .await
            .map_err(Error::Gatt)?;
        if cp_chars.is_empty() {
            return Err(ClientError::MissingControlPointCharacteristic.into());
        }
        if cp_chars.len() > 1 {
            return Err(ClientError::ExtraControlPointCharacteristic.into());
        }
        let cp_handle = *cp_chars[0].handle();
        let cp_stream = gatt_client.subscribe(&cp_handle);
        Ok(AseControlPoint { handle: cp_handle, notification_stream: cp_stream.boxed() })
    }

    async fn discover_all_endpoints(
        gatt_client: &T::PeerService,
    ) -> Result<SinkAndSourceEndpoints, Error>
    where
        <T as bt_gatt::GattTypes>::NotificationStream: std::marker::Send,
    {
        let sink_chars = ServiceCharacteristic::<T>::find(gatt_client, SINK_ASE_UUID)
            .await
            .map_err(Error::Gatt)?;

        let source_chars = ServiceCharacteristic::<T>::find(gatt_client, SOURCE_ASE_UUID)
            .await
            .map_err(Error::Gatt)?;

        if sink_chars.is_empty() && source_chars.is_empty() {
            return Err(ClientError::MissingAudioStreamEndpoints.into());
        }

        let mut sink_endpoints = HashMap::new();
        for c in sink_chars {
            let handle = *c.handle();
            let endpoint =
                Self::read_and_create_endpoint(gatt_client, handle, AudioDirection::Sink).await?;
            let notification_stream = gatt_client.subscribe(&handle).boxed();
            sink_endpoints
                .insert(handle, AudioStreamEndpointHandle { endpoint, notification_stream });
        }

        let mut source_endpoints = HashMap::new();
        for c in source_chars {
            let handle = *c.handle();
            let endpoint =
                Self::read_and_create_endpoint(gatt_client, handle, AudioDirection::Source).await?;
            let notification_stream = gatt_client.subscribe(&handle).boxed();
            source_endpoints
                .insert(handle, AudioStreamEndpointHandle { endpoint, notification_stream });
        }

        Ok((sink_endpoints, source_endpoints))
    }

    async fn read_and_create_endpoint(
        gatt_client: &T::PeerService,
        handle: Handle,
        direction: AudioDirection,
    ) -> Result<AudioStreamEndpoint, Error> {
        let mut buf = vec![0; 255];
        let (read_bytes, _truncated) =
            gatt_client.read_characteristic(&handle, 0, &mut buf[..]).await.map_err(Error::Gatt)?;
        let endpoint = AudioStreamEndpoint::from_char_value(handle, direction, &buf[0..read_bytes])
            .map_err(|e| {
                bt_gatt::types::Error::other(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    format!("Failed to decode AudioStreamEndpoint: {:?}", e),
                ))
            })?;
        Ok(endpoint)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bt_gatt::test_utils::{FakePeerService, FakeTypes};
    use bt_gatt::types::{AttributePermissions, CharacteristicProperty};
    use bt_gatt::Characteristic;

    const CONTROL_POINT_HANDLE: Handle = Handle(1);
    const SINK_ASE_HANDLE: Handle = Handle(2);
    const SOURCE_ASE_HANDLE: Handle = Handle(3);

    fn setup_fake_service() -> FakePeerService {
        let mut service = FakePeerService::new();
        // Add Control Point
        service.add_characteristic(
            Characteristic {
                handle: CONTROL_POINT_HANDLE,
                uuid: ASE_CONTROL_POINT_UUID,
                properties: CharacteristicProperty::Write
                    | CharacteristicProperty::WriteWithoutResponse
                    | CharacteristicProperty::Notify,
                permissions: AttributePermissions::default(),
                descriptors: vec![],
            },
            vec![],
        );
        // Add Sink ASE
        service.add_characteristic(
            Characteristic {
                handle: SINK_ASE_HANDLE,
                uuid: SINK_ASE_UUID,
                properties: CharacteristicProperty::Read | CharacteristicProperty::Notify,
                permissions: AttributePermissions::default(),
                descriptors: vec![],
            },
            vec![0x01, 0x00], // ASE ID 1, State Idle
        );
        // Add Source ASE
        service.add_characteristic(
            Characteristic {
                handle: SOURCE_ASE_HANDLE,
                uuid: SOURCE_ASE_UUID,
                properties: CharacteristicProperty::Read | CharacteristicProperty::Notify,
                permissions: AttributePermissions::default(),
                descriptors: vec![],
            },
            vec![0x02, 0x00], // ASE ID 2, State Idle
        );
        service
    }

    fn run_to_completion<F: std::future::Future>(fut: F) -> F::Output {
        let mut fut = std::pin::pin!(fut);
        let mut cx = futures::task::Context::from_waker(futures::task::noop_waker_ref());
        match fut.as_mut().poll(&mut cx) {
            std::task::Poll::Ready(res) => res,
            std::task::Poll::Pending => panic!("Future did not complete synchronously"),
        }
    }

    #[test]
    fn client_creation_success() {
        let service = setup_fake_service();
        let client_fut = AudioStreamControlServiceClient::<FakeTypes>::create(service.clone());
        let client = run_to_completion(client_fut).expect("client creation should succeed");

        assert_eq!(client.control_point.handle, CONTROL_POINT_HANDLE);
        assert_eq!(client.sink_endpoints.len(), 1);
        assert_eq!(client.source_endpoints.len(), 1);

        let sink = &client.sink_endpoints[&SINK_ASE_HANDLE].endpoint;
        assert_eq!(sink.ase_id, AseId(1));
        assert_eq!(sink.state, AseState::Idle);
        assert_eq!(sink.direction, AudioDirection::Sink);

        let source = &client.source_endpoints[&SOURCE_ASE_HANDLE].endpoint;
        assert_eq!(source.ase_id, AseId(2));
        assert_eq!(source.state, AseState::Idle);
        assert_eq!(source.direction, AudioDirection::Source);
    }

    #[test]
    fn client_creation_missing_control_point() {
        let mut service = FakePeerService::new();
        // Only add Sink (no Control Point)
        service.add_characteristic(
            Characteristic {
                handle: SINK_ASE_HANDLE,
                uuid: SINK_ASE_UUID,
                properties: CharacteristicProperty::Read | CharacteristicProperty::Notify,
                permissions: AttributePermissions::default(),
                descriptors: vec![],
            },
            vec![0x01, 0x00],
        );

        let client_fut = AudioStreamControlServiceClient::<FakeTypes>::create(service);
        let err = match run_to_completion(client_fut) {
            Err(e) => e,
            Ok(_) => panic!("expected error, got Ok"),
        };
        assert!(matches!(err, Error::Client(ClientError::MissingControlPointCharacteristic)));
    }

    #[test]
    fn client_creation_missing_endpoints() {
        let mut service = FakePeerService::new();
        // Only add Control Point (no endpoints)
        service.add_characteristic(
            Characteristic {
                handle: CONTROL_POINT_HANDLE,
                uuid: ASE_CONTROL_POINT_UUID,
                properties: CharacteristicProperty::Write | CharacteristicProperty::Notify,
                permissions: AttributePermissions::default(),
                descriptors: vec![],
            },
            vec![],
        );

        let client_fut = AudioStreamControlServiceClient::<FakeTypes>::create(service);
        let err = match run_to_completion(client_fut) {
            Err(e) => e,
            Ok(_) => panic!("expected error, got Ok"),
        };
        assert!(matches!(err, Error::Client(ClientError::MissingAudioStreamEndpoints)));
    }

    #[test]
    fn client_creation_extra_control_point() {
        let mut service = FakePeerService::new();
        // Add first Control Point
        service.add_characteristic(
            Characteristic {
                handle: CONTROL_POINT_HANDLE,
                uuid: ASE_CONTROL_POINT_UUID,
                properties: CharacteristicProperty::Write | CharacteristicProperty::Notify,
                permissions: AttributePermissions::default(),
                descriptors: vec![],
            },
            vec![],
        );
        // Add second Control Point
        service.add_characteristic(
            Characteristic {
                handle: Handle(4),
                uuid: ASE_CONTROL_POINT_UUID,
                properties: CharacteristicProperty::Write | CharacteristicProperty::Notify,
                permissions: AttributePermissions::default(),
                descriptors: vec![],
            },
            vec![],
        );

        let client_fut = AudioStreamControlServiceClient::<FakeTypes>::create(service);
        let err = match run_to_completion(client_fut) {
            Err(e) => e,
            Ok(_) => panic!("expected error, got Ok"),
        };
        assert!(matches!(err, Error::Client(ClientError::ExtraControlPointCharacteristic)));
    }
}