// 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 bt_common::PeerId;
use bt_gatt::GattTypes;
use std::collections::HashMap;

use crate::client::BapUnicastClient;
use crate::types::{Error, SupportedCapabilities};

/// Manages multiple BAP Unicast peers.
pub struct BapUnicastManager<T: GattTypes> {
    peers: HashMap<PeerId, BapUnicastClient<T>>,
}

impl<T: GattTypes> BapUnicastManager<T> {
    pub fn new() -> Self {
        Self { peers: HashMap::new() }
    }

    /// Add a peer to the manager.
    /// Triggers the connection and discovery of PACS and ASCS on the remote
    /// peer.
    pub async fn add_peer(
        &mut self,
        peer_id: PeerId,
        gatt_client: T::Client,
    ) -> Result<SupportedCapabilities, Error>
    where
        <T as GattTypes>::NotificationStream: Send,
    {
        if self.peers.contains_key(&peer_id) {
            return Err(Error::PeerAlreadyExists(peer_id));
        }

        let mut client = BapUnicastClient::new(peer_id, gatt_client);
        let mapping = client.discover().await?;
        self.peers.insert(peer_id, client);
        Ok(mapping)
    }

    /// Remove a peer from the manager.
    pub fn remove_peer(&mut self, peer_id: PeerId) -> Result<(), Error> {
        if self.peers.remove(&peer_id).is_none() {
            return Err(Error::PeerNotFound(peer_id));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{setup_ascs_service, setup_pacs_service};
    use crate::Configuration;
    use bt_common::PeerId;
    use bt_gatt::test_utils::{FakeClient, FakeTypes};

    #[test]
    fn test_add_remove_peer() {
        let mut manager = BapUnicastManager::<FakeTypes>::new();
        let peer_id = PeerId(1);
        let mut client = FakeClient::new();

        // 1. Setup PACS
        let _pacs_service = setup_pacs_service(&mut client);

        // 2. Setup ASCS
        let _ascs_service = setup_ascs_service(&mut client, 1, 1);

        let mapping =
            futures::executor::block_on(manager.add_peer(peer_id, client.clone())).unwrap();
        assert!(mapping.contains_key(&Configuration::MonoAudioSink));
        assert!(mapping.contains_key(&Configuration::Conversational));

        assert!(futures::executor::block_on(manager.add_peer(peer_id, client.clone())).is_err()); // Duplicate

        assert!(manager.remove_peer(peer_id).is_ok());
        assert!(manager.remove_peer(peer_id).is_err()); // Not found
    }

    #[test]
    fn test_add_peer_fails_non_unique_services() {
        let mut manager = BapUnicastManager::<FakeTypes>::new();
        let peer_id = PeerId(1);

        // 0 PACS services: add_peer fails and peer is not stored.
        let mut client = FakeClient::new();
        let _ascs_service = setup_ascs_service(&mut client, 1, 1);
        assert!(futures::executor::block_on(manager.add_peer(peer_id, client.clone())).is_err());
        assert!(manager.remove_peer(peer_id).is_err());

        // Multiple PACS services: add_peer fails and peer is not stored.
        let mut client = FakeClient::new();
        let _pacs1 = setup_pacs_service(&mut client);
        let _pacs2 = setup_pacs_service(&mut client);
        let _ascs = setup_ascs_service(&mut client, 1, 1);
        assert!(futures::executor::block_on(manager.add_peer(peer_id, client.clone())).is_err());
        assert!(manager.remove_peer(peer_id).is_err());

        // 0 ASCS services: add_peer fails and peer is not stored.
        let mut client = FakeClient::new();
        let _pacs = setup_pacs_service(&mut client);
        assert!(futures::executor::block_on(manager.add_peer(peer_id, client.clone())).is_err());
        assert!(manager.remove_peer(peer_id).is_err());

        // Multiple ASCS services: add_peer fails and peer is not stored.
        let mut client = FakeClient::new();
        let _pacs = setup_pacs_service(&mut client);
        let _ascs1 = setup_ascs_service(&mut client, 1, 1);
        let _ascs2 = setup_ascs_service(&mut client, 1, 1);
        assert!(futures::executor::block_on(manager.add_peer(peer_id, client.clone())).is_err());
        assert!(manager.remove_peer(peer_id).is_err());
    }
}