use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use bytes::Bytes;
use folk_api::{BoxFuture, Executor, PluginContext, RpcMethodDef, ServerPlugin};
use tracing::{error, info, warn};

use crate::config::{DriverKind, JobsConfig, QueueConfig};
use crate::driver::Driver;

pub struct JobsPlugin {
    config: JobsConfig,
}

impl JobsPlugin {
    pub fn new(config: JobsConfig) -> Self {
        Self { config }
    }
}

#[async_trait]
impl ServerPlugin for JobsPlugin {
    fn name(&self) -> &'static str {
        "jobs"
    }

    async fn run(&self, ctx: PluginContext) -> Result<()> {
        let driver: Arc<dyn Driver> = match self.config.driver {
            DriverKind::Memory => crate::driver::MemoryDriver::new(),
            DriverKind::Redis => {
                crate::redis_driver::RedisDriver::new(&self.config.redis_url).await?
            }
        };

        // One consumer task per queue × concurrency
        let mut tasks = Vec::new();
        for q in &self.config.queues {
            for _ in 0..q.concurrency {
                let d = driver.clone();
                let exec = ctx.executor.clone();
                let q_cfg = q.clone();
                let sd = ctx.shutdown.clone();
                tasks.push(tokio::spawn(async move {
                    consume_loop(d, exec, q_cfg, sd).await;
                }));
            }
        }

        info!(
            queues = self.config.queues.len(),
            driver = ?self.config.driver,
            "jobs plugin started"
        );

        // Register RPC stats
        if let Some(reg) = &ctx.rpc_registrar {
            let d = driver.clone();
            let queues: Vec<String> = self.config.queues.iter().map(|q| q.name.clone()).collect();
            reg.register_raw(
                "jobs.stats".into(),
                Arc::new(move |_: Bytes| -> BoxFuture<'static, Result<Bytes>> {
                    let d = d.clone();
                    let queues = queues.clone();
                    Box::pin(async move {
                        let mut stats = vec![];
                        for q in &queues {
                            let depth = d.depth(q).await.unwrap_or(0);
                            stats.push(format!("{q}: depth={depth}"));
                        }
                        Ok(Bytes::from(rmp_serde::to_vec(&stats)?))
                    })
                }),
            )
            .await;
        }

        let mut sd = ctx.shutdown.clone();
        sd.changed().await.ok();
        for t in tasks {
            let _ = t.await;
        }
        Ok(())
    }

    fn rpc_methods(&self) -> Vec<RpcMethodDef> {
        vec![RpcMethodDef::new(
            "jobs.stats",
            "queue depth and processing counters",
        )]
    }
}

async fn consume_loop(
    driver: Arc<dyn Driver>,
    executor: Arc<dyn Executor>,
    queue: QueueConfig,
    mut shutdown: tokio::sync::watch::Receiver<bool>,
) {
    loop {
        let job = tokio::select! {
            job = driver.pop(&queue.name) => job,
            _ = shutdown.changed() => {
                if *shutdown.borrow() { return; }
                continue;
            }
        };

        let payload = match job {
            Ok(Some(p)) => p,
            Ok(None) => continue,
            Err(e) => {
                error!(error = ?e, queue = %queue.name, "driver pop error");
                continue;
            }
        };

        let mut retries = 0u32;
        loop {
            match executor.execute_method("jobs.process", payload.clone()).await {
                Ok(_) => break,
                Err(e) => {
                    retries += 1;
                    if retries >= queue.max_retries {
                        error!(
                            queue = %queue.name,
                            retries,
                            "job failed after max retries; discarding"
                        );
                        break;
                    }
                    warn!(
                        queue = %queue.name,
                        retries,
                        error = ?e,
                        "job failed; retrying"
                    );
                    tokio::time::sleep(Duration::from_secs(1 << retries.min(5))).await;
                }
            }
        }
    }
}
