use std::{
    fmt::Debug,
    sync::{Arc, atomic::AtomicU16},
};

use anyhow::Result;
use log::{info, warn};
use tici_grpc::{
    S3Location, ShardManifestHeader, WorkerNodeHeartbeatRequest, WorkerNodeHeartbeatResponse,
    WorkerNodeStatus,
};
use tici_import::ImportTask;
use tokio::sync::mpsc::UnboundedSender;

use crate::{
    ShardInfoUpdate,
    meta::{
        MetaStore,
        writer::{
            CompactionState, CompactionStateKey, FTSIndex, ShardBuildProgress, ShardStats,
            ShardTask, WorkerNodeMgr,
        },
    },
};

pub struct WorkerNodeTaskCtx<T: Send + Sync> {
    pub addr: String,
    // TODO(P2) make these callbacks return directly the grpc response, can use closures to do
    // that.
    callback: tokio::sync::oneshot::Sender<Result<T>>,
}

impl<T: Send + Sync> Debug for WorkerNodeTaskCtx<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WorkerNodeTaskCtx")
            .field("addr", &self.addr)
            .finish()
    }
}

impl<T: Send + Sync> WorkerNodeTaskCtx<T> {
    pub fn new(addr: String, callback: tokio::sync::oneshot::Sender<Result<T>>) -> Self {
        Self { addr, callback }
    }

    pub fn create(addr: String) -> (Self, tokio::sync::oneshot::Receiver<Result<T>>) {
        let (callback, receiver) = tokio::sync::oneshot::channel();
        (Self::new(addr, callback), receiver)
    }

    pub fn report_error(self, err: anyhow::Error) {
        info!("report from wn addr={} {:?}", self.addr, err);
        let _ = self.callback.send(Err(err));
    }

    pub fn report_ok(self, t: T) {
        let _ = self.callback.send(Ok(t));
    }
}

#[derive(Debug)]

pub struct ModifyShardTaskAdd {
    pub header: ShardManifestHeader,
    pub progress: ShardBuildProgress,
    pub upload_location: S3Location,
    pub shard_stats: ShardStats,
    pub index: Arc<FTSIndex>,
    pub notified: bool,
}

#[derive(Debug)]

pub struct ModifyShardTaskRemove {
    pub shard_id: u64,
}

#[derive(Debug)]
pub struct ModifyShardTask {
    pub ctx: WorkerNodeTaskCtx<()>,
    pub to_add: Vec<ModifyShardTaskAdd>,
    pub to_remove: Vec<ModifyShardTaskRemove>,
}

#[derive(Debug)]
pub struct AddCompactionTask {
    pub ctx: WorkerNodeTaskCtx<()>,
    pub header: ShardManifestHeader,
    pub key: CompactionStateKey,
    pub work: Arc<CompactionWork>,
}

pub struct CreateWorkerNodeTask {
    pub ctx: WorkerNodeTaskCtx<()>,
    pub status: Option<WorkerNodeStatus>,
    pub meta_store: Arc<dyn MetaStore>,
}

impl Debug for CreateWorkerNodeTask {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CreateWorkerNodeTask")
            .field("ctx", &self.ctx)
            .field("status", &self.status)
            .finish()
    }
}

#[derive(Clone, Copy, Debug)]
pub struct HeartbeatTaskPagination {
    pub response_size_limit: usize,
}

#[derive(Debug)]
pub struct HeartbeatTask {
    pub ctx: WorkerNodeTaskCtx<Vec<WorkerNodeHeartbeatResponse>>,
    pub heartbeat: WorkerNodeHeartbeatRequest,
    pub shard_task_tx: UnboundedSender<ShardTask>,
    pub pagination: HeartbeatTaskPagination,
}

#[derive(Debug)]
pub struct UpdateTask {
    pub ctx: WorkerNodeTaskCtx<()>,
    pub update: ShardInfoUpdate,
}

#[allow(clippy::type_complexity)]
#[derive(Debug)]
pub enum WorkerNodeTask {
    Hup,
    // Scheduled from grpc server, with a callback.
    Heartbeat(HeartbeatTask),
    Update(UpdateTask),
    // add/remove shard
    ModifyShard(ModifyShardTask),
    AddCompaction(AddCompactionTask),
    CreateWorkerNode(CreateWorkerNodeTask),
    // Scheduled with no callback.
    AddImportFragBuildTask(ImportTask),
}

impl WorkerNodeTask {
    pub fn report_error(self, err: anyhow::Error) {
        match self {
            WorkerNodeTask::Heartbeat(HeartbeatTask { ctx, .. }) => {
                ctx.report_error(err);
            }
            WorkerNodeTask::Update(UpdateTask { ctx, .. }) => {
                ctx.report_error(err);
            }
            WorkerNodeTask::ModifyShard(ModifyShardTask { ctx, .. }) => {
                ctx.report_error(err);
            }
            WorkerNodeTask::AddCompaction(AddCompactionTask { ctx, .. }) => {
                ctx.report_error(err);
            }
            WorkerNodeTask::CreateWorkerNode(CreateWorkerNodeTask { ctx, .. }) => {
                ctx.report_error(err);
            }
            WorkerNodeTask::AddImportFragBuildTask(..) => {
                // TODO(P1) Use callback
            }
            WorkerNodeTask::Hup => {
                // do nothing
            }
        }
    }
}

#[derive(Debug)]
pub struct CompactionWork {
    pub state: CompactionState,
    pub tick: AtomicU16,
}

#[derive(Debug, Clone)]
pub struct ShardSchedule {
    pub shard_header: ShardManifestHeader,
    pub latest_progress: ShardBuildProgress,
    pub upload_location: S3Location,
    pub shard_stats: ShardStats,
    index: Arc<FTSIndex>,
}

impl ShardSchedule {
    pub fn new(
        shard_header: ShardManifestHeader,
        latest_progress: ShardBuildProgress,
        upload_location: S3Location,
        shard_stats: ShardStats,
        index: Arc<FTSIndex>,
    ) -> Self {
        Self {
            shard_header,
            latest_progress,
            upload_location,
            shard_stats,
            index,
        }
    }

    pub fn get_shard_header(&self) -> &ShardManifestHeader {
        &self.shard_header
    }

    pub fn get_index(&self) -> &FTSIndex {
        self.index.as_ref()
    }

    pub fn index_cloned(&self) -> Arc<FTSIndex> {
        self.index.clone()
    }
}

impl WorkerNodeMgr {
    // Worker nodes are registered by:
    // - heartbeat
    // - restored messages
    // TODO(P2) Migrate if the restored worker node is no longer available.
    pub async fn send_modify_shard(
        wn_task_tx: UnboundedSender<WorkerNodeTask>,
        addr: String,
        to_add: Vec<ModifyShardTaskAdd>,
        to_remove: Vec<ModifyShardTaskRemove>,
    ) -> Result<()> {
        let (tx, rx) = tokio::sync::oneshot::channel();
        if let Err(e) = wn_task_tx.send(WorkerNodeTask::ModifyShard(ModifyShardTask {
            ctx: WorkerNodeTaskCtx::new(addr.clone(), tx),
            to_add,
            to_remove,
        })) {
            warn!(
                "can not send send_modify_shard, addr={}, task={:?}",
                addr, e.0
            );
        }
        let x = rx.await?;
        x?;
        Ok(())
    }

    pub async fn send_add_compaction(
        wn_task_tx: UnboundedSender<WorkerNodeTask>,
        addr: String,
        header: ShardManifestHeader,
        key: CompactionStateKey,
        compaction_work: Arc<CompactionWork>,
    ) -> Result<()> {
        let (tx, rx) = tokio::sync::oneshot::channel();
        if let Err(e) = wn_task_tx.send(WorkerNodeTask::AddCompaction(AddCompactionTask {
            ctx: WorkerNodeTaskCtx::new(addr.clone(), tx),
            header,
            key,
            work: compaction_work,
        })) {
            warn!(
                "can not send send_add_compaction, addr={}, task={:?}",
                addr, e.0
            );
        }
        let x = rx.await?;
        x?;
        Ok(())
    }

    pub async fn send_create_worker_node(
        wn_task_tx: UnboundedSender<WorkerNodeTask>,
        addr: String,
        status: Option<WorkerNodeStatus>,
        meta_store: Arc<dyn MetaStore>,
    ) -> Result<()> {
        let (tx, rx) = tokio::sync::oneshot::channel();
        if let Err(e) = wn_task_tx.send(WorkerNodeTask::CreateWorkerNode(CreateWorkerNodeTask {
            ctx: WorkerNodeTaskCtx::new(addr.clone(), tx),
            status,
            meta_store,
        })) {
            warn!(
                "can not send send_create_worker_node, addr={}, task={:?}",
                addr, e.0
            );
        }
        let x = rx.await?;
        x?;
        Ok(())
    }
}