use std::sync::Arc;

use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use futures::stream::StreamExt;
use log::info;
use tokio::{
    sync::{Semaphore, mpsc, oneshot},
    time::{self, Duration, Instant},
};
use tokio_stream::wrappers::ReceiverStream;

use crate::utils::metrics::{
    BATCH_QUEUE_FLUSH_REASON_DRAIN, BATCH_QUEUE_FLUSH_REASON_MAX_LATENCY,
    BATCH_QUEUE_FLUSH_REASON_MAX_WEIGHT, BATCH_QUEUE_FLUSH_REASON_NEXT_OVERFLOW,
    BATCH_QUEUE_FLUSH_REASON_OVERSIZED_HEAD, BatchQueueMetrics,
};

// Helper function to create a vector of identical errors.
pub fn broadcast_error<T>(e: anyhow::Error, count: usize) -> Vec<Result<T>> {
    let err_msg = format!("{:?}", e);
    (0..count).map(|_| Err(anyhow!(err_msg.clone()))).collect()
}

// Trait for requests that have an associated weight used for batching.
//
// Types that will be processed by a [`BatchExecutor`] should implement this
// trait to describe their relative "cost" (e.g., size, resource usage, or
// complexity). The [`BatchQueue`] uses these weights to limit batch size based
// on the total weight of a batch rather than the simple count of requests.
pub trait Weighted {
    fn weight(&self) -> usize;
}

// Defines the contract for the actual logic that processes a batch of requests.
// This trait allows the BatchQueue to be generic over different execution
// strategies.
#[async_trait]
pub trait BatchExecutor: Send + Sync + 'static {
    type Request: Send + Weighted + 'static;
    type Response: Send + 'static;

    // Process a vector of requests and return a corresponding vector of responses.
    // The order of responses must match the order of requests.
    async fn execute_batch(&self, requests: Vec<Self::Request>) -> Vec<Result<Self::Response>>;
}

// Internal wrapper struct to couple a request with its specific return channel.
struct Job<R, T> {
    req: R,
    weight: u64,
    submit_started_at: Instant,
    enqueued_at: Instant,
    tx: oneshot::Sender<Result<T>>,
}

struct RequestTiming {
    submit_started_at: Instant,
    enqueued_at: Instant,
}

// A generic batching queue.
// It aggregates individual requests into batches based on weight or time
// (latency) limits and executes them concurrently with a limit.
pub struct BatchQueue<E: BatchExecutor> {
    sender: mpsc::Sender<Job<E::Request, E::Response>>,
    metrics: BatchQueueMetrics,
}

impl<E> BatchQueue<E>
where
    E: BatchExecutor,
{
    pub fn new(
        component: &'static str,
        executor: E,
        max_batch_weight: usize,
        max_latency: Duration,
        concurrency_limit: usize,
    ) -> Self {
        let (tx, rx) = mpsc::channel(1000);
        let executor = Arc::new(executor);
        let metrics = BatchQueueMetrics::new(component);
        let semaphore = Arc::new(Semaphore::new(concurrency_limit));

        let background_metrics = metrics;
        tokio::spawn(async move {
            let res = Self::background_task(
                rx,
                executor,
                semaphore,
                background_metrics,
                max_batch_weight as u64,
                max_latency,
            )
            .await;
            info!(
                "BatchQueue background_task exit, component={}, res={:?}",
                background_metrics.component(),
                res
            );
        });

        Self {
            sender: tx,
            metrics,
        }
    }

    pub async fn submit(&self, req: E::Request) -> Result<oneshot::Receiver<Result<E::Response>>> {
        let _submit_guard = self.metrics.begin_submit();
        let submit_started_at = Instant::now();
        let (resp_tx, resp_rx) = oneshot::channel();
        let weight = req.weight() as u64;
        let reserve_started_at = Instant::now();
        let permit = match self.sender.reserve().await {
            Ok(permit) => permit,
            Err(_) => {
                self.metrics
                    .observe_enqueue_block_duration(reserve_started_at.elapsed());
                self.metrics.record_submit_result(true);
                return Err(anyhow!("Failed to reserve slot in BatchQueue"));
            }
        };
        self.metrics
            .observe_enqueue_block_duration(reserve_started_at.elapsed());
        let enqueued_at = Instant::now();
        let job = Job {
            req,
            weight,
            submit_started_at,
            enqueued_at,
            tx: resp_tx,
        };
        self.metrics.add_backlog(1, weight);
        permit.send(job);
        self.metrics.record_submit_result(false);
        Ok(resp_rx)
    }

    async fn background_task(
        rx: mpsc::Receiver<Job<E::Request, E::Response>>,
        executor: Arc<E>,
        semaphore: Arc<Semaphore>,
        metrics: BatchQueueMetrics,
        max_batch_weight: u64,
        max_latency: Duration,
    ) -> Result<()> {
        // Convert the Receiver into a Peekable Stream.
        // This allows us to inspect the next item without consuming it.
        let stream = ReceiverStream::new(rx);
        let peekable_stream = stream.peekable();
        tokio::pin!(peekable_stream);

        loop {
            // Rate Limiting: Acquire a permit before collecting a new batch.
            // If we have reached `concurrency_limit`, this will block until a permit is
            // available. This ensures that we never exceed the concurrency limit.
            let permit = semaphore
                .clone()
                .acquire_owned()
                .await
                .context("Semaphore is closed")?;

            let mut buffer = Vec::new();
            let mut current_weight = 0u64;

            // Wait for the first job.
            let batch_collect_started_at = match peekable_stream.next().await {
                Some(job) => {
                    current_weight += job.weight;
                    buffer.push(job);
                    Instant::now()
                }
                None => break, // Channel closed, exit loop
            };

            // Batch Accumulation.
            // Try to collect more items until weight limit or timeout.
            let deadline = batch_collect_started_at + max_latency;
            let flush_reason = loop {
                if current_weight > max_batch_weight {
                    break BATCH_QUEUE_FLUSH_REASON_OVERSIZED_HEAD;
                }

                // If the batch is already full, stop collecting.
                if current_weight == max_batch_weight {
                    break BATCH_QUEUE_FLUSH_REASON_MAX_WEIGHT;
                }

                // Check the next item with a timeout.
                // We use `peek` here to decide whether to include the item.
                let peek_result = time::timeout_at(deadline, peekable_stream.as_mut().peek()).await;

                match peek_result {
                    Ok(Some(job_ref)) => {
                        // Pre-check: If adding the next job exceeds the weight limit, stop
                        // collecting. We leave this item in the stream for
                        // the next batch.
                        if current_weight.saturating_add(job_ref.weight) > max_batch_weight {
                            break BATCH_QUEUE_FLUSH_REASON_NEXT_OVERFLOW;
                        }

                        // Weight is okay, actually consume the item.
                        if let Some(job) = peekable_stream.next().await {
                            current_weight += job.weight;
                            buffer.push(job);
                            continue;
                        }

                        break BATCH_QUEUE_FLUSH_REASON_DRAIN;
                    }
                    Ok(None) => break BATCH_QUEUE_FLUSH_REASON_DRAIN,
                    Err(_) => break BATCH_QUEUE_FLUSH_REASON_MAX_LATENCY,
                }
            };

            // Dispatch Execution
            if !buffer.is_empty() {
                let mut reqs = Vec::with_capacity(buffer.len());
                let mut request_timings = Vec::with_capacity(buffer.len());
                let mut txs = Vec::with_capacity(buffer.len());
                for job in buffer.drain(..) {
                    request_timings.push(RequestTiming {
                        submit_started_at: job.submit_started_at,
                        enqueued_at: job.enqueued_at,
                    });
                    reqs.push(job.req);
                    txs.push(job.tx);
                }

                let executor_clone = executor.clone();
                let batch_size = reqs.len();
                let batch_weight = current_weight;
                let batch_metrics = metrics;

                // Spawn a new task for the actual execution.
                tokio::spawn(async move {
                    // Permit is dropped when this task finishes
                    let _permit = permit;
                    let batch_started_at = Instant::now();
                    batch_metrics.observe_batch_collect_duration(
                        batch_started_at.duration_since(batch_collect_started_at),
                    );
                    batch_metrics.observe_batch_size(batch_size);
                    batch_metrics.observe_batch_weight(batch_weight);
                    batch_metrics.record_batch_flush(flush_reason);
                    batch_metrics.record_merged_requests(batch_size.saturating_sub(1));
                    batch_metrics.sub_backlog(batch_size, batch_weight);
                    for timing in &request_timings {
                        batch_metrics.observe_request_queue_duration(
                            batch_started_at.duration_since(timing.enqueued_at),
                        );
                    }
                    let _execution_guard = batch_metrics.begin_batch_execution(batch_size);
                    let results = executor_clone.execute_batch(reqs).await;
                    let completed_at = Instant::now();
                    if results.len() != txs.len() {
                        let err = anyhow!(
                            "BatchExecutor returned {} results, expected {}",
                            results.len(),
                            txs.len()
                        );
                        let error_responses = broadcast_error::<E::Response>(err, txs.len());
                        for ((tx, timing), res) in txs
                            .into_iter()
                            .zip(request_timings.into_iter())
                            .zip(error_responses)
                        {
                            batch_metrics.observe_request_completion(
                                completed_at.duration_since(timing.submit_started_at),
                                true,
                            );
                            let _ = tx.send(res);
                        }
                        return;
                    }
                    // Map results back to their specific callers.
                    for ((tx, timing), res) in txs
                        .into_iter()
                        .zip(request_timings.into_iter())
                        .zip(results)
                    {
                        batch_metrics.observe_request_completion(
                            completed_at.duration_since(timing.submit_started_at),
                            res.is_err(),
                        );
                        let _ = tx.send(res);
                    }
                });
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use tokio::{sync::Notify, time::sleep};

    use super::*;

    const TEST_COMPONENT: &str = "ut_batch_queue";

    impl Weighted for i32 {
        fn weight(&self) -> usize {
            1
        }
    }

    #[derive(Clone)]
    struct MockExecutor {
        pub calls: Arc<Mutex<Vec<Vec<i32>>>>,
        pub notify: Arc<Notify>,
        pub wait_for_signal: bool,
    }

    impl MockExecutor {
        fn new() -> Self {
            Self {
                calls: Arc::new(Mutex::new(Vec::new())),
                notify: Arc::new(Notify::new()),
                wait_for_signal: false,
            }
        }

        fn new_blocking() -> Self {
            Self {
                calls: Arc::new(Mutex::new(Vec::new())),
                notify: Arc::new(Notify::new()),
                wait_for_signal: true,
            }
        }
    }

    #[async_trait]
    impl BatchExecutor for MockExecutor {
        type Request = i32;
        type Response = String;

        async fn execute_batch(&self, requests: Vec<Self::Request>) -> Vec<Result<Self::Response>> {
            {
                let mut calls = self.calls.lock().unwrap();
                calls.push(requests.clone());
            }

            if self.wait_for_signal {
                self.notify.notified().await;
            }

            requests
                .into_iter()
                .map(|i| Ok(format!("res_{}", i)))
                .collect()
        }
    }

    #[derive(Clone, Debug, PartialEq)]
    struct HeavyRequest {
        id: i32,
        w: usize,
    }

    impl Weighted for HeavyRequest {
        fn weight(&self) -> usize {
            self.w
        }
    }

    #[derive(Clone)]
    struct HeavyMockExecutor {
        pub calls: Arc<Mutex<Vec<Vec<HeavyRequest>>>>,
    }

    impl HeavyMockExecutor {
        fn new() -> Self {
            Self {
                calls: Arc::new(Mutex::new(Vec::new())),
            }
        }
    }

    #[async_trait]
    impl BatchExecutor for HeavyMockExecutor {
        type Request = HeavyRequest;
        type Response = String;

        async fn execute_batch(&self, requests: Vec<Self::Request>) -> Vec<Result<Self::Response>> {
            {
                let mut calls = self.calls.lock().unwrap();
                calls.push(requests.clone());
            }
            requests
                .into_iter()
                .map(|r| Ok(format!("id_{}", r.id)))
                .collect()
        }
    }

    #[tokio::test]
    async fn test_basic_flow() -> Result<()> {
        let executor = MockExecutor::new();
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            10,
            Duration::from_millis(100),
            1,
        );

        let rx1 = queue.submit(1).await?;
        let rx2 = queue.submit(2).await?;

        let res1 = rx1.await??;
        let res2 = rx2.await??;

        assert_eq!(res1, "res_1");
        assert_eq!(res2, "res_2");

        let calls = executor.calls.lock().unwrap();
        assert!(!calls.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_batch_size_trigger() -> Result<()> {
        let executor = MockExecutor::new();
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            3,
            Duration::from_secs(60),
            1,
        );

        let rx1 = queue.submit(1).await?;
        let rx2 = queue.submit(2).await?;

        {
            let calls = executor.calls.lock().unwrap();
            assert_eq!(calls.len(), 0);
        }

        let rx3 = queue.submit(3).await?;

        let _ = rx3.await?;
        let _ = rx1.await?;
        let _ = rx2.await?;

        let calls = executor.calls.lock().unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0], vec![1, 2, 3]);

        Ok(())
    }

    #[tokio::test]
    async fn test_latency_trigger() -> Result<()> {
        let executor = MockExecutor::new();
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            10,
            Duration::from_millis(50),
            1,
        );

        let rx = queue.submit(1).await?;
        {
            let calls = executor.calls.lock().unwrap();
            assert_eq!(calls.len(), 0);
        }

        let res = rx.await??;
        assert_eq!(res, "res_1");

        let calls = executor.calls.lock().unwrap();
        assert_eq!(calls.len(), 1);

        Ok(())
    }

    #[tokio::test]
    async fn test_latency_starts_after_first_job_arrives() -> Result<()> {
        let executor = MockExecutor::new();
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            10,
            Duration::from_millis(200),
            1,
        );

        sleep(Duration::from_millis(300)).await;

        let rx1 = queue.submit(1).await?;
        sleep(Duration::from_millis(50)).await;

        let rx2 = queue.submit(2).await?;

        let res1 = rx1.await??;
        let res2 = rx2.await??;

        assert_eq!(res1, "res_1");
        assert_eq!(res2, "res_2");

        let calls = executor.calls.lock().unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0], vec![1, 2]);

        Ok(())
    }

    #[tokio::test]
    async fn test_backpressure_and_accumulation() -> Result<()> {
        let executor = MockExecutor::new_blocking();
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            5,
            Duration::from_millis(10),
            1,
        );

        let rx1 = queue.submit(100).await?;

        sleep(Duration::from_millis(100)).await;

        let mut rxs = vec![];
        for i in 1..=6 {
            rxs.push(queue.submit(i).await?);
        }

        executor.notify.notify_one();
        let _ = rx1.await?;

        sleep(Duration::from_millis(100)).await;
        executor.notify.notify_one();

        sleep(Duration::from_millis(100)).await;
        executor.notify.notify_one();

        for rx in rxs {
            let _ = rx.await?;
        }

        let calls = executor.calls.lock().unwrap();
        assert_eq!(calls.len(), 3);
        assert_eq!(calls[0], vec![100]);
        assert_eq!(calls[1], vec![1, 2, 3, 4, 5]);
        assert_eq!(calls[2], vec![6]);

        Ok(())
    }

    #[tokio::test]
    async fn test_fire_and_forget() -> Result<()> {
        let executor = MockExecutor::new();
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            10,
            Duration::from_millis(10),
            1,
        );

        // Submit a task but don't await the result (fire-and-forget)
        let rx = queue.submit(999).await?;
        drop(rx);

        // Wait for the task to be processed
        sleep(Duration::from_millis(100)).await;

        let calls = executor.calls.lock().unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0], vec![999]);

        Ok(())
    }

    #[tokio::test]
    async fn test_queue_closed_error() -> Result<()> {
        let executor = MockExecutor::new();
        let queue = BatchQueue::new(TEST_COMPONENT, executor, 10, Duration::from_millis(10), 1);
        let rx_rsp = queue.submit(1).await?;

        // Drop queue
        drop(queue);

        // If the queue is dropped, the rx may receive an error or a success value,
        // it depends on whether the task is processed before the queue is dropped.
        let _ = rx_rsp.await;
        Ok(())
    }

    #[tokio::test]
    async fn test_weight_overflow_splits_batch() -> Result<()> {
        let executor = HeavyMockExecutor::new();
        // Set max weight to 10
        // Set a long latency to ensure batch split is triggered by weight, not timeout
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            10,
            Duration::from_secs(10),
            1,
        );

        // Task 1: Weight 6 (Current accumulated: 6/10)
        let rx1 = queue.submit(HeavyRequest { id: 1, w: 6 }).await?;

        // Task 2: Weight 5 (If added, accumulated: 11/10 -> Overflow)
        // Expectation: Task 2 should stay in queue and be part of the next batch
        let rx2 = queue.submit(HeavyRequest { id: 2, w: 5 }).await?;

        // Task 3: Weight 2 (Should follow Task 2 in the second batch)
        let rx3 = queue.submit(HeavyRequest { id: 3, w: 2 }).await?;

        // Wait for results
        let _ = rx1.await?;
        let _ = rx2.await?;
        let _ = rx3.await?;

        let calls = executor.calls.lock().unwrap();

        // We expect execution in 2 batches:
        assert_eq!(
            calls.len(),
            2,
            "Should be split into 2 batches due to weight limit"
        );

        // Batch 1: Only id 1 (Weight 6)
        assert_eq!(calls[0].len(), 1);
        assert_eq!(calls[0][0].id, 1);
        assert_eq!(calls[0][0].w, 6);

        // Batch 2: id 2 (Weight 5) + id 3 (Weight 2) = Total 7 < 10
        assert_eq!(calls[1].len(), 2);
        assert_eq!(calls[1][0].id, 2);
        assert_eq!(calls[1][1].id, 3);

        Ok(())
    }

    #[tokio::test]
    async fn test_exact_weight_fit() -> Result<()> {
        let executor = HeavyMockExecutor::new();
        // Max weight 10
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            10,
            Duration::from_secs(10),
            1,
        );

        // Task 1: Weight 5
        let rx1 = queue.submit(HeavyRequest { id: 1, w: 5 }).await?;
        // Task 2: Weight 5 (Total 10, fits exactly)
        let rx2 = queue.submit(HeavyRequest { id: 2, w: 5 }).await?;
        // Task 3: Weight 1 (Total 11, moves to next batch)
        let rx3 = queue.submit(HeavyRequest { id: 3, w: 1 }).await?;

        let _ = rx1.await?;
        let _ = rx2.await?;
        let _ = rx3.await?;

        let calls = executor.calls.lock().unwrap();
        assert_eq!(calls.len(), 2);

        // Batch 1 contains id 1 and id 2
        assert_eq!(calls[0].len(), 2);
        assert_eq!(calls[0][0].id, 1);
        assert_eq!(calls[0][1].id, 2);

        // Batch 2 contains id 3
        assert_eq!(calls[1].len(), 1);
        assert_eq!(calls[1][0].id, 3);

        Ok(())
    }

    #[tokio::test]
    async fn test_single_giant_item() -> Result<()> {
        // Test logic: If the queue is empty, the first task must be processed
        // even if it exceeds max_batch_weight, to prevent deadlocks.
        let executor = HeavyMockExecutor::new();
        // Max weight 5
        let queue = BatchQueue::new(
            TEST_COMPONENT,
            executor.clone(),
            5,
            Duration::from_secs(10),
            1,
        );

        // Submit a task with weight 10 ( > max_batch_weight )
        let rx = queue.submit(HeavyRequest { id: 99, w: 10 }).await?;

        // It should succeed and not hang
        let res = rx.await??;
        assert_eq!(res, "id_99");

        let calls = executor.calls.lock().unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0][0].w, 10);

        Ok(())
    }
}