use std::{fmt::Debug, sync::Arc, time::Instant};

use anyhow::Result;
use tici_common::index::IndexKey;
use tici_grpc::{
    AppendDeltaFragMetaRequest, AppendFragMetaResponse, ErrorCode, FinishCompactFragResponse,
    FragMeta, ShardManifestHeader,
};

use crate::{
    RangeBound, ShardInfo, ShardInfoUpdate,
    meta::{
        metrics::SHARD_TASK_DURATION_HISTOGRAM,
        writer::{CompactionState, CompactionStateKey, FTSIndex, ShardSchedule},
    },
};

#[derive(Debug, Clone)]
pub struct ShardTaskHeader {
    pub index_key: IndexKey,
    pub start_key: RangeBound,
    pub shard_id: u64,
}

impl From<&CompactionState> for ShardTaskHeader {
    fn from(state: &CompactionState) -> Self {
        ShardTaskHeader {
            index_key: state.index_key.clone(),
            start_key: state.range.sortable_start(),
            shard_id: state.shard_id,
        }
    }
}

impl From<&ShardInfo> for ShardTaskHeader {
    fn from(info: &ShardInfo) -> Self {
        ShardTaskHeader {
            index_key: info.index_key.clone(),
            start_key: info.range.sortable_start(),
            shard_id: info.shard_id,
        }
    }
}

impl From<&ShardSchedule> for ShardTaskHeader {
    fn from(schedule: &ShardSchedule) -> Self {
        ShardTaskHeader {
            index_key: schedule.get_index().index_key(),
            start_key: RangeBound::make_start(schedule.get_shard_header().start_key.clone()),
            shard_id: schedule.get_shard_header().shard_id,
        }
    }
}

impl ShardTaskHeader {
    pub fn new(index_key: IndexKey, start_key: RangeBound, shard_id: u64) -> Self {
        Self {
            index_key,
            start_key,
            shard_id,
        }
    }

    pub fn to_long_ident(&self) -> String {
        format!(
            "({:?}, shard_id={}, start_key={:?})",
            self.index_key, self.shard_id, self.start_key
        )
    }
}

pub struct ShardTaskCtx<T: Send + Sync> {
    callback: Option<tokio::sync::oneshot::Sender<Result<T>>>,
    create_at: Instant,
    metric_label: Option<String>,
}

impl<T: Send + Sync> Debug for ShardTaskCtx<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("STC")
            .field("has_callback", &self.callback.is_some())
            .field("create_at", &self.create_at)
            .field("elapsed", &self.create_at.elapsed().as_secs_f32())
            .finish()
    }
}

impl<T: Send + Sync> ShardTaskCtx<T> {
    pub fn new_no_cb(metric_label: Option<String>) -> Self {
        Self {
            callback: None,
            create_at: Instant::now(),
            metric_label,
        }
    }

    pub fn new(
        metric_label: Option<String>,
        callback: tokio::sync::oneshot::Sender<Result<T>>,
    ) -> Self {
        Self {
            callback: Some(callback),
            create_at: Instant::now(),
            metric_label,
        }
    }

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

    pub fn queued_secs(&self) -> f64 {
        self.create_at.elapsed().as_secs_f64()
    }

    pub fn metric_label(&self) -> Option<&str> {
        self.metric_label.as_deref()
    }

    pub fn report_ok(self, r: T) {
        if let Some(s) = self.callback {
            let _ = s.send(Ok(r));
        }
        if let Some(s) = self.metric_label {
            let e = Instant::now().duration_since(self.create_at).as_millis() as f64;
            SHARD_TASK_DURATION_HISTOGRAM
                .with_label_values(&[s.as_str()])
                .observe(e);
        }
    }

    pub fn report_err(self, r: anyhow::Error) {
        if let Some(s) = self.callback {
            let _ = s.send(Err(r));
        }
    }
}

#[derive(Debug)]
pub struct AppendDeltaFragTask {
    pub ctx: ShardTaskCtx<(AppendFragMetaResponse, Option<ShardInfoUpdate>)>,
    pub header: ShardTaskHeader,
    pub index: Arc<FTSIndex>,
    pub req: AppendDeltaFragMetaRequest,
}

#[derive(Debug)]

pub struct AppendBaseFragTask {
    pub ctx: ShardTaskCtx<(AppendFragMetaResponse, Option<ShardInfoUpdate>)>,
    pub header: ShardTaskHeader,
    pub index: Arc<FTSIndex>,
    pub frags: Vec<FragMeta>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExplicitSplitReason {
    PresplitImport,
}

impl ExplicitSplitReason {
    pub fn to_metrics_label(&self) -> &str {
        match self {
            ExplicitSplitReason::PresplitImport => "split_presplit_import",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ExplicitSplitResult {
    pub split_count: u32,
}

#[derive(Debug)]
pub struct ExplicitSplitTask {
    pub ctx: ShardTaskCtx<ExplicitSplitResult>,
    pub header: ShardTaskHeader,
    pub index: Arc<FTSIndex>,
    pub split_keys: Vec<Vec<u8>>,
    pub reason: ExplicitSplitReason,
}

#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum CompactShardReason {
    Uni,
    UniChain,
    MergeOverlap,
    Manual,
    // If there are too many fragments, commit append will fail.
    FragCount,
}

impl CompactShardReason {
    pub fn to_metrics_label(&self) -> &str {
        match self {
            CompactShardReason::Uni => "compact_uni",
            CompactShardReason::UniChain => "compact_uni_chain",
            CompactShardReason::MergeOverlap => "compact_merge_overlap",
            CompactShardReason::Manual => "compact_manual",
            CompactShardReason::FragCount => "compact_fragcount",
        }
    }
}

#[derive(Debug)]

pub struct CompactionFinishTask {
    pub ctx: ShardTaskCtx<FinishCompactFragResponse>,
    pub header: ShardManifestHeader,
    pub index: Arc<FTSIndex>,
    pub compaction_key: CompactionStateKey,
    pub compaction_state: CompactionState,
    pub output_frags: Vec<FragMeta>,
}

#[derive(Debug)]
pub struct CompactionAbortTask {
    pub ctx: ShardTaskCtx<()>,
    pub header: ShardTaskHeader,
    pub key: CompactionStateKey,
}

#[derive(Debug)]
pub enum MergeShardReason {
    Auto,
    Manual,
}

impl MergeShardReason {
    pub fn to_metrics_label(&self) -> &str {
        match self {
            MergeShardReason::Auto => "merge_auto",
            MergeShardReason::Manual => "merge_manual",
        }
    }
}

#[derive(Debug)]

pub struct MergeTask {
    pub ctx: ShardTaskCtx<()>,
    pub index: Arc<FTSIndex>,
    pub start_key: RangeBound,
    pub shard_id: u64,
    pub reason: MergeShardReason,
}

#[derive(Debug)]
pub enum MaybeSplitOrCompactReason {
    AfterAppend,
    AfterCompaction {
        trigger_frag_path: String,
        remaining_iterations: usize,
    },
    ManualSplit,
    StallTooManyFrags,
}

#[derive(Debug)]

pub struct MaybeSplitOrCompactTask {
    pub ctx: ShardTaskCtx<()>,
    pub header: ShardTaskHeader,
    pub reason: MaybeSplitOrCompactReason,
}

#[derive(Debug, Clone)]
pub enum RescheduleShardReason {
    Manual,
    Dead,
    ScaleIn,
    Blocked,
    ProgressUnwinded,
    // TODO(P2) Support auto balance by size or load
}

impl RescheduleShardReason {
    pub fn to_metrics_label(&self) -> &str {
        match self {
            RescheduleShardReason::Manual => "resche_manual",
            RescheduleShardReason::Dead => "resche_dead",
            RescheduleShardReason::ScaleIn => "resche_scale_in",
            RescheduleShardReason::Blocked => "blocked_worker",
            RescheduleShardReason::ProgressUnwinded => "progress_unwinded",
        }
    }
}

#[derive(Debug)]
pub struct NeedRescheduleTask {
    pub ctx: ShardTaskCtx<()>,
    pub header: ShardTaskHeader,
    pub designated_writer: Option<String>,
    pub reason: RescheduleShardReason,
}

#[derive(Debug)]

pub struct ManualCompactionTask {
    pub ctx: ShardTaskCtx<ErrorCode>,
    pub header: ShardTaskHeader,
    pub range: String,
}

// TODO(P2) improve the codes if possible
// - Make `ctx`` a common field
// - Remove the `metric_label` passed to `ShardTaskCtx`

#[allow(clippy::type_complexity)]
#[derive(Debug)]
pub enum ShardTask {
    Hup,
    // Scheduled from WNMgr, with a callback
    CompactionFinish(CompactionFinishTask),
    AppendDeltaFrag(AppendDeltaFragTask),
    AppendBaseFrag(AppendBaseFragTask),
    ExplicitSplit(ExplicitSplitTask),
    // Scheduled without callback
    MaybeSplitOrCompact(MaybeSplitOrCompactTask),
    NeedReschedule(NeedRescheduleTask),
    CompactionAborted(CompactionAbortTask),
    TryMerge(MergeTask),
    // Scheduled by operators, with a callback
    ManualCompaction(ManualCompactionTask),
}

impl ShardTask {
    pub fn get_local_id(&self, n: usize) -> Option<usize> {
        match &self {
            // Hup message must be handled by self.
            ShardTask::Hup => None,
            ShardTask::CompactionFinish(c) => Some(c.header.shard_id as usize % n),
            ShardTask::AppendDeltaFrag(a) => Some(a.header.shard_id as usize % n),
            ShardTask::AppendBaseFrag(a) => Some(a.header.shard_id as usize % n),
            ShardTask::ExplicitSplit(t) => Some(t.header.shard_id as usize % n),
            ShardTask::MaybeSplitOrCompact(t) => Some(t.header.shard_id as usize % n),
            ShardTask::NeedReschedule(t) => Some(t.header.shard_id as usize % n),
            ShardTask::CompactionAborted(c) => Some(c.header.shard_id as usize % n),
            ShardTask::ManualCompaction(c) => Some(c.header.shard_id as usize % n),
            ShardTask::TryMerge(m) => Some(m.shard_id as usize % n),
        }
    }
}