use crate::network::{NetworkTransport, NetError};
use iroh::{Endpoint, EndpointId};
use iroh::endpoint::Connection;
use bytes::Bytes;
use tokio::sync::{mpsc, Mutex};
use tokio::sync::mpsc::Receiver;
use std::sync::Arc;
use std::collections::HashMap;
use async_trait::async_trait;

#[derive(Debug, Clone)]
pub struct AudioProtocol {
    incoming_tx: mpsc::Sender<(EndpointId, Bytes)>,
    connections: Arc<Mutex<HashMap<EndpointId, Connection>>>,
}

impl iroh::protocol::ProtocolHandler for AudioProtocol {
    fn accept(
        &self,
        connection: Connection,
    ) -> impl std::future::Future<Output = Result<(), iroh::protocol::AcceptError>> + Send {
        let peer_id = connection.remote_id();
        let incoming_tx = self.incoming_tx.clone();
        let connections = self.connections.clone();

        async move {
            connections.lock().await.insert(peer_id, connection.clone());
            loop {
                match connection.read_datagram().await {
                    Ok(bytes) => {
                        if incoming_tx.send((peer_id, bytes)).await.is_err() {
                            break;
                        }
                    }
                    Err(_) => {
                        connections.lock().await.remove(&peer_id);
                        break;
                    }
                }
            }
            Ok(())
        }
    }
}

pub struct IrohTransport {
    endpoint: Endpoint,
    connections: Arc<Mutex<HashMap<EndpointId, Connection>>>,
    incoming_tx: mpsc::Sender<(EndpointId, Bytes)>,
    incoming_rx: Mutex<Option<mpsc::Receiver<(EndpointId, Bytes)>>>,
}

impl IrohTransport {
    pub fn new(endpoint: Endpoint) -> (Self, AudioProtocol) {
        let (incoming_tx, incoming_rx) = mpsc::channel(1000);
        let connections = Arc::new(Mutex::new(HashMap::new()));

        let audio_proto = AudioProtocol {
            incoming_tx: incoming_tx.clone(),
            connections: connections.clone(),
        };

        let transport = Self {
            endpoint,
            connections,
            incoming_tx,
            incoming_rx: Mutex::new(Some(incoming_rx)),
        };

        (transport, audio_proto)
    }
}

#[async_trait]
impl NetworkTransport for IrohTransport {
    async fn send_datagram(&self, peer_id: EndpointId, data: Bytes) -> Result<(), NetError> {
        let mut conns = self.connections.lock().await;
        let conn = if let Some(conn) = conns.get(&peer_id) {
            conn.clone()
        } else {
            // Establish a new connection.
            // We use the same audio ALPN: b"peerspeak-audio"
            let alpn = b"peerspeak-audio";
            let conn = self.endpoint.connect(peer_id, alpn).await
                .map_err(|e| NetError::Connection(e.to_string()))?;

            conns.insert(peer_id, conn.clone());

            let incoming_tx_inner = self.incoming_tx.clone();
            let connections_inner = self.connections.clone();
            let conn_clone = conn.clone();

            tokio::spawn(async move {
                loop {
                    match conn_clone.read_datagram().await {
                        Ok(bytes) => {
                            if let Err(_) = incoming_tx_inner.send((peer_id, bytes)).await {
                                break;
                            }
                        }
                        Err(_) => {
                            connections_inner.lock().await.remove(&peer_id);
                            break;
                        }
                    }
                }
            });

            conn
        };

        conn.send_datagram(data)
            .map_err(|e| NetError::Connection(e.to_string()))?;
        Ok(())
    }

    async fn receive_datagrams(&self) -> Result<Receiver<(EndpointId, Bytes)>, NetError> {
        let mut rx_guard = self.incoming_rx.lock().await;
        if let Some(rx) = rx_guard.take() {
            Ok(rx)
        } else {
            Err(NetError::Other("Datagram receiver already subscribed".to_string()))
        }
    }
}