use crate::{
    FusionBackend, FusionDevice, FusionHandle, FusionRuntime, FusionServer, FusionTensor,
    FusionUtilities, UnfusedOp,
    stream::{StreamId, execution::Operation},
};
#[cfg(feature = "distributed")]
use burn_backend::distributed::DistributedBackend;
use burn_backend::{Device, DeviceHandle, DeviceId, DeviceService, DeviceServiceStage};
#[cfg(feature = "distributed")]
use burn_std::CommunicationId;

use burn_backend::{TensorData, backend::ExecutionError};
use burn_ir::{OperationIr, TensorId, TensorIr};
use burn_std::stub::RwLock;
use hashbrown::HashSet;
use std::sync::atomic::{AtomicU64, Ordering};

/// Use a mutex to communicate with the fusion server.
pub struct GlobalFusionClient<R: FusionRuntime> {
    server: DeviceHandle<FusionServer<R>>,
    device: FusionDevice<R>,
}

impl<R: FusionRuntime> DeviceService for FusionServer<R> {
    fn init(device_id: DeviceId) -> Self {
        let device = FusionDevice::<R>::from_id(device_id);
        let utilities = FusionUtilities {
            initialized_comms: RwLock::new(HashSet::default()),
        };
        FusionServer::new(device, utilities)
    }

    fn utilities(&self) -> burn_backend::ServerUtilitiesHandle {
        self.utilities.clone()
    }
    fn stage() -> DeviceServiceStage {
        DeviceServiceStage::Upstream
    }
}

impl<R> Clone for GlobalFusionClient<R>
where
    R: FusionRuntime,
{
    fn clone(&self) -> Self {
        Self {
            server: self.server.clone(),
            device: self.device.clone(),
        }
    }
}
impl<R> GlobalFusionClient<R>
where
    R: FusionRuntime + 'static,
{
    /// Loads the client from the given device.
    pub fn load(device: &FusionDevice<R>) -> Self {
        Self {
            device: device.clone(),
            server: DeviceHandle::new(device.to_id()),
        }
    }
}

static COUNTER: AtomicU64 = AtomicU64::new(0);

impl<R> GlobalFusionClient<R>
where
    R: FusionRuntime + 'static,
{
    /// Create a new client for the given [device](FusionRuntime::FusionDevice).
    pub fn new(device: FusionDevice<R>) -> Self {
        Self {
            device: device.clone(),
            server: DeviceHandle::new(device.to_id()),
        }
    }

    pub(crate) fn tag_shared_view(&self, src_stream: StreamId, src: TensorId, dst: TensorId) {
        let _ = self
            .server
            .submit_blocking(move |server| {
                server.tag_shared_view(src_stream, src, dst);
            })
            .ok();
    }

    /// Register a new [tensor operation intermediate representation](OperationIr) on `stream`.
    ///
    /// Returns the new (uninitialized) output tensor(s) generated by the registered operation.
    ///
    /// Cross-stream tensor sharing is resolved out-of-band by `tag_shared_view` before
    /// this call (see [`MultiStream`](crate::stream::MultiStream) docs for the full
    /// strategy), so every input tensor in `repr` is assumed to live on `stream`. Pass
    /// `StreamId::current()` for normal ops and the tensor's own stream for `Drop`.
    pub fn register<O>(
        &self,
        stream: StreamId,
        repr: OperationIr,
        operation: O,
    ) -> Vec<FusionTensor<R>>
    where
        O: Operation<R> + 'static,
    {
        // Create output tensors returned by this operation
        let outputs = repr
            .outputs()
            .map(|output| {
                FusionTensor::new(
                    output.id,
                    output.shape.clone(),
                    output.dtype,
                    self.clone(),
                    stream,
                )
            })
            .collect();

        // By doing this comparison, we reduce the number of bytes transferred in the device handle
        // queue.
        if size_of::<O>() < size_of::<UnfusedOp<R>>() {
            // Here the [`O`] type is smaller than the [`UnfusedOp`] type, so it's better to
            // transfer it directly.
            self.server.submit(move |server| {
                let operation = UnfusedOp::new(operation, stream);
                server.register(stream, repr, operation);
            });
        } else {
            // Here the [`O`] type is larger than the [`UnfusedOp`] type, so it's better to
            // first create the [`UnfusedOp`] before transferring it to the server.
            let operation = UnfusedOp::new(operation, stream);

            self.server.submit(move |server| {
                server.register(stream, repr, operation);
            });
        }

        outputs
    }

    /// Register all lazy computation.
    pub fn sync<Re: Send + 'static>(&self, sync_fn: impl FnOnce() -> Re + Send + 'static) -> Re {
        let id = StreamId::current();
        self.server
            .submit_blocking(move |server| {
                server.drain_stream(id);
                sync_fn()
            })
            .unwrap()
    }

    /// Flush the operations queue.
    pub fn flush_queue(&self) {
        self.server.flush_queue();
    }

    /// Create a new (uninitialized) empty tensor handle and returns its corresponding [tensor id](TensorId).
    pub fn create_empty_handle(&self) -> TensorId {
        let value = COUNTER.fetch_add(1, Ordering::Relaxed);
        TensorId::new(value)
    }

    /// Get the current device used by all operations handled by this client.
    pub fn device(&self) -> &FusionDevice<R> {
        &self.device
    }

    /// Create a tensor with the given handle and returns its corresponding [tensor id](TensorId).
    pub fn register_tensor_handle(&self, handle: FusionHandle<R>) -> TensorId {
        let id = self.create_empty_handle();

        self.server
            .submit(move |server| server.handles.register_handle(id, handle));

        id
    }

    /// Read the values contained by a float tensor.
    pub fn read_tensor_float<B>(
        self,
        tensor: TensorIr,
        stream: StreamId,
    ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        self.server
            .submit_blocking(move |server| server.float_data::<B>(tensor, stream))
            .unwrap()
    }

    /// Read the values contained by an int tensor.
    pub fn read_tensor_int<B>(
        self,
        tensor: TensorIr,
        stream: StreamId,
    ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        self.server
            .submit_blocking(move |server| server.int_data::<B>(tensor, stream))
            .unwrap()
    }

    /// Read the values contained by a bool tensor.
    pub fn read_tensor_bool<B>(
        self,
        tensor: TensorIr,
        stream: StreamId,
    ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        self.server
            .submit_blocking(move |server| server.bool_data::<B>(tensor, stream))
            .unwrap()
    }

    /// Read the values contained by a quantized tensor.
    pub fn read_tensor_quantized<B>(
        self,
        tensor: TensorIr,
        stream: StreamId,
    ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        self.server
            .submit_blocking(move |server| server.quantized_data::<B>(tensor, stream))
            .unwrap()
    }

    /// Change the client of the given float tensor.
    pub fn change_client_float<B>(
        tensor: TensorIr,
        client_src: Self,
        client_dst: Self,
        stream: StreamId,
    ) -> FusionTensor<R>
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        let dtype = tensor.dtype;
        let client_dst_cloned = client_dst.clone();
        let shape = tensor.shape.clone();
        let id = client_src.create_empty_handle();

        let float_tensor = client_src
            .server
            .submit_blocking(move |server_src| server_src.read_float::<B>(tensor, stream))
            .unwrap();
        let float_tensor = B::float_to_device(float_tensor, client_dst.device());
        client_dst.server.submit(move |server_dst| {
            server_dst
                .handles
                .register_float_tensor::<B>(&id, float_tensor.clone());
        });

        FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current())
    }

    /// Change the client of the given int tensor.
    pub fn change_client_int<B>(
        tensor: TensorIr,
        client_src: Self,
        client_dst: Self,
        stream: StreamId,
    ) -> FusionTensor<R>
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        let dtype = tensor.dtype;
        let client_dst_cloned = client_dst.clone();
        let shape = tensor.shape.clone();
        let id = client_src.create_empty_handle();

        let int_tensor = client_src
            .server
            .submit_blocking(move |server_src| server_src.read_int::<B>(tensor, stream))
            .unwrap();
        let int_tensor = B::int_to_device(int_tensor, client_dst.device());
        client_dst.server.submit(move |server_dst| {
            server_dst
                .handles
                .register_int_tensor::<B>(&id, int_tensor.clone());
        });

        FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current())
    }

    /// Change the client of the given bool tensor.
    pub fn change_client_bool<B>(
        tensor: TensorIr,
        client_src: Self,
        client_dst: Self,
        stream: StreamId,
    ) -> FusionTensor<R>
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        let dtype = tensor.dtype;
        let client_dst_cloned = client_dst.clone();
        let shape = tensor.shape.clone();
        let id = client_src.create_empty_handle();

        let bool_tensor = client_src
            .server
            .submit_blocking(move |server_src| server_src.read_bool::<B>(tensor, stream))
            .unwrap();
        let bool_tensor = B::bool_to_device(bool_tensor, client_dst.device());
        client_dst.server.submit(move |server_dst| {
            server_dst
                .handles
                .register_bool_tensor::<B>(&id, bool_tensor.clone());
        });

        FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current())
    }

    /// Change the client of the given quantized tensor.
    pub fn change_client_quantized<B>(
        tensor: TensorIr,
        client_src: Self,
        client_dst: Self,
        stream: StreamId,
    ) -> FusionTensor<R>
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        let dtype = tensor.dtype;
        let client_dst_cloned = client_dst.clone();
        let shape = tensor.shape.clone();
        let id = client_src.create_empty_handle();

        let q_tensor = client_src
            .server
            .submit_blocking(move |server_src| server_src.read_quantized::<B>(tensor, stream))
            .unwrap();
        let q_tensor = B::q_to_device(q_tensor, client_dst.device());
        client_dst.server.submit(move |server_dst| {
            server_dst
                .handles
                .register_quantized_tensor::<B>(&id, q_tensor.clone());
        });

        FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current())
    }

    /// Resolve the given float tensor to a primitive tensor.
    pub fn resolve_tensor_float<B>(&self, tensor: FusionTensor<R>) -> B::FloatTensorPrimitive
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        // Must evaluate `into_ir()` on the frontend thread. If evaluated inside
        // the closure (server thread), a cross-stream tensor will trigger `shared_view` ->
        // `submit_blocking`, causing a fatal channel re-entrancy panic (`BorrowMutError`).
        let stream = tensor.stream;
        let tensor = tensor.into_ir();
        self.server
            .submit_blocking(move |server| {
                server.drain_stream(stream);
                server.resolve_server_float::<B>(&tensor)
            })
            .unwrap()
    }

    /// Resolve the given int tensor to a primitive tensor.
    pub fn resolve_tensor_int<B>(&self, tensor: FusionTensor<R>) -> B::IntTensorPrimitive
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        // Must evaluate `into_ir()` on the frontend thread. If evaluated inside
        // the closure (server thread), a cross-stream tensor will trigger `shared_view` ->
        // `submit_blocking`, causing a fatal channel re-entrancy panic (`BorrowMutError`).
        let stream = tensor.stream;
        let tensor = tensor.into_ir();
        self.server
            .submit_blocking(move |server| {
                server.drain_stream(stream);
                server.resolve_server_int::<B>(&tensor)
            })
            .unwrap()
    }

    /// Resolve the given bool tensor to a primitive tensor.
    pub fn resolve_tensor_bool<B>(&self, tensor: FusionTensor<R>) -> B::BoolTensorPrimitive
    where
        B: FusionBackend<FusionRuntime = R>,
    {
        // Must evaluate `into_ir()` on the frontend thread. If evaluated inside
        // the closure (server thread), a cross-stream tensor will trigger `shared_view` ->
        // `submit_blocking`, causing a fatal channel re-entrancy panic (`BorrowMutError`).
        let stream = tensor.stream;
        let tensor = tensor.into_ir();
        self.server
            .submit_blocking(move |server| {
                server.drain_stream(stream);
                server.resolve_server_bool::<B>(&tensor)
            })
            .unwrap()
    }

    /// Synchronize the collective operations.
    #[cfg(feature = "distributed")]
    pub fn sync_collective<B>(&self, device: &B::Device)
    where
        B: FusionBackend<FusionRuntime = R> + DistributedBackend,
    {
        println!("fusion sync_collective");
        // Ensure that all operations are resolved before calling sync_collective.
        self.sync(|| ());
        B::sync_collective(device)
    }

    /// Ensure that communication between the given devices is initialized.
    /// Initializing communication is generally blocking, so we make sure to flush those operations.
    #[cfg(feature = "distributed")]
    pub fn ensure_collective_init<B>(&self, device_ids: Vec<DeviceId>)
    where
        B: FusionBackend<FusionRuntime = R> + DistributedBackend,
    {
        println!("fusion ensure_init");
        println!("devices: {device_ids:?}");

        let utilities = self
            .server
            .utilities()
            .downcast::<FusionUtilities>()
            .expect("Can downcast to `FusionUtilities`");
        let id = CommunicationId::from(device_ids);
        println!("devices: {:?}", utilities.initialized_comms.read().unwrap());

        if !utilities.initialized_comms.read().unwrap().contains(&id) {
            println!("fusion ensure_init flush");

            // Communication initialization is blocking for the server, so we need to flush right away to make sure other devices
            // aren't waiting indefinitely on this initialization call.
            // This is already handled by cubecl, but since fusion adds another layer of streams and asynchronous submits,
            // we also needed to add some logic here to flush the fusion server.
            self.flush_queue();
            let mut initialized_comms = utilities.initialized_comms.write().unwrap();
            initialized_comms.insert(id);
        }
    }
}
