// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{
    cell::Cell,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

use pin_project::pin_project;
use trace_event::types::TraceContext;

use crate::{GLOBAL_TRACKERS, INVALID_TRACKER_TOKEN, Tracker, slab::TrackerToken};

thread_local! {
    static TLS_TRACKER_TOKEN: Cell<TrackerToken> = const { Cell::new(INVALID_TRACKER_TOKEN) };
    // TODO: This might be better to be moved to trace_event module and properly referenced by the
    //   `TrackedFuture`.
    static TLS_TRACE_CTX: Cell<TraceContext> = Cell::new(TraceContext::default());
}

pub fn set_tls_tracker_token(token: TrackerToken) {
    TLS_TRACKER_TOKEN.with(|c| {
        c.set(token);
    })
}

pub fn set_tls_trace_ctx(v: TraceContext) {
    TLS_TRACE_CTX.with(|c| c.set(v))
}

pub fn clear_tls_tracker_token() {
    set_tls_tracker_token(INVALID_TRACKER_TOKEN);
}

pub fn clear_tls_trace_ctx() {
    set_tls_trace_ctx(TraceContext::default());
}

pub fn get_tls_tracker_token() -> TrackerToken {
    TLS_TRACKER_TOKEN.with(|c| c.get())
}

pub fn get_tls_trace_ctx() -> TraceContext {
    TLS_TRACE_CTX.with(|c| {
        let v = c.take();
        let ret = v.clone();
        c.set(v);
        ret
    })
}

pub fn with_tls_tracker<F>(mut f: F)
where
    F: FnMut(&mut Tracker),
{
    TLS_TRACKER_TOKEN.with(|c| {
        GLOBAL_TRACKERS.with_tracker(c.get(), &mut f);
    });
}

struct TlsContextGuard {
    prev_tracker: TrackerToken,
    prev_trace_ctx: TraceContext,
}

impl TlsContextGuard {
    fn install(tracker: TrackerToken, trace_ctx: TraceContext) -> Self {
        let prev_tracker = get_tls_tracker_token();
        let prev_trace_ctx = get_tls_trace_ctx();
        set_tls_tracker_token(tracker);
        set_tls_trace_ctx(trace_ctx);
        Self {
            prev_tracker,
            prev_trace_ctx,
        }
    }
}

impl Drop for TlsContextGuard {
    fn drop(&mut self) {
        set_tls_tracker_token(self.prev_tracker);
        set_tls_trace_ctx(std::mem::take(&mut self.prev_trace_ctx));
    }
}

#[pin_project]
pub struct TrackedFuture<F> {
    #[pin]
    future: F,
    tracker: TrackerToken,
    trace_ctx: TraceContext,
}

impl<F> TrackedFuture<F> {
    pub fn new(future: F) -> TrackedFuture<F> {
        TrackedFuture {
            future,
            tracker: get_tls_tracker_token(),
            trace_ctx: get_tls_trace_ctx(),
        }
    }
}

impl<F: Future> Future for TrackedFuture<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let _tls_guard = TlsContextGuard::install(*this.tracker, this.trace_ctx.clone());
        this.future.poll(cx)
    }
}

#[cfg(test)]
mod tests {
    use std::{
        panic::{AssertUnwindSafe, catch_unwind},
        pin::Pin,
        task::{RawWaker, RawWakerVTable, Waker},
    };

    use super::*;
    use crate::{RequestInfo, RequestType};

    fn new_test_tracker(task_id: u64) -> TrackerToken {
        GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo {
            region_id: 1,
            start_ts: 0,
            task_id,
            resource_group_tag: vec![],
            request_type: RequestType::Unknown,
            is_remote: false,
        }))
    }

    fn noop_waker() -> Waker {
        fn clone(_: *const ()) -> RawWaker {
            RawWaker::new(std::ptr::null(), &VTABLE)
        }
        fn wake(_: *const ()) {}
        fn wake_by_ref(_: *const ()) {}
        fn drop(_: *const ()) {}

        static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);

        unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) }
    }

    fn poll_once<F: Future>(future: Pin<&mut F>) -> Poll<F::Output> {
        let waker = noop_waker();
        let mut cx = Context::from_waker(&waker);
        future.poll(&mut cx)
    }

    fn assert_trace_ctx_eq(lhs: TraceContext, rhs: &TraceContext) {
        assert_eq!(lhs.trace_id, rhs.trace_id);
        assert_eq!(lhs.control_flags(), rhs.control_flags());
    }

    #[test]
    fn test_tracked_future_restores_outer_tls_after_nested_poll() {
        let outer_tracker = new_test_tracker(1);
        let outer_trace_ctx = TraceContext::from_proto(b"outer", 1);
        set_tls_tracker_token(outer_tracker);
        set_tls_trace_ctx(outer_trace_ctx.clone());

        let mut future = std::pin::pin!(TrackedFuture::new({
            let outer_trace_ctx = outer_trace_ctx.clone();
            async move {
                assert_eq!(get_tls_tracker_token(), outer_tracker);
                assert_trace_ctx_eq(get_tls_trace_ctx(), &outer_trace_ctx);

                let inner_trace_ctx = outer_trace_ctx.clone();
                TrackedFuture::new(async move {
                    assert_eq!(get_tls_tracker_token(), outer_tracker);
                    assert_trace_ctx_eq(get_tls_trace_ctx(), &inner_trace_ctx);
                })
                .await;

                assert_eq!(get_tls_tracker_token(), outer_tracker);
                assert_trace_ctx_eq(get_tls_trace_ctx(), &outer_trace_ctx);
            }
        }));

        assert!(matches!(poll_once(future.as_mut()), Poll::Ready(())));
        assert_eq!(get_tls_tracker_token(), outer_tracker);
        assert_trace_ctx_eq(get_tls_trace_ctx(), &outer_trace_ctx);

        clear_tls_tracker_token();
        clear_tls_trace_ctx();
        GLOBAL_TRACKERS.remove(outer_tracker);
    }

    #[test]
    fn test_tracked_future_restores_previous_tls_after_panic() {
        let captured_tracker = new_test_tracker(1);
        let prev_tracker = new_test_tracker(2);
        let captured_trace_ctx = TraceContext::from_proto(b"captured", 1);
        let prev_trace_ctx = TraceContext::from_proto(b"prev", 2);

        set_tls_tracker_token(captured_tracker);
        set_tls_trace_ctx(captured_trace_ctx.clone());
        let mut future = std::pin::pin!(TrackedFuture::new(async {
            panic!("tracked future panic");
        }));

        set_tls_tracker_token(prev_tracker);
        set_tls_trace_ctx(prev_trace_ctx.clone());

        let result = catch_unwind(AssertUnwindSafe(|| poll_once(future.as_mut())));
        result.unwrap_err();
        assert_eq!(get_tls_tracker_token(), prev_tracker);
        assert_trace_ctx_eq(get_tls_trace_ctx(), &prev_trace_ctx);

        clear_tls_tracker_token();
        clear_tls_trace_ctx();
        GLOBAL_TRACKERS.remove(captured_tracker);
        GLOBAL_TRACKERS.remove(prev_tracker);
    }
}