rs-zero 0.2.6

use std::fmt;

use serde_json::Value;
use tracing::{Event, Subscriber};
use tracing_subscriber::{
    fmt::{
        FmtContext, FormattedFields,
        format::{DefaultFields, Format, FormatEvent, FormatFields, Json, JsonFields, Writer},
    },
    registry::LookupSpan,
};

use crate::core::logging::redaction::{RedactionConfig, is_sensitive_key, redact_text};

/// Text event formatter that redacts event and span fields before output.
#[derive(Debug, Clone)]
pub(crate) struct RedactingTextFormat {
    inner: Format,
    redaction: RedactionConfig,
}

impl RedactingTextFormat {
    pub(crate) fn new(inner: Format, redaction: RedactionConfig) -> Self {
        Self { inner, redaction }
    }
}

impl<S> FormatEvent<S, DefaultFields> for RedactingTextFormat
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    fn format_event(
        &self,
        ctx: &FmtContext<'_, S, DefaultFields>,
        mut writer: Writer<'_>,
        event: &Event<'_>,
    ) -> fmt::Result {
        let mut line = String::new();
        self.inner
            .format_event(ctx, Writer::new(&mut line), event)?;
        writer.write_str(&redact_text(&line, &self.redaction))
    }
}

/// JSON event formatter that redacts parsed JSON before final output.
#[derive(Debug, Clone)]
pub(crate) struct RedactingJsonFormat {
    inner: Format<Json>,
    redaction: RedactionConfig,
}

impl RedactingJsonFormat {
    pub(crate) fn new(inner: Format<Json>, redaction: RedactionConfig) -> Self {
        Self { inner, redaction }
    }
}

impl<S> FormatEvent<S, RedactingJsonFields> for RedactingJsonFormat
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    fn format_event(
        &self,
        ctx: &FmtContext<'_, S, RedactingJsonFields>,
        mut writer: Writer<'_>,
        event: &Event<'_>,
    ) -> fmt::Result {
        let mut line = String::new();
        self.inner
            .format_event(ctx, Writer::new(&mut line), event)?;
        let Ok(mut value) = serde_json::from_str::<Value>(line.trim_end()) else {
            return writer.write_str(&redact_text(&line, &self.redaction));
        };

        redact_json_value(None, &mut value, &self.redaction);
        serde_json::to_writer(JsonFmtWriter(writer.by_ref()), &value).map_err(|_| fmt::Error)?;
        writer.write_char('\n')
    }
}

/// JSON field formatter that stores redacted span fields in subscriber state.
#[derive(Debug, Clone)]
pub(crate) struct RedactingJsonFields {
    redaction: RedactionConfig,
}

impl RedactingJsonFields {
    pub(crate) fn new(redaction: RedactionConfig) -> Self {
        Self { redaction }
    }
}

impl<'writer> FormatFields<'writer> for RedactingJsonFields {
    fn format_fields<R>(&self, mut writer: Writer<'writer>, fields: R) -> fmt::Result
    where
        R: tracing_subscriber::field::RecordFields,
    {
        let mut output = String::new();
        JsonFields::new().format_fields(Writer::new(&mut output), fields)?;
        write_redacted_json_object(&mut writer, &output, &self.redaction)
    }

    fn add_fields(
        &self,
        current: &'writer mut FormattedFields<Self>,
        fields: &tracing::span::Record<'_>,
    ) -> fmt::Result {
        let existing = current.as_str();
        let mut output = String::new();
        if existing.is_empty() {
            JsonFields::new().format_fields(Writer::new(&mut output), fields)?;
        } else {
            let mut merged: serde_json::Map<String, Value> =
                serde_json::from_str(existing).map_err(|_| fmt::Error)?;
            let mut added = String::new();
            JsonFields::new().format_fields(Writer::new(&mut added), fields)?;
            let Value::Object(added) =
                serde_json::from_str::<Value>(&added).map_err(|_| fmt::Error)?
            else {
                return Err(fmt::Error);
            };
            merged.extend(added);
            output = Value::Object(merged).to_string();
        }

        let mut redacted = String::new();
        write_redacted_json_object(&mut redacted, &output, &self.redaction)?;
        current.fields = redacted;
        Ok(())
    }
}

fn write_redacted_json_object<W: fmt::Write>(
    writer: &mut W,
    output: &str,
    redaction: &RedactionConfig,
) -> fmt::Result {
    let mut value = serde_json::from_str::<Value>(output).map_err(|_| fmt::Error)?;
    redact_json_value(None, &mut value, redaction);
    writer.write_str(&value.to_string())
}

fn redact_json_value(key: Option<&str>, value: &mut Value, config: &RedactionConfig) {
    if key.is_some_and(|key| is_sensitive_key(key, config)) {
        *value = Value::String(config.replacement.clone());
        return;
    }

    match value {
        Value::String(text) => *text = redact_text(text, config),
        Value::Array(items) => {
            for item in items {
                redact_json_value(None, item, config);
            }
        }
        Value::Object(fields) => {
            for (field, value) in fields {
                redact_json_value(Some(field), value, config);
            }
        }
        Value::Null | Value::Bool(_) | Value::Number(_) => {}
    }
}

struct JsonFmtWriter<'a>(Writer<'a>);

impl std::io::Write for JsonFmtWriter<'_> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let text = std::str::from_utf8(buf)
            .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?;
        self.0
            .write_str(text)
            .map_err(|_| std::io::Error::other("format writer failed"))?;
        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}