// Copyright 2023 PingCAP, Inc.

use std::{error::Error as StdError, future::Future, sync::Arc, time::Duration};

use bytes::{Buf, Bytes};
use cloud_encryption::MasterKey;
use cloud_server::StatusServer as CloudStatusServer;
use dashmap::DashMap;
use flate2::{Compression, write::GzEncoder};
use http::{
    HeaderValue, Method, Request, Response, StatusCode,
    header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_TYPE},
    request::Parts,
};
use hyper::{
    Body,
    server::accept::Accept,
    service::{make_service_fn, service_fn},
};
use kvengine::{
    CURRENT_COMPACTOR_VERSION, CompactionCtx, CompactionRequest, CompactionType,
    INCOMPATIBLE_COMPACTOR_ERROR_CODE, IdAllocator, SCHEMA_OUT_OF_DATE_ERROR_CODE, SnapAccess,
    context::{IaCtx, MetaFileCacheWeighter, PrepareType, SnapCtx},
    dfs,
    dfs::{DFS_REMOTE_CACHE_ADDR_HEADER, Dfs},
    local_compact,
    metrics::ENGINE_REMOTE_COMPACT_EXCEED_MEMORY_LIMIT_COUNTER,
    table::{
        ChecksumType,
        columnar::ColumnarMetaCache,
        file::File,
        fts::{FtsCache, FtsDeltaCache},
        schema_file::SchemaFile,
        sstable::BlockCache,
    },
    txn_chunk_manager::TxnChunkManager,
};
use pd_client::PdClient;
use prometheus::TEXT_FORMAT;
use protobuf::Message;
use replication_worker::ReplicationScheduler;
use resource_metering::{FutureExt, ResourceTagFactory};
use rfstore::store::{PdIdAllocator, RegionSnapshot};
use security::HttpClient;
use tikv::{
    coprocessor::{
        REQ_TYPE_ANALYZE, REQ_TYPE_CHECKSUM, REQ_TYPE_DAG, RemoteCopRequestType,
        decode_query_digest, remote_cop_request_type,
        remote_dispatcher::{decode_remote_cop_request, large_topn_paging_size},
    },
    server::status_server::StatusServer,
};
use tikv_util::{
    deadline::Deadline,
    defer, error,
    http::{CONTENT_TYPE_PROTOBUF, HeaderExt},
    http_response, info,
    memory::MemoryLimiter,
    metrics::{dump, dump_to},
    time::Instant,
    warn,
};
use tipb::DagRequest;
use tokio::io::{AsyncRead, AsyncWrite};
use tracker::{
    GLOBAL_TRACKERS, RequestInfo, RequestType, TrackedFuture, clear_tls_tracker_token,
    set_tls_tracker_token,
};

use crate::{
    Config,
    encryption::KeyspaceEncryptionInfoCache,
    load_data::{self, LoadDataManager},
    metrics::*,
    native_br::{self, NativeBrManager},
    remote_cop::{record_remote_cop_outcome, record_remote_cop_scan_stats},
    txn_chunk,
    txn_chunk::TxnChunkHandler,
    worker_limiter::WorkerLimiter,
    write_sst,
    write_sst::WriteSstManager,
};

pub(crate) struct Context {
    pub config: Config,
    pub cluster_id: u64,
    pub block_size: usize,
    pub compression_lvl: i32,
    pub checksum_type: ChecksumType,
    pub thread_pool: tokio::runtime::Handle,
    pub dfs: Arc<dyn Dfs>,
    pub load_manager: Arc<LoadDataManager>,
    pub br_manager: Arc<NativeBrManager>,
    pub rep_scheduler: Option<ReplicationScheduler>,
    pub txn_chunk_handler: Arc<TxnChunkHandler>,
    pub pd: Arc<dyn PdClient>,
    pub master_key: MasterKey,
    pub encryption_key_manager: Arc<cloud_encryption::EncryptionKeyManager>,
    pub keyspace_encryption_info_cache: KeyspaceEncryptionInfoCache,
    pub block_cache: BlockCache,
    pub columnar_meta_cache: ColumnarMetaCache,
    pub schema_files: Option<Arc<DashMap<u64, SchemaFile>>>,
    pub coprocessor_limiter: WorkerLimiter,
    pub compaction_limiter: WorkerLimiter,
    pub memory_limiter: MemoryLimiter,
    pub txn_chunk_manager: TxnChunkManager,
    pub ia_ctx: IaCtx,
    pub read_columnar: bool,
    pub meta_file_cache: Arc<quick_cache::sync::Cache<u64, Arc<dyn File>, MetaFileCacheWeighter>>,
    pub http_client: Arc<HttpClient>,
    pub remote_cache_ttl: Duration,
    pub write_sst_manager: WriteSstManager,
    pub resource_tag_factory: ResourceTagFactory,
    pub rsmeter_pubsub_service: resource_metering::PubSubService,
    pub strict_file_memory_quota: Arc<tikv_util::memory::MemoryQuota>,
    pub strict_file_memory_quota_wait_timeout: Duration,
}

impl Context {
    pub(crate) fn get_snap_ctx(&self, dfs_remote_cache_addr: Option<&str>) -> SnapCtx {
        let dfs: Arc<dyn dfs::Dfs> = if let Some(dfs_cache_addr) = dfs_remote_cache_addr {
            Arc::new(dfs::RemoteCachedDfs::new(
                self.dfs.clone(),
                dfs_cache_addr.to_string(),
                self.http_client.clone(),
                self.remote_cache_ttl,
            ))
        } else {
            self.dfs.clone()
        };
        let prepare_type = if self.read_columnar {
            PrepareType::All
        } else {
            PrepareType::SstOnly
        };
        SnapCtx {
            dfs,
            master_key: self.master_key.clone(),
            block_cache: self.block_cache.clone(),
            columnar_meta_cache: self.columnar_meta_cache.clone(),
            vector_index_cache: None,
            columnar_file_cache: None,
            fts_cache: FtsCache::disabled(),
            fts_delta_cache: FtsDeltaCache::disabled(),
            schema_files: self.schema_files.clone(),
            txn_chunk_manager: self.txn_chunk_manager.clone(),
            ia_ctx: self.ia_ctx.clone(),
            prepare_type,
            read_columnar: self.read_columnar,
            encryption_key_manager: self.encryption_key_manager.clone(),
            meta_file_cache: self.meta_file_cache.clone(),
            strict_file_memory_quota: self.strict_file_memory_quota.clone(),
            strict_file_memory_quota_wait_timeout: self.strict_file_memory_quota_wait_timeout,
        }
    }
}

#[macro_export]
macro_rules! start_serve {
    ($ctx:expr, $acceptor:expr) => {{
        match $acceptor {
            tikv_util::Either::Left(acceptor) => $crate::server::start(
                $ctx,
                hyper::server::Server::builder(acceptor)
                    .http1_header_read_timeout(SERVER_READ_TIMEOUT),
            ),
            tikv_util::Either::Right(acceptor) => $crate::server::start(
                $ctx,
                hyper::server::Server::builder(acceptor)
                    .http1_header_read_timeout(SERVER_READ_TIMEOUT),
            ),
        }
    }};
}

pub(crate) fn start<I, C>(
    ctx: Arc<Context>,
    builder: hyper::server::Builder<I>,
) -> Box<dyn Future<Output = hyper::Result<()>> + Send + Unpin>
where
    I: Accept<Conn = C, Error = std::io::Error> + Send + Unpin + 'static,
    I::Error: Into<Box<dyn StdError + Send + Sync>>,
    I::Conn: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    let server = builder.serve(make_service_fn(move |_| {
        let ctx = ctx.clone();
        async move {
            Ok::<_, hyper::Error>(service_fn(move |req: hyper::Request<hyper::Body>| {
                let ctx = ctx.clone();
                tikv_util::init_task_local(async move {
                    let path = req.uri().path().to_owned();
                    match path.as_ref() {
                        "/healthz" => http_response!(hyper::StatusCode::OK, "ok"),
                        "/compact" => handle_compaction(ctx, req).await,
                        path if path.starts_with("/cdc") => {
                            replication_worker::handle_cdc_request(ctx.rep_scheduler.as_ref(), req)
                                .await
                        }
                        "/coprocessor" => handle_coprocessor(ctx, req).await,
                        "/load_data" => handle_load_data(ctx, req).await,
                        "/metrics" => handle_get_metrics(req).await,
                        "/debug/pprof/profile" => {
                            StatusServer::<u8, u8>::dump_cpu_prof_to_resp(req).await
                        }
                        "/debug/pprof/heap" => CloudStatusServer::dump_heap_prof_to_resp(req),
                        "/debug/pprof/cmdline" => CloudStatusServer::get_cmdline(req),
                        "/debug/pprof/symbol" => {
                            if req.method() == Method::GET {
                                CloudStatusServer::get_symbol_count(req)
                            } else if req.method() == Method::POST {
                                CloudStatusServer::get_symbol(req).await
                            } else {
                                http_response!(hyper::StatusCode::NOT_FOUND, "Not Found")
                            }
                        }
                        #[cfg(debug_assertions)]
                        "/debug/sleep" => handle_sleep(ctx).await,
                        #[cfg(debug_assertions)]
                        "/debug/load-data-sleep" => handle_load_data_sleep(ctx).await,
                        native_br::BACKUPS_API_PATH => {
                            native_br::handle_backup(ctx.br_manager.clone(), req).await
                        }
                        path if path.starts_with(native_br::RESTORE_KEYSPACE_API_PATH) => {
                            handle_restore_keyspace(ctx, req).await
                        }
                        "/txn_chunk" => handle_txn_chunk(ctx, req).await,
                        "/write_sst" => handle_write_sst(ctx, req).await,

                        path if path.starts_with(native_br::v1x::API_V1X) => {
                            native_br::v1x::serve(ctx.br_manager.clone(), req).await
                        }
                        _ => http_response!(hyper::StatusCode::NOT_FOUND, "Not Found"),
                    }
                })
            }))
        }
    }));
    Box::new(server)
}

// Helper function to spawn a task on the thread pool and await its result
async fn spawn_and_await<F, T>(thread_pool: tokio::runtime::Handle, f: F) -> T
where
    F: Future<Output = T> + Send + 'static,
    T: Send + 'static + std::fmt::Debug,
{
    let handle = thread_pool.spawn(tikv_util::init_task_local(f));
    match handle.await {
        Ok(res) => res,
        Err(err) => {
            panic!("spawn and await failed: {:?}", err);
        }
    }
}

async fn handle_compaction(
    ctx: Arc<Context>,
    req: Request<Body>,
) -> hyper::Result<hyper::Response<Body>> {
    // Parse the request body to get the compaction request.
    let body_bytes = hyper::body::to_bytes(req.into_body()).await?;
    let result: Result<CompactionRequest, _> = serde_json::from_slice(&body_bytes);
    if let Err(err) = &result {
        return http_response!(http::StatusCode::BAD_REQUEST, err.to_string());
    }
    let comp_req = result.unwrap();

    let keyspace_id = 0; // Keyspace is not supported yet.
    // Try to acquire a permit with timeout
    let timeout = ctx.compaction_limiter.wait_timeout();
    let res =
        tokio::time::timeout(timeout, ctx.compaction_limiter.acquire_permit(keyspace_id)).await;

    match res {
        Ok(Some(_permit)) => handle_compaction_internal(ctx, comp_req).await,
        Ok(None) => {
            // Queue is full
            REMOTE_COMPACT_FAILED_REQUESTS_COUNTER_VEC.queue_full.inc();
            http_response!(
                http::StatusCode::SERVICE_UNAVAILABLE,
                "compaction request wait queue is full"
            )
        }
        Err(_) => {
            // Timeout
            REMOTE_COMPACT_FAILED_REQUESTS_COUNTER_VEC
                .wait_timeout
                .inc();
            http_response!(http::StatusCode::SERVICE_UNAVAILABLE, "wait permit timeout")
        }
    }
}

// Extract common compaction processing logic into a separate function
async fn handle_compaction_internal(
    ctx: Arc<Context>,
    comp_req: CompactionRequest,
) -> hyper::Result<hyper::Response<Body>> {
    REMOTE_COMPACT_PROCESSING_REQUESTS_COUNTER.inc();
    defer! {
        REMOTE_COMPACT_PROCESSING_REQUESTS_COUNTER.dec()
    }

    let handle_start = Instant::now_coarse();
    let allocator = Arc::new(PdIdAllocator::new(ctx.pd.clone()));
    let resp = execute_remote_compaction(ctx, comp_req, allocator).await;

    if resp.is_ok() && resp.as_ref().unwrap().status().is_success() {
        // Record processing time histogram
        REMOTE_COMPACT_REQ_HANDLE_HISTOGRAM
            .observe(handle_start.saturating_elapsed().as_secs_f64());
    }

    resp
}

/// Execute a remote compaction request submitted by a tikv node.
///
/// Performs three layers of admission control before building a
/// `CompactionCtx` and delegating to `kvengine::local_compact` on a
/// thread-pool worker.
async fn execute_remote_compaction(
    ctx: Arc<Context>,
    comp_req: CompactionRequest,
    id_allocator: Arc<dyn IdAllocator>,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    let thread_pool = ctx.thread_pool.clone();
    let dfs = ctx.dfs.clone();
    let compression_lvl = ctx.compression_lvl;
    let checksum_type = ctx.checksum_type;
    let master_key = ctx.master_key.clone();

    if comp_req.compactor_version > CURRENT_COMPACTOR_VERSION {
        warn!(
            "received incompatible compactor-version({}). Upgrade tikv-worker (version:{}). Request: {:?}",
            comp_req.compactor_version, CURRENT_COMPACTOR_VERSION, comp_req,
        );
        let err_str = format!("incompatible compactor version ({CURRENT_COMPACTOR_VERSION})");
        return Ok(hyper::Response::builder()
            .status(INCOMPATIBLE_COMPACTOR_ERROR_CODE)
            .body(err_str.into())
            .unwrap());
    }

    if comp_req.input_size == 0 {
        warn!("input_size not set: {:?}", comp_req);
        // TODO: `debug_assert!(false)`.
    }

    // Vector-index throttle (only for `CompactionType::VectorIndex`): tries a
    // non-blocking permit acquire; returns HTTP 503 immediately on failure so
    // the tikv node can retry another worker instance.
    let _vector_index_permit =
        if let CompactionType::VectorIndex(vec_idx_update) = &comp_req.compaction_tp {
            // Try acquire permit without waiting. If failed, return ASAP to make the client
            // retry request to another worker instance.
            match ctx.compaction_limiter.try_acquire_vector_index_permit() {
                Ok(permit) => Some(permit),
                Err(_) => {
                    let err_str = "vector index compaction requests throttled";
                    warn!("{}: {:?}", err_str, vec_idx_update);
                    let body = hyper::Body::from(err_str);
                    return Ok(hyper::Response::builder()
                        .status(StatusCode::SERVICE_UNAVAILABLE)
                        .body(body)
                        .unwrap());
                }
            }
        } else {
            None
        };

    // Memory limit: acquires `input_size * 2` from the worker-wide `memory_limiter`
    // for both reading and writing. If the limit is exceeded, returns HTTP 503.
    let memory_limiter = ctx.memory_limiter.clone();
    let request_size = comp_req.input_size * 2;
    let mem_limiter_guard = match memory_limiter.acquire(request_size) {
        Ok(guard) => guard,
        Err(exceeded_size) => {
            ENGINE_REMOTE_COMPACT_EXCEED_MEMORY_LIMIT_COUNTER.inc();
            warn!("{} memory limit exceeded", comp_req.get_tag();
                "request_size" => request_size, "exceeded" => exceeded_size,
                "limiter" => ?memory_limiter);
            let body = hyper::Body::from("memory limit exceeded");
            return Ok(hyper::Response::builder()
                .status(StatusCode::SERVICE_UNAVAILABLE)
                .body(body)
                .unwrap());
        }
    };

    let encryption_key = if !comp_req.exported_encryption_meta.is_empty() {
        match cloud_encryption::current_data_key_from_encryption_meta_bytes(
            ctx.encryption_key_manager.clone(),
            &comp_req.exported_encryption_meta,
            Some(comp_req.keyspace_id()),
        ) {
            Ok(key) => Some(key),
            Err(e) => {
                let is_retryable = e.is_retryable();
                warn!(
                    "failed to parse exported encryption meta in remote compaction request";
                    "tag" => comp_req.get_tag(),
                    "keyspace_id" => comp_req.keyspace_id(),
                    "err" => %e,
                );
                if is_retryable {
                    return Ok(hyper::Response::builder()
                        .status(StatusCode::SERVICE_UNAVAILABLE)
                        .body(format!("encryption key is not ready: {}", e).into())
                        .unwrap());
                } else {
                    return Ok(hyper::Response::builder()
                        .status(StatusCode::BAD_REQUEST)
                        .body(format!("invalid exported encryption meta: {}", e).into())
                        .unwrap());
                }
            }
        }
    } else if !comp_req.exported_encryption_key.is_empty() {
        Some(
            master_key
                .decrypt_encryption_key(&comp_req.exported_encryption_key)
                .unwrap(),
        )
    } else {
        None
    };
    let comp_ctx = CompactionCtx {
        req: Arc::new(comp_req),
        dfs,
        compression_lvl,
        id_allocator,
        encryption_key,
        local_dirs: vec![],
        for_restore: false,
        checksum_type,
        columnar_meta_cache: ctx.columnar_meta_cache.clone(),
    };

    // The actual compaction runs asynchronously on the thread pool.
    let task = thread_pool.spawn(tikv_util::init_task_local(async move {
        tikv_util::set_current_region(comp_ctx.req.shard_id);
        let _guard = mem_limiter_guard;
        local_compact(&comp_ctx).await
    }));

    match task.await {
        Ok(Ok(cs)) => {
            // serialized as `pb::ChangeSet` in the HTTP response body.
            let data = cs.write_to_bytes().unwrap();
            Ok(hyper::Response::builder()
                .status(StatusCode::OK)
                .body(data.into())
                .unwrap())
        }
        Ok(Err(e)) => {
            let err_str = e.to_string();
            error!("compaction failed {}", err_str);
            let (status, body) = match &e {
                kvengine::Error::TableError(kvengine::table::Error::SchemaOutOfDate(info)) => (
                    SCHEMA_OUT_OF_DATE_ERROR_CODE,
                    hyper::Body::from(serde_json::to_string(info).unwrap()),
                ),
                _ => (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    hyper::Body::from(err_str),
                ),
            };
            Ok(hyper::Response::builder()
                .status(status)
                .body(body)
                .unwrap())
        }
        Err(join_err) => {
            let err_str = join_err.to_string();
            error!("compaction failed {}", err_str);
            let body = hyper::Body::from(err_str);
            Ok(hyper::Response::builder()
                .status(StatusCode::INTERNAL_SERVER_ERROR)
                .body(body)
                .unwrap())
        }
    }
}

async fn handle_coprocessor(
    ctx: Arc<Context>,
    req: hyper::Request<hyper::Body>,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    let (parts, body) = req.into_parts();
    let body = hyper::body::to_bytes(body).await?;
    spawn_and_await(
        ctx.thread_pool.clone(),
        handle_remote_coprocessor(ctx, parts, body),
    )
    .await
}

async fn handle_load_data(
    ctx: Arc<Context>,
    req: hyper::Request<hyper::Body>,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    let load_data_manager = ctx.load_manager.clone();
    let (parts, body) = req.into_parts();
    let body = hyper::body::to_bytes(body).await?;
    spawn_and_await(
        ctx.thread_pool.clone(),
        load_data::handle_load_data(load_data_manager, parts, body),
    )
    .await
}

async fn handle_restore_keyspace(
    ctx: Arc<Context>,
    req: hyper::Request<hyper::Body>,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    let br_manager = ctx.br_manager.clone();
    let (parts, _) = req.into_parts();
    spawn_and_await(
        ctx.thread_pool.clone(),
        native_br::handle_restore_keyspace(br_manager, parts),
    )
    .await
}

async fn handle_txn_chunk(
    ctx: Arc<Context>,
    req: hyper::Request<hyper::Body>,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    let (parts, body) = req.into_parts();
    let body = hyper::body::to_bytes(body).await?;
    spawn_and_await(
        ctx.thread_pool.clone(),
        txn_chunk::handle_txn_chunk(ctx, parts, body),
    )
    .await
}

async fn handle_write_sst(
    ctx: Arc<Context>,
    req: hyper::Request<hyper::Body>,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    spawn_and_await(
        ctx.thread_pool.clone(),
        write_sst::handle_write_sst(ctx, req),
    )
    .await
}

// Debug API to sleep for 5 seconds
#[cfg(debug_assertions)]
async fn handle_sleep(_ctx: Arc<Context>) -> hyper::Result<hyper::Response<hyper::Body>> {
    info!("sleep for 5 seconds");
    std::thread::sleep(Duration::from_secs(5));
    info!("sleep done");
    http_response!(http::StatusCode::OK, "OK")
}

#[cfg(debug_assertions)]
async fn handle_load_data_sleep(ctx: Arc<Context>) -> hyper::Result<hyper::Response<hyper::Body>> {
    let load_data_manager = ctx.load_manager.clone();
    spawn_and_await(ctx.thread_pool.clone(), async move {
        info!("sleep for 5 seconds on load-data worker runtime");
        tokio::time::sleep(Duration::from_secs(5)).await;
        drop(load_data_manager);
        info!("load-data worker runtime sleep done");
        http_response!(http::StatusCode::OK, "OK")
    })
    .await
}

/// Ref: https://docs.pingcap.com/tidb/stable/tidb-configuration-file/#copr-req-timeout-new-in-v750
const DEFAULT_COP_TIMEOUT: Duration = Duration::from_secs(60);

// Return `StatusCode::SERVICE_UNAVAILABLE` to indicate that the request can be
// retried later.
async fn handle_remote_coprocessor(
    ctx: Arc<Context>,
    parts: Parts,
    req_body: Bytes,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    let accept_pb = parts.headers.is_accept_protobuf();
    let decode_res = decode_remote_cop_request(req_body.chunk());
    if let Err(err) = decode_res {
        return http_response!(http::StatusCode::BAD_REQUEST, format!("{:?}", err));
    }
    let (req_data, mem_data, snap_data) = decode_res.unwrap();
    let mut cop_req = kvproto::coprocessor::Request::default();
    if let Err(err) = cop_req.merge_from_bytes(req_data) {
        return http_response!(http::StatusCode::BAD_REQUEST, format!("{:?}", err));
    }
    let cop_ctx = cop_req.get_context();
    let keyspace_id = cop_ctx.get_keyspace_id();
    let handle_start = Instant::now_coarse();
    let request_type = remote_cop_request_type(&cop_req);
    REMOTE_COPR_RANGE_COUNT_HISTOGRAM.observe(cop_req.get_ranges().len() as f64);
    let tag = get_cop_req_tag(&cop_req);
    if ctx.config.black_list.rejects_request(&cop_req) {
        warn!("request rejected by blacklist"; "tag" => &tag);
        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.black_list.inc();
        record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
        return http_response!(http::StatusCode::FORBIDDEN, "rejected");
    }
    let mut timeout = Duration::from_millis(cop_ctx.get_max_execution_duration_ms());
    if timeout.is_zero() {
        timeout = DEFAULT_COP_TIMEOUT;
    }
    let deadline = Deadline::from_now(timeout);

    let res =
        tokio::time::timeout(timeout, ctx.coprocessor_limiter.acquire_permit(keyspace_id)).await;
    match res {
        Ok(Some(_permit)) => {
            let resource_tag = ctx.resource_tag_factory.new_tag(cop_req.get_context());
            handle_remote_coprocessor_internal(
                ctx,
                cop_req,
                req_body.len(),
                get_dfs_remote_cache_addr(&parts),
                mem_data,
                snap_data,
                handle_start,
                timeout,
                deadline,
                tag,
                accept_pb,
                request_type,
            )
            .in_resource_metering_tag(resource_tag)
            .await
        }
        Ok(None) => {
            // Queue is full
            REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.queue_full.inc();
            record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
            http_response!(
                http::StatusCode::SERVICE_UNAVAILABLE,
                "coprocessor request wait queue is full"
            )
        }
        Err(_) => {
            // Timeout
            REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.wait_timeout.inc();
            record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
            http_response!(http::StatusCode::SERVICE_UNAVAILABLE, "wait permit timeout")
        }
    }
}

async fn handle_remote_coprocessor_internal(
    ctx: Arc<Context>,
    mut cop_req: kvproto::coprocessor::Request,
    req_len: usize,
    dfs_remote_cache_addr: Option<&str>,
    mem_data: &[u8],
    snap_data: &[u8],
    handle_start: Instant,
    timeout: Duration,
    deadline: Deadline,
    tag: String,
    accept_pb: bool,
    request_type: RemoteCopRequestType,
) -> hyper::Result<hyper::Response<hyper::Body>> {
    REMOTE_COPR_PROCESSING_REQ_COUNTER.inc();
    defer! {
        REMOTE_COPR_PROCESSING_REQ_COUNTER.dec()
    }

    if set_topn_paging_size(&mut cop_req, ctx.config.cop_topn_max_pushdown_rows) {
        REMOTE_COPR_TOPN_PAGING_REWRITE_TOTAL.inc();
    }
    let req_type = cop_req.get_tp();
    let query_digest = decode_query_digest(cop_req.get_context().get_resource_group_tag());
    let snap_start = Instant::now_coarse();
    let snap_ctx = ctx.get_snap_ctx(dfs_remote_cache_addr);
    let mem_limiter = ctx.memory_limiter.clone();
    let snap_access_res =
        SnapAccess::construct_snapshot(&tag, &snap_ctx, mem_data, snap_data, mem_limiter).await;
    if let Err(err) = snap_access_res.as_ref() {
        let status_code = if matches!(err, kvengine::Error::MemoryLimitExceeded(_)) {
            REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.memory_limiter.inc();
            http::StatusCode::SERVICE_UNAVAILABLE
        } else {
            REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.snapshot_error.inc();
            http::StatusCode::INTERNAL_SERVER_ERROR
        };
        record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
        return http_response!(status_code, format!("{:?}", err));
    }
    if handle_start.saturating_elapsed() > timeout {
        info!("construct snapshot timeout"; "tag" => tag);
        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
            .snapshot_timeout
            .inc();
        record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
        return http_response!(
            http::StatusCode::SERVICE_UNAVAILABLE,
            "construct snapshot timeout"
        );
    }
    let (snap_access, mem_limiter_guard) = snap_access_res.unwrap();

    let tracker = GLOBAL_TRACKERS.insert(tracker::Tracker::new(RequestInfo::new(
        cop_req.get_context(),
        RequestType::Unknown,
        cop_req.start_ts,
    )));
    set_tls_tracker_token(tracker);
    defer! {{
        clear_tls_tracker_token();
        GLOBAL_TRACKERS.remove(tracker);
    }};
    TrackedFuture::new(async move {
        let prefetch_start = Instant::now_coarse();
        if matches!(req_type, REQ_TYPE_DAG if ctx.ia_ctx.is_enabled()) {
            let prefetch_result = tikv::coprocessor::prefetch_ia_remote_segments(
                &tag,
                &snap_ctx,
                &snap_access,
                cop_req.get_ranges(),
                deadline,
            )
            .await;
            match prefetch_result {
                Ok(cache_hit) => {
                    if let Some(cache_hit) = cache_hit {
                        REMOTE_COPR_PREFETCH_CACHE_HIT_PERCENT_HISTOGRAM.observe(cache_hit * 100.0);
                    }
                }
                Err(err) => {
                    error!("{} prefetch failed, error {:?}", tag, err);
                    if matches!(err, tikv::coprocessor::Error::DeadlineExceeded) {
                        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
                            .prefetch_timeout
                            .inc();
                    } else {
                        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.prefetch_error.inc();
                    }
                    record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                    // Encode as a coprocessor pb error so DeadlineExceeded maps
                    // to a server_is_busy region error and tidb-server retries.
                    if accept_pb {
                        let cop_resp = tikv::coprocessor::make_error_response(err);
                        let body = hyper::Body::from(cop_resp.write_to_bytes().unwrap());
                        return Ok(hyper::Response::builder()
                            .status(http::StatusCode::INTERNAL_SERVER_ERROR)
                            .header(CONTENT_TYPE, CONTENT_TYPE_PROTOBUF)
                            .body(body)
                            .unwrap());
                    }
                    return http_response!(
                        http::StatusCode::INTERNAL_SERVER_ERROR,
                        format!("prefetch segments failed: {:?}", err)
                    );
                }
            }
            if deadline.check().is_err() {
                info!("prefetch timeout"; "tag" => tag);
                REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
                    .prefetch_timeout
                    .inc();
                record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                return http_response!(
                    http::StatusCode::SERVICE_UNAVAILABLE,
                    "prefetch segments timeout"
                );
            }
        }
        let snap = RegionSnapshot::from_snapshot(snap_access, None, None);
        let process_start = Instant::now_coarse();
        let before_duration = process_start.saturating_duration_since(handle_start);
        // Calculate the remaining duration of the request timeout for cop processing
        let remaining_duration = match timeout.checked_sub(before_duration) {
            Some(dur) => dur,
            None => {
                info!("before process timeout"; "tag" => tag);
                REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
                    .process_timeout
                    .inc();
                record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                let body = hyper::Body::from("before process timeout");
                return Ok(hyper::Response::builder()
                    .status(StatusCode::SERVICE_UNAVAILABLE)
                    .body(body)
                    .unwrap());
            }
        };
        let result = tikv::coprocessor::parse_request_and_handle_remote_cop(
            cop_req,
            None,
            remaining_duration,
            ctx.config.cop_max_resp_size.0,
            snap,
        )
        .await;
        drop(mem_limiter_guard);
        if let Err(err) = result {
            // Strict-file-memory-quota wait timeouts surface as
            // `Error::Other(String)` after passing through the coprocessor /
            // iterator layers. Return HTTP 503 so TiKV's remote_dispatcher
            // maps it to RemoteServiceUnavailable and retries with backoff.
            // This is an expected fast-fail path, so log at warn level.
            if is_strict_file_memory_quota_err(&err) {
                warn!("remote coprocessor hit strict file memory quota limit";
                    "tag" => &tag,
                    "err" => ?err,
                );
                REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.memory_limiter.inc();
                record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                let body = hyper::Body::from(format!("{}", err));
                return Ok(hyper::Response::builder()
                    .status(StatusCode::SERVICE_UNAVAILABLE)
                    .body(body)
                    .unwrap());
            }
            error!("{} remote coprocessor failed, error {:?}", tag, err);
            if matches!(&err, tikv::coprocessor::Error::DeadlineExceeded) {
                REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
                    .process_timeout
                    .inc();
            } else {
                REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.process_error.inc();
            }
            record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
            if accept_pb {
                let cop_resp = tikv::coprocessor::make_error_response(err);
                let body = hyper::Body::from(cop_resp.write_to_bytes().unwrap());
                return Ok(hyper::Response::builder()
                    .status(http::StatusCode::INTERNAL_SERVER_ERROR)
                    .header(CONTENT_TYPE, CONTENT_TYPE_PROTOBUF)
                    .body(body)
                    .unwrap());
            } else {
                return http_response!(
                    http::StatusCode::INTERNAL_SERVER_ERROR,
                    format!("{:?}", err)
                );
            }
        }

        let finish_time = Instant::now_coarse();
        let permit_wait_duration = snap_start.saturating_duration_since(handle_start);
        let snap_duration = prefetch_start.saturating_duration_since(snap_start);
        let prefetch_duration = process_start.saturating_duration_since(prefetch_start);
        let process_duration = finish_time.saturating_duration_since(process_start);
        let handle_duration = finish_time.saturating_duration_since(handle_start);
        let result = result.unwrap();
        REMOTE_COPR_SNAPSHOT_HISTOGRAM.observe(snap_duration.as_secs_f64());
        REMOTE_COPR_PREFETCH_HISTOGRAM.observe(prefetch_duration.as_secs_f64());
        REMOTE_COPR_PROCESS_HISTOGRAM.observe(process_duration.as_secs_f64());
        REMOTE_COPR_REQ_HANDLE_HISTOGRAM.observe(handle_duration.as_secs_f64());
        record_remote_cop_scan_stats(result.scan_stats);
        let response = result.response;
        info!(
            "finished remote coprocessor";
            "tag" => &tag,
            "query_digest" => &query_digest,
            "req_size" => req_len,
            "resp_size" => response.data.len(),
            "timeout" => ?timeout,
            "permit_wait_duration" => ?permit_wait_duration,
            "snap_duration" => ?snap_duration,
            "prefetch_duration" => ?prefetch_duration,
            "process_duration" => ?process_duration,
            "handle_duration" => ?handle_duration,
        );

        match req_type {
            REQ_TYPE_DAG => {
                REMOTE_COPR_DAG_REQ_COUNTER.inc();
                REMOTE_COPR_DAG_RESP_SIZE.inc_by(response.data.len() as u64);
            }
            REQ_TYPE_ANALYZE => {
                REMOTE_ANALYZE_REQ_COUNTER.inc();
                REMOTE_ANALYZE_RESP_SIZE.inc_by(response.data.len() as u64);
            }
            REQ_TYPE_CHECKSUM => {
                REMOTE_CHECKSUM_REQ_COUNTER.inc();
                REMOTE_CHECKSUM_RESP_SIZE.inc_by(response.data.len() as u64);
            }
            _ => {}
        }
        match response.write_to_bytes() {
            Ok(response_data) => {
                record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::success);
                http_response!(http::StatusCode::OK, response_data)
            }
            Err(err) => {
                error!("{} serialize response failed, error {:?}", tag, err);
                REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
                    .response_encode_error
                    .inc();
                record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                http_response!(
                    http::StatusCode::INTERNAL_SERVER_ERROR,
                    format!("{:?}", err)
                )
            }
        }
    })
    .await
}

/// Applies the compatibility TopN paging policy to remote requests from TiKV
/// versions that still send the paging sentinel. A zero threshold or a finite
/// page returns before decoding the nested DAG.
fn set_topn_paging_size(
    cop_req: &mut kvproto::coprocessor::Request,
    max_pushdown_rows: u64,
) -> bool {
    if max_pushdown_rows == 0
        || cop_req.get_tp() != REQ_TYPE_DAG
        || cop_req.get_paging_size() != u32::MAX as u64
    {
        return false;
    }

    // Leave malformed DAG requests unchanged so the normal request parser
    // produces the canonical coprocessor error response.
    let mut dag = DagRequest::default();
    if dag.merge_from_bytes(cop_req.get_data()).is_err() {
        return false;
    }
    if let Some(paging_size) = large_topn_paging_size(&dag, max_pushdown_rows) {
        cop_req.set_paging_size(paging_size);
        return true;
    }
    false
}

fn get_dfs_remote_cache_addr(parts: &http::request::Parts) -> Option<&str> {
    parts
        .headers
        .get(DFS_REMOTE_CACHE_ADDR_HEADER)?
        .to_str()
        .ok()
}

pub fn get_cop_req_tag(cop_req: &kvproto::coprocessor::Request) -> String {
    let cop_ctx = cop_req.get_context();
    let keyspace_id = cop_ctx.get_keyspace_id();
    let start_ts = cop_req.get_start_ts();
    let req_type = cop_req.get_tp();
    let req_type_str = match req_type {
        REQ_TYPE_DAG => "dag".to_string(),
        REQ_TYPE_ANALYZE => "analyze".to_string(),
        REQ_TYPE_CHECKSUM => "checksum".to_string(),
        _ => "".to_string(),
    };
    format!(
        "{} ks{}:{}:{}:{}",
        req_type_str,
        keyspace_id,
        cop_ctx.get_region_id(),
        cop_ctx.get_region_epoch().get_version(),
        start_ts,
    )
}

async fn handle_get_metrics(req: Request<Body>) -> hyper::Result<Response<Body>> {
    let gz_encoding = client_accept_gzip(&req);
    let metrics = if gz_encoding {
        // gzip can reduce the body size to less than 1/10.
        let mut encoder = GzEncoder::new(vec![], Compression::default());
        dump_to(&mut encoder, true);
        encoder.finish().unwrap()
    } else {
        dump(true).into_bytes()
    };
    let mut resp = Response::new(metrics.into());
    resp.headers_mut()
        .insert(CONTENT_TYPE, HeaderValue::from_static(TEXT_FORMAT));
    if gz_encoding {
        resp.headers_mut()
            .insert(CONTENT_ENCODING, HeaderValue::from_static("gzip"));
    }
    Ok(resp)
}

// Reports whether the client accepts gzip compression. Adapted from Prometheus:
// https://github.com/prometheus/client_golang/blob/24172847e35ba46025c49d90b8846b59eb5d9ead/prometheus/promhttp/http.go#L367-L378
// Copyright 2016 The Prometheus Authors
// Modifications Copyright 2023 PingCAP, Inc.
fn client_accept_gzip(req: &Request<Body>) -> bool {
    let encoding = req
        .headers()
        .get(ACCEPT_ENCODING)
        .map(|enc| enc.to_str().unwrap_or_default())
        .unwrap_or_default();
    encoding
        .split(',')
        .map(|s| s.trim())
        .any(|s| s == "gzip" || s.starts_with("gzip;"))
}

// Detect strict-file-memory-quota wait timeouts. These surface as
// `Error::Other(String)` after passing through coprocessor / iterator layers,
// so we match on the Display string. The string is deliberately unique
// (component-prefixed) to avoid colliding with unrelated "memory" errors
// elsewhere in the stack. If you change `dfs::Error::MemoryQuotaWaitTimeout`'s
// Display, update the substring here.
//
// TODO: Replace string matching with a typed `MemoryQuotaExceeded` variant on
// `coprocessor::Error`.
fn is_strict_file_memory_quota_err<E: std::fmt::Display>(err: &E) -> bool {
    format!("{}", err).contains("strict file memory quota wait timeout")
}

#[cfg(test)]
mod tests {
    use kvproto::coprocessor::Request;
    use tipb::{ExecType, Executor, TopN};

    use super::*;

    fn topn_request(limit: u64, partitioned: bool) -> Request {
        let mut topn = TopN::default();
        topn.set_limit(limit);
        if partitioned {
            topn.mut_partition_by().push(Default::default());
        }

        let mut executor = Executor::default();
        executor.set_tp(ExecType::TypeTopN);
        executor.set_top_n(topn);

        let mut dag = DagRequest::default();
        dag.mut_executors().push(executor);

        let mut req = Request::default();
        req.set_tp(REQ_TYPE_DAG);
        req.set_data(dag.write_to_bytes().unwrap());
        req.set_paging_size(u32::MAX as u64);
        req
    }

    #[test]
    fn test_set_topn_paging_size_for_large_topn() {
        let mut req = topn_request(50_001, false);
        assert!(set_topn_paging_size(&mut req, 50_000));
        assert_eq!(req.get_paging_size(), 50_000);
    }

    #[test]
    fn test_keep_paging_size_for_topn_at_threshold() {
        let mut req = topn_request(50_000, false);
        assert!(!set_topn_paging_size(&mut req, 50_000));
        assert_eq!(req.get_paging_size(), u32::MAX as u64);
    }

    #[test]
    fn test_keep_paging_size_for_partition_topn() {
        let mut req = topn_request(50_001, true);
        assert!(!set_topn_paging_size(&mut req, 50_000));
        assert_eq!(req.get_paging_size(), u32::MAX as u64);
    }

    #[test]
    fn test_keep_caller_paging_size() {
        let mut req = topn_request(50_001, false);
        req.set_paging_size(128);

        assert!(!set_topn_paging_size(&mut req, 50_000));
        assert_eq!(req.get_paging_size(), 128);
    }

    #[test]
    fn test_zero_threshold_disables_topn_bypass() {
        let mut req = topn_request(50_001, false);

        assert!(!set_topn_paging_size(&mut req, 0));
        assert_eq!(req.get_paging_size(), u32::MAX as u64);
    }

    #[test]
    fn test_zero_threshold_keeps_malformed_dag_untouched() {
        let mut req = Request::default();
        req.set_tp(REQ_TYPE_DAG);
        req.set_paging_size(u32::MAX as u64);
        req.set_data(vec![0xff]);

        assert!(!set_topn_paging_size(&mut req, 0));
        assert_eq!(req.get_paging_size(), u32::MAX as u64);
        assert_eq!(req.get_data(), &[0xff]);
    }

    #[test]
    fn test_finite_page_keeps_malformed_dag_untouched() {
        let mut req = Request::default();
        req.set_tp(REQ_TYPE_DAG);
        req.set_paging_size(50_000);
        req.set_data(vec![0xff]);

        assert!(!set_topn_paging_size(&mut req, 50_000));
        assert_eq!(req.get_paging_size(), 50_000);
        assert_eq!(req.get_data(), &[0xff]);
    }
}