//! gRPC 请求日志：包装 tonic 的服务栈，为每次调用打印一行结构化日志。
//!
//! gRPC 的错误码放在 HTTP/2 的 trailer 里，所以响应体也必须包一层，等到 trailer
//! 到达后才能拿到真实的 `grpc-status`。这里同时兼容把状态写在响应头的情况
//! （反射、健康检查与部分异常路径）。

use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Instant;

use bytes::Bytes;
use envelope::{LEGACY_REQUEST_ID_HEADER, REQUEST_ID_HEADER, new_request_id, validate_request_id};
use http::header::{HeaderName, HeaderValue};
use http::{HeaderMap, Request, Response};
use http_body::{Body, Frame, SizeHint};
use tower::{Layer, Service};
use tracing::Instrument;

/// 调用方身份头（可选，缺省表示匿名调用）。
const SERVICE_ID_HEADER: &str = "gitdata-service-id";
const GRPC_STATUS: &str = "grpc-status";

/// gRPC 请求日志层，通过 `tonic::transport::Server::layer` 装配。
#[derive(Clone, Copy, Debug, Default)]
pub struct GrpcRequestLog;

impl<S> Layer<S> for GrpcRequestLog {
    type Service = GrpcRequestLogService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        GrpcRequestLogService { inner }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct GrpcRequestLogService<S> {
    inner: S,
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for GrpcRequestLogService<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
    S::Future: Send + 'static,
    ResBody: Body<Data = Bytes> + Send + 'static,
{
    type Response = Response<LoggedBody<ResBody>>;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(context)
    }

    fn call(&mut self, mut request: Request<ReqBody>) -> Self::Future {
        let request_id = {
            let headers = request.headers();
            headers
                .get(REQUEST_ID_HEADER)
                .or_else(|| headers.get(LEGACY_REQUEST_ID_HEADER))
                .and_then(|value| value.to_str().ok())
                .and_then(validate_request_id)
                .unwrap_or_else(new_request_id)
        };
        if let Ok(value) = HeaderValue::from_str(&request_id) {
            request
                .headers_mut()
                .insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
        }
        let call = Call::new(&request, request_id.clone());
        let span = tracing::info_span!(
            "grpc.server.request",
            "otel.kind" = "server",
            "rpc.system" = "grpc",
            "rpc.service" = %call.service,
            "rpc.method" = %call.method,
            "gitdata.service_id" = %call.caller,
            request_id = %request_id,
        );
        let traceparent = request
            .headers()
            .get(crate::TRACEPARENT)
            .and_then(|value| value.to_str().ok())
            .map(str::to_owned);
        crate::set_parent_from_traceparent(&span, traceparent.as_deref());

        let started = Instant::now();
        let future = self.inner.call(request);
        let context_id = request_id.clone();
        Box::pin(envelope::context::with_request_id(
            Some(context_id),
            async move {
                let response = future.await?;
                let (mut parts, body) = response.into_parts();
                if let Ok(value) = HeaderValue::from_str(&request_id) {
                    parts
                        .headers
                        .insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
                }
                let header_status = status_of(&parts.headers);
                // 错误响应（含未认证、反射失败）不带响应体，trailer 永远不会被轮询到，
                // 所以结束流的情况必须在包装时就记一笔。
                let ended = body.is_end_stream();
                let mut logged = LoggedBody {
                    inner: Box::pin(body),
                    call,
                    started,
                    header_status,
                    logged: false,
                };
                if ended {
                    logged.finish(header_status);
                }
                Ok(Response::from_parts(parts, logged))
            }
            .instrument(span),
        ))
    }
}

/// 一次调用的静态信息，在请求进栈时抓取。
#[derive(Clone, Debug)]
struct Call {
    method: String,
    service: String,
    request_id: String,
    caller: String,
}

impl Call {
    fn new<B>(request: &Request<B>, request_id: String) -> Self {
        let path = request.uri().path();
        let (service, method) = match path.rsplit_once('/') {
            Some((service, method)) => (
                service.trim_start_matches('/').to_owned(),
                method.to_owned(),
            ),
            None => (String::new(), path.to_owned()),
        };
        Self {
            method,
            service,
            request_id,
            caller: header(request.headers(), SERVICE_ID_HEADER),
        }
    }

    fn emit(&self, status: Option<i64>, elapsed_ms: f64) {
        let code = status.unwrap_or(-1);
        let duration_ms = (elapsed_ms * 1_000.0).round() / 1_000.0;
        let caller = if self.caller.is_empty() {
            "anonymous"
        } else {
            self.caller.as_str()
        };
        let request_id = if self.request_id.is_empty() {
            "-"
        } else {
            self.request_id.as_str()
        };
        if code == 0 {
            tracing::info!(
                "rpc.service" = %self.service,
                "rpc.method" = %self.method,
                "rpc.system" = "grpc",
                "rpc.grpc.status_code" = code,
                "duration_ms" = duration_ms,
                "gitdata.service_id" = caller,
                "request_id" = request_id,
                "grpc request"
            );
        } else {
            tracing::warn!(
                "rpc.service" = %self.service,
                "rpc.method" = %self.method,
                "rpc.system" = "grpc",
                "rpc.grpc.status_code" = code,
                "duration_ms" = duration_ms,
                "gitdata.service_id" = caller,
                "request_id" = request_id,
                "grpc request failed"
            );
        }
    }
}

/// 透传响应体，等 trailer 出来时补上一条日志。
pub struct LoggedBody<B> {
    inner: Pin<Box<B>>,
    call: Call,
    started: Instant,
    header_status: Option<i64>,
    logged: bool,
}

impl<B: Body<Data = Bytes>> LoggedBody<B> {
    fn finish(&mut self, status: Option<i64>) {
        if self.logged {
            return;
        }
        self.logged = true;
        let elapsed_ms = self.started.elapsed().as_secs_f64() * 1_000.0;
        self.call.emit(status, elapsed_ms);
    }
}

impl<B: Body<Data = Bytes>> Body for LoggedBody<B> {
    type Data = Bytes;
    type Error = B::Error;

    fn poll_frame(
        self: Pin<&mut Self>,
        context: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        // `Pin<Box<B>>` 让整体始终 Unpin，无需 unsafe 投影。
        let this = self.get_mut();
        let polled = this.inner.as_mut().poll_frame(context);
        match &polled {
            Poll::Ready(Some(Ok(frame))) => {
                if let Some(trailers) = frame.trailers_ref() {
                    this.finish(status_of(trailers).or(this.header_status));
                }
            }
            Poll::Ready(Some(Err(_))) => this.finish(this.header_status),
            Poll::Ready(None) => this.finish(this.header_status),
            Poll::Pending => {}
        }
        polled
    }

    fn is_end_stream(&self) -> bool {
        self.inner.is_end_stream()
    }

    fn size_hint(&self) -> SizeHint {
        self.inner.size_hint()
    }
}

impl<B> fmt::Debug for LoggedBody<B> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("LoggedBody")
            .field("call", &self.call)
            .field("header_status", &self.header_status)
            .field("logged", &self.logged)
            .finish_non_exhaustive()
    }
}

fn header(headers: &HeaderMap, name: &str) -> String {
    headers
        .get(name)
        .and_then(|value| value.to_str().ok())
        .unwrap_or_default()
        .to_owned()
}

fn status_of(headers: &HeaderMap) -> Option<i64> {
    headers
        .get(GRPC_STATUS)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.trim().parse().ok())
}