// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use std::{
    convert::Infallible,
    net::TcpListener,
    sync::{Arc, Mutex},
    thread::sleep,
    time::Duration,
};

use http::{Method, StatusCode};
use hyper::{
    Body, Request, Response, Server,
    service::{make_service_fn, service_fn},
};
use kvproto::metapb::{self, PeerRole};
use pd_client::PdClient;
use serde::Deserialize;
use slog_global::warn;
use tokio::runtime::Runtime;

use crate::{PdClientExt, TestPdClient};

type KeyspaceLookup = dyn Fn(&str) -> Option<u32> + Send + Sync;

/// Minimal PD HTTP control surface backed by an in-memory [`TestPdClient`].
///
/// TestPD exposes the gRPC APIs but not PD's HTTP APIs. This server supports
/// the keyspace lookup, scheduling rule, and transfer-region endpoints needed
/// by native BR integration and random tests.
pub struct TestPdControlServer {
    endpoint: String,
    shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
}

impl TestPdControlServer {
    /// Starts a PD HTTP control server.
    ///
    /// `keyspace_lookup` must reflect keyspaces created after the server
    /// starts.
    pub fn start(
        runtime: &Runtime,
        pd_client: Arc<TestPdClient>,
        keyspace_lookup: impl Fn(&str) -> Option<u32> + Send + Sync + 'static,
    ) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let endpoint = listener.local_addr().unwrap().to_string();
        let keyspace_lookup: Arc<KeyspaceLookup> = Arc::new(keyspace_lookup);
        let transfer_lock = Arc::new(Mutex::new(()));
        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
        runtime.spawn(async move {
            let make_service = make_service_fn(move |_| {
                let pd_client = pd_client.clone();
                let keyspace_lookup = keyspace_lookup.clone();
                let transfer_lock = transfer_lock.clone();
                async move {
                    Ok::<_, Infallible>(service_fn(move |request| {
                        handle_request(
                            request,
                            pd_client.clone(),
                            keyspace_lookup.clone(),
                            transfer_lock.clone(),
                        )
                    }))
                }
            });
            let server = Server::from_tcp(listener)
                .unwrap()
                .serve(make_service)
                .with_graceful_shutdown(async {
                    let _ = shutdown_rx.await;
                });
            if let Err(err) = server.await {
                warn!("test PD control server stopped with error: {:?}", err);
            }
        });
        Self {
            endpoint,
            shutdown_tx: Some(shutdown_tx),
        }
    }

    /// Returns the HTTP endpoint without a URL scheme.
    pub fn endpoint(&self) -> &str {
        &self.endpoint
    }
}

impl Drop for TestPdControlServer {
    fn drop(&mut self) {
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            let _ = shutdown_tx.send(());
        }
    }
}

#[derive(Deserialize)]
struct TransferRegionRequest {
    name: String,
    region_id: u64,
    to_store_ids: Vec<u64>,
    peer_roles: Vec<String>,
}

async fn handle_request(
    request: Request<Body>,
    pd_client: Arc<TestPdClient>,
    keyspace_lookup: Arc<KeyspaceLookup>,
    transfer_lock: Arc<Mutex<()>>,
) -> Result<Response<Body>, Infallible> {
    const KEYSPACE_PREFIX: &str = "/pd/api/v2/keyspaces/";
    const LABEL_RULE_PATH: &str = "/pd/api/v1/config/region-label/rule";
    const OPERATORS_PATH: &str = "/pd/api/v1/operators";

    let method = request.method().clone();
    let path = request.uri().path().to_string();
    let response = if method == Method::GET && path.starts_with(KEYSPACE_PREFIX) {
        let keyspace_name = &path[KEYSPACE_PREFIX.len()..];
        match keyspace_lookup(keyspace_name) {
            Some(keyspace_id) => json_response(
                StatusCode::OK,
                serde_json::json!({
                    "id": keyspace_id,
                    "name": keyspace_name,
                    "state": "ENABLED",
                }),
            ),
            None => text_response(StatusCode::NOT_FOUND, "keyspace not found"),
        }
    } else if method == Method::POST && path == LABEL_RULE_PATH {
        match hyper::body::to_bytes(request.into_body()).await {
            Ok(body) => match serde_json::from_slice::<pd_client::pd_control::LabelRule>(&body) {
                Ok(rule) => match pd_client.set_region_label_rule(&rule) {
                    Ok(()) => json_response(StatusCode::OK, serde_json::json!("ok")),
                    Err(err) => text_response(StatusCode::BAD_REQUEST, &err),
                },
                Err(err) => text_response(
                    StatusCode::BAD_REQUEST,
                    &format!("invalid Region label rule: {err}"),
                ),
            },
            Err(err) => text_response(
                StatusCode::BAD_REQUEST,
                &format!("failed to read Region label rule: {err}"),
            ),
        }
    } else if method == Method::GET && path == OPERATORS_PATH {
        json_response(StatusCode::OK, serde_json::json!([]))
    } else if method == Method::POST && path == OPERATORS_PATH {
        match hyper::body::to_bytes(request.into_body()).await {
            Ok(body) => match serde_json::from_slice::<TransferRegionRequest>(&body) {
                Ok(transfer) => {
                    let pd_client = pd_client.clone();
                    let result = tokio::task::spawn_blocking(move || {
                        let _guard = transfer_lock.lock().unwrap();
                        apply_transfer_region(&pd_client, transfer)
                    })
                    .await;
                    match result {
                        Ok(Ok(())) => json_response(StatusCode::OK, serde_json::json!("ok")),
                        Ok(Err(err)) => text_response(StatusCode::INTERNAL_SERVER_ERROR, &err),
                        Err(err) => text_response(
                            StatusCode::INTERNAL_SERVER_ERROR,
                            &format!("transfer-region task failed: {err}"),
                        ),
                    }
                }
                Err(err) => text_response(
                    StatusCode::BAD_REQUEST,
                    &format!("invalid transfer-region request: {err}"),
                ),
            },
            Err(err) => text_response(
                StatusCode::BAD_REQUEST,
                &format!("failed to read transfer-region request: {err}"),
            ),
        }
    } else {
        text_response(
            StatusCode::NOT_FOUND,
            &format!("unsupported test PD control request: {method} {path}"),
        )
    };
    Ok(response)
}

fn apply_transfer_region(
    pd_client: &TestPdClient,
    transfer: TransferRegionRequest,
) -> Result<(), String> {
    if transfer.name != "transfer-region" {
        return Err(format!("unsupported operator {}", transfer.name));
    }
    if transfer.to_store_ids.len() != transfer.peer_roles.len() {
        return Err("transfer-region Store IDs and peer roles have different lengths".to_string());
    }
    let desired_leaders = transfer
        .to_store_ids
        .iter()
        .zip(&transfer.peer_roles)
        .filter_map(|(&store_id, role)| (role == "leader").then_some(store_id))
        .collect::<Vec<_>>();
    let [desired_leader_store] = desired_leaders.as_slice() else {
        return Err(format!(
            "transfer-region requires exactly one leader, got {:?}",
            desired_leaders
        ));
    };

    for _ in 0..3 {
        let (region, _) =
            futures::executor::block_on(pd_client.get_region_leader_by_id(transfer.region_id))
                .map_err(|err| format!("get Region {} failed: {err}", transfer.region_id))?
                .ok_or_else(|| format!("Region {} not found", transfer.region_id))?;

        for &store_id in &transfer.to_store_ids {
            if region
                .get_peers()
                .iter()
                .any(|peer| peer.get_store_id() == store_id)
            {
                continue;
            }
            let mut peer = metapb::Peer::default();
            peer.set_id(
                pd_client
                    .alloc_id()
                    .map_err(|err| format!("allocate peer ID failed: {err}"))?,
            );
            peer.set_store_id(store_id);
            peer.set_role(PeerRole::Voter);
            pd_client.add_peer(transfer.region_id, peer);
            wait_for_region_state(
                pd_client,
                transfer.region_id,
                &format!("add peer on Store {store_id}"),
                |region, _| {
                    region
                        .get_peers()
                        .iter()
                        .any(|peer| peer.get_store_id() == store_id)
                },
            )?;
        }

        let (region, leader) =
            futures::executor::block_on(pd_client.get_region_leader_by_id(transfer.region_id))
                .map_err(|err| format!("refresh Region {} failed: {err}", transfer.region_id))?
                .ok_or_else(|| format!("Region {} disappeared", transfer.region_id))?;
        let desired_leader = region
            .get_peers()
            .iter()
            .find(|peer| peer.get_store_id() == *desired_leader_store)
            .cloned()
            .ok_or_else(|| {
                format!(
                    "Region {} has no peer on desired leader Store {}",
                    transfer.region_id, desired_leader_store
                )
            })?;
        if leader.get_store_id() != *desired_leader_store {
            pd_client.transfer_leader(
                transfer.region_id,
                desired_leader.clone(),
                vec![desired_leader],
            );
            wait_for_region_state(
                pd_client,
                transfer.region_id,
                &format!("transfer leader to Store {desired_leader_store}"),
                |_, leader| leader.get_store_id() == *desired_leader_store,
            )?;
        }

        let (region, _) =
            futures::executor::block_on(pd_client.get_region_leader_by_id(transfer.region_id))
                .map_err(|err| {
                    format!(
                        "refresh Region {} before removing peers failed: {err}",
                        transfer.region_id
                    )
                })?
                .ok_or_else(|| format!("Region {} disappeared", transfer.region_id))?;
        for peer in region.get_peers() {
            if !transfer.to_store_ids.contains(&peer.get_store_id()) {
                let store_id = peer.get_store_id();
                pd_client.remove_peer(transfer.region_id, peer.clone());
                wait_for_region_state(
                    pd_client,
                    transfer.region_id,
                    &format!("remove peer from Store {store_id}"),
                    |region, _| {
                        region
                            .get_peers()
                            .iter()
                            .all(|peer| peer.get_store_id() != store_id)
                    },
                )?;
            }
        }

        let (region, leader) =
            futures::executor::block_on(pd_client.get_region_leader_by_id(transfer.region_id))
                .map_err(|err| {
                    format!(
                        "verify Region {} placement failed: {err}",
                        transfer.region_id
                    )
                })?
                .ok_or_else(|| format!("Region {} disappeared", transfer.region_id))?;
        let mut actual_stores = region
            .get_peers()
            .iter()
            .filter(|peer| peer.get_role() == PeerRole::Voter)
            .map(|peer| peer.get_store_id())
            .collect::<Vec<_>>();
        let mut desired_stores = transfer.to_store_ids.clone();
        actual_stores.sort_unstable();
        desired_stores.sort_unstable();
        if actual_stores == desired_stores && leader.get_store_id() == *desired_leader_store {
            return Ok(());
        }
    }

    Err(format!(
        "Region {} placement did not converge to {:?}",
        transfer.region_id, transfer.to_store_ids
    ))
}

fn wait_for_region_state(
    pd_client: &TestPdClient,
    region_id: u64,
    operation: &str,
    predicate: impl Fn(&metapb::Region, &metapb::Peer) -> bool,
) -> Result<(metapb::Region, metapb::Peer), String> {
    let mut last_region = None;
    for _ in 0..300 {
        match futures::executor::block_on(pd_client.get_region_leader_by_id(region_id)) {
            Ok(Some((region, leader))) => {
                if predicate(&region, &leader) {
                    return Ok((region, leader));
                }
                last_region = Some(region);
            }
            Ok(None) => {
                return Err(format!(
                    "Region {region_id} disappeared while waiting to {operation}"
                ));
            }
            Err(err) => {
                return Err(format!(
                    "get Region {region_id} while waiting to {operation} failed: {err}"
                ));
            }
        }
        sleep(Duration::from_millis(10));
    }
    Err(format!(
        "Region {region_id} did not finish {operation}, last state: {last_region:?}"
    ))
}

fn json_response(status: StatusCode, value: serde_json::Value) -> Response<Body> {
    Response::builder()
        .status(status)
        .header(http::header::CONTENT_TYPE, "application/json")
        .body(Body::from(value.to_string()))
        .unwrap()
}

fn text_response(status: StatusCode, message: &str) -> Response<Body> {
    Response::builder()
        .status(status)
        .body(Body::from(message.to_string()))
        .unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_schedule_deny_rule_is_visible_to_test_pd_client() {
        let runtime = Runtime::new().unwrap();
        let pd_client = Arc::new(TestPdClient::new(1, false));
        let rule = pd_client::pd_control::LabelRule {
            id: "deny".to_string(),
            labels: vec![pd_client::pd_control::RegionLabel {
                key: "schedule".to_string(),
                value: "deny".to_string(),
                ttl: "1m".to_string(),
            }],
            rule_type: "key-range".to_string(),
            data: vec![pd_client::pd_control::KeyRangeRule {
                start_key: hex::encode(b"a"),
                end_key: hex::encode(b"c"),
            }],
        };
        let request = Request::post("/pd/api/v1/config/region-label/rule")
            .body(Body::from(serde_json::to_vec(&rule).unwrap()))
            .unwrap();
        let keyspace_lookup: Arc<KeyspaceLookup> = Arc::new(|_| None);
        let response = runtime
            .block_on(handle_request(
                request,
                pd_client.clone(),
                keyspace_lookup,
                Arc::new(Mutex::new(())),
            ))
            .unwrap();

        let mut region = metapb::Region::default();
        region.set_start_key(b"b".to_vec());
        region.set_end_key(b"d".to_vec());
        assert_eq!(response.status(), StatusCode::OK);
        assert!(pd_client.is_scheduling_denied(&region));
    }
}