knut-bifrost 0.1.8

//! **Apache Doris** [`TableSink`] — delivery over Stream Load.
//!
//! Feature `doris` (off by default). Enabling it adds only `reqwest`, which the
//! `airflow` feature already pulls, so the workspace dependency tree does not
//! grow. See `.nornir/doris-sink-design.md`.
//!
//! # The transport is a seam on purpose
//!
//! Every correctness property this module has — label derivation, replay
//! rejection, delete-mode refusal, redirect diagnosis, schema refusal, chunking
//! — is testable without a Doris cluster, because [`DorisSink`] talks to a
//! [`StreamLoadTransport`] rather than to an HTTP client. [`ReqwestTransport`]
//! is the real one; the tests use a fake that encodes the *documented* protocol.
//! Only throughput needs a live cluster.
//!
//! # Topology, and the error that looks like auth and is not
//!
//! Stream Load answers on the **FE** (`http_port`, default 8030) at
//! `PUT /api/{db}/{table}/_stream_load`, and the FE replies **307** with a
//! `Location` naming a **BE** (`webserver_port`, default 8040) that does the
//! actual work. So the loader must reach the FE **and every BE it may be
//! redirected to**, and the redirect carries the BE's *configured* address,
//! which on a split-addressing cluster need not be routable from here.
//!
//! **Measured against the live cluster (Doris 4.0.8):** when the FE answers but
//! the BEs are unreachable, the load does not fail — it **hangs** after the 307
//! until something times out. `/api/health` on the FE stays 200 the whole time,
//! so nothing looks wrong. Without an explicit timeout a single stuck window
//! blocks the stream forever with no error at all.
//!
//! So this module (a) follows the redirect itself, automatic following disabled
//! on the real client, and (b) puts a hard connect **and** request timeout on
//! every hop, so the stall becomes a reportable error naming the BE.
//!
//! # `Expect: 100-continue` is mandatory
//!
//! Also measured: without it Doris answers
//! `{"status":"FAILED","msg":"There is no 100-continue header"}`. This is a
//! nasty one because `curl` omits the header for small bodies and sends it for
//! large ones — so small windows fail while large ones succeed. The sink sets it
//! on every load, unconditionally, and a test asserts it is on every PUT.

use std::collections::BTreeSet;

use arrow_array::RecordBatch;
use arrow_schema::Schema;
use async_trait::async_trait;

use crate::{
    error::{BifrostError, Result},
    sink::table::{AckToken, LoadAck, TableSink, TableWindow, stream_load_label},
};

// ── the transport seam ──────────────────────────────────────────────────────

/// One HTTP exchange, reduced to what Stream Load actually needs.
#[derive(Debug, Clone)]
pub struct StreamLoadRequest {
    /// `"PUT"` for a load, `"GET"` for the schema probe.
    pub method: &'static str,
    pub url: String,
    /// Header name/value pairs, in the order the sink set them.
    pub headers: Vec<(String, String)>,
    /// Shared, not copied: following the FE→BE redirect re-sends the same bytes.
    pub body: bytes::Bytes,
}

impl StreamLoadRequest {
    /// The value of a header, case-insensitively.
    pub fn header(&self, name: &str) -> Option<&str> {
        self.headers
            .iter()
            .find(|(k, _)| k.eq_ignore_ascii_case(name))
            .map(|(_, v)| v.as_str())
    }
}

/// What came back.
#[derive(Debug, Clone)]
pub struct StreamLoadResponse {
    pub status: u16,
    /// The `Location` header of a 3xx. Redirects are **not** followed by the
    /// transport — [`DorisSink`] follows them so it can diagnose an unreachable
    /// backend precisely.
    pub location: Option<String>,
    pub body: String,
}

impl StreamLoadResponse {
    pub fn ok(body: impl Into<String>) -> Self {
        Self {
            status: 200,
            location: None,
            body: body.into(),
        }
    }

    pub fn redirect(location: impl Into<String>) -> Self {
        Self {
            status: 307,
            location: Some(location.into()),
            body: String::new(),
        }
    }

    pub fn is_redirect(&self) -> bool {
        matches!(self.status, 301 | 302 | 303 | 307 | 308) && self.location.is_some()
    }
}

/// How bytes get to Doris. Implemented for real by [`ReqwestTransport`], and by
/// a fake in the tests.
#[async_trait]
pub trait StreamLoadTransport: Send + Sync {
    async fn send(&self, req: &StreamLoadRequest) -> Result<StreamLoadResponse>;
}

// ── configuration ───────────────────────────────────────────────────────────

/// What a `DELETE` row means for this destination.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LoadMode {
    /// Duplicate-key table, append-only source. A `DELETE` row is a **hard
    /// error**: on an append-only source an unexpected delete means the premise
    /// of the deployment is wrong, and a loud failure beats a silent divergence
    /// between the lake and the engine.
    Append,
    /// Unique-key table. Delete rows are shipped as a second Stream Load with
    /// `merge_type: DELETE`, which is how Doris sets its hidden
    /// `__DORIS_DELETE_SIGN__` column. Not available on duplicate-key or
    /// aggregate-key tables — hence the two modes.
    Upsert,
}

/// Stream Load wire encoding. Doris accepts `csv`, `json`, `parquet`, `orc` and
/// `arrow`; only JSON is implemented, because it needs no dependency bifröst
/// does not already have and is unambiguous about nulls and types. The enum
/// exists so parquet/arrow can be added as a throughput follow-up without
/// changing any caller — whether they win is an empirical question, not an
/// assumption.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum LoadFormat {
    #[default]
    Json,
    /// Arrow IPC stream (`format: arrow`, Doris ≥ 2.1): the batches go as they
    /// are, no text encoding on this side and no parsing on the BE's.
    Arrow,
}

impl LoadFormat {
    const fn header_value(self) -> &'static str {
        match self {
            LoadFormat::Json => "json",
            LoadFormat::Arrow => "arrow",
        }
    }
}

/// Where to load, as what, with which credentials.
#[derive(Debug, Clone)]
pub struct DorisConfig {
    /// FE base URL including the HTTP port, e.g. `http://doris-fe:8030`.
    pub fe_url: String,
    /// Destination database and table (NOT the source table — that is the
    /// window's, and it is what the cursor is keyed on).
    pub db: String,
    pub table: String,
    pub user: String,
    pub password: String,
    pub mode: LoadMode,
    pub format: LoadFormat,
    /// Byte budget for one Stream Load body. The BE caps a single load at
    /// `streaming_load_max_mb` (default 10 GB) and Doris's own guidance is to
    /// stay well under it — a dropped connection restarts the whole request.
    /// Default 64 MiB.
    pub max_chunk_bytes: usize,
    /// TCP connect timeout for each hop. An unreachable BE after the 307 is a
    /// **hang**, not an error (measured), so this is what turns a silently stuck
    /// stream into a reportable failure. Default 10 s.
    pub connect_timeout: std::time::Duration,
    /// Whole-request timeout for each hop, the load body included. Must exceed
    /// the time a real BE needs to ingest one chunk. Default 300 s.
    pub request_timeout: std::time::Duration,
    /// Skip the FE and PUT straight at a known-healthy BE. The escape hatch for
    /// a cluster whose FE hands out addresses that are not routable from here.
    pub be_url: Option<String>,
    /// Chunks of one window loaded concurrently (each its own label). Default 4.
    pub max_parallel_loads: usize,
}

/// Default body budget per Stream Load: 64 MiB. Two orders of magnitude below
/// the BE cap, so a retry is cheap.
pub const DEFAULT_MAX_CHUNK_BYTES: usize = 64 * 1024 * 1024;

/// Default chunks of one window in flight at once.
pub const DEFAULT_MAX_PARALLEL_LOADS: usize = 4;

/// Default TCP connect timeout per hop. See [`DorisConfig::connect_timeout`] —
/// an unreachable BE hangs rather than refusing, so this bound is the only thing
/// standing between a routing mistake and a permanently stalled stream.
pub const DEFAULT_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);

/// Default whole-request timeout per hop.
pub const DEFAULT_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);

impl DorisConfig {
    /// Credentials from the environment: `DORIS_USER`, `DORIS_PASSWORD`.
    ///
    /// There is no constructor that takes a password as a flag, and there will
    /// not be one: flags land in `/proc/<pid>/cmdline`, readable by every user
    /// on the host, and in shell history and process listings.
    pub fn from_env(
        fe_url: impl Into<String>,
        db: impl Into<String>,
        table: impl Into<String>,
    ) -> Result<Self> {
        let user = std::env::var("DORIS_USER")
            .map_err(|_| BifrostError::Other("DORIS_USER is not set".into()))?;
        let password = std::env::var("DORIS_PASSWORD")
            .map_err(|_| BifrostError::Other("DORIS_PASSWORD is not set".into()))?;
        Ok(Self {
            fe_url: fe_url.into(),
            db: db.into(),
            table: table.into(),
            user,
            password,
            mode: LoadMode::Append,
            format: LoadFormat::default(),
            max_chunk_bytes: DEFAULT_MAX_CHUNK_BYTES,
            connect_timeout: DEFAULT_CONNECT_TIMEOUT,
            request_timeout: DEFAULT_REQUEST_TIMEOUT,
            be_url: None,
            max_parallel_loads: DEFAULT_MAX_PARALLEL_LOADS,
        })
    }

    /// For tests and for callers that source credentials elsewhere.
    pub fn new(
        fe_url: impl Into<String>,
        db: impl Into<String>,
        table: impl Into<String>,
        user: impl Into<String>,
        password: impl Into<String>,
    ) -> Self {
        Self {
            fe_url: fe_url.into(),
            db: db.into(),
            table: table.into(),
            user: user.into(),
            password: password.into(),
            mode: LoadMode::Append,
            format: LoadFormat::default(),
            max_chunk_bytes: DEFAULT_MAX_CHUNK_BYTES,
            connect_timeout: DEFAULT_CONNECT_TIMEOUT,
            request_timeout: DEFAULT_REQUEST_TIMEOUT,
            be_url: None,
            max_parallel_loads: DEFAULT_MAX_PARALLEL_LOADS,
        }
    }

    #[must_use]
    pub fn with_timeouts(
        mut self,
        connect: std::time::Duration,
        request: std::time::Duration,
    ) -> Self {
        self.connect_timeout = connect;
        self.request_timeout = request;
        self
    }

    #[must_use]
    pub fn with_mode(mut self, mode: LoadMode) -> Self {
        self.mode = mode;
        self
    }

    #[must_use]
    pub fn with_max_chunk_bytes(mut self, n: usize) -> Self {
        self.max_chunk_bytes = n.max(1);
        self
    }

    #[must_use]
    pub fn with_max_parallel_loads(mut self, n: usize) -> Self {
        self.max_parallel_loads = n.max(1);
        self
    }

    #[must_use]
    pub fn with_be_url(mut self, url: impl Into<String>) -> Self {
        self.be_url = Some(url.into());
        self
    }

    fn stream_load_url(&self, base: &str) -> String {
        format!(
            "{}/api/{}/{}/_stream_load",
            base.trim_end_matches('/'),
            self.db,
            self.table
        )
    }

    fn schema_url(&self) -> String {
        format!(
            "{}/api/{}/{}/_schema",
            self.fe_url.trim_end_matches('/'),
            self.db,
            self.table
        )
    }

    /// HTTP Basic, precomputed. Doris Stream Load takes credentials this way.
    fn basic_auth(&self) -> String {
        format!(
            "Basic {}",
            base64_std(format!("{}:{}", self.user, self.password).as_bytes())
        )
    }
}

// ── the sink ────────────────────────────────────────────────────────────────

/// A [`TableSink`] that delivers CDC windows into a Doris table over Stream
/// Load.
pub struct DorisSink<T: StreamLoadTransport> {
    cfg: DorisConfig,
    transport: T,
    probed: bool,
}

impl<T: StreamLoadTransport> DorisSink<T> {
    pub fn new(cfg: DorisConfig, transport: T) -> Self {
        Self {
            cfg,
            transport,
            probed: false,
        }
    }

    pub fn config(&self) -> &DorisConfig {
        &self.cfg
    }

    /// Column names the destination table actually has, from the FE's
    /// `_schema` endpoint.
    async fn remote_columns(&self) -> Result<BTreeSet<String>> {
        let req = StreamLoadRequest {
            method: "GET",
            url: self.cfg.schema_url(),
            headers: vec![("Authorization".into(), self.cfg.basic_auth())],
            body: bytes::Bytes::new(),
        };
        let resp = self.transport.send(&req).await?;
        if resp.status != 200 {
            return Err(BifrostError::Other(format!(
                "doris: schema probe of {}.{} returned HTTP {} — {}",
                self.cfg.db, self.cfg.table, resp.status, resp.body
            )));
        }
        let v: serde_json::Value = serde_json::from_str(&resp.body).map_err(|e| {
            BifrostError::Other(format!("doris: schema probe returned non-JSON: {e}"))
        })?;
        let props = v
            .get("data")
            .and_then(|d| d.get("properties"))
            .and_then(serde_json::Value::as_array)
            .ok_or_else(|| {
                BifrostError::Other(format!(
                    "doris: schema probe of {}.{} has no data.properties — {}",
                    self.cfg.db, self.cfg.table, resp.body
                ))
            })?;
        Ok(props
            .iter()
            .filter_map(|p| p.get("name").and_then(serde_json::Value::as_str))
            .map(str::to_owned)
            .collect())
    }

    /// PUT one already-encoded chunk, following the FE→BE redirect by hand.
    ///
    /// Returns `(rows_loaded, was_duplicate)`.
    async fn put_chunk(
        &self,
        label: &str,
        merge_type: &str,
        columns: &str,
        body: bytes::Bytes,
    ) -> Result<(usize, bool)> {
        let base = self.cfg.be_url.as_deref().unwrap_or(&self.cfg.fe_url);
        let mut url = self.cfg.stream_load_url(base);
        let mut headers = vec![
            // The body's columns, by name. MEASURED on Doris 4.1.4 (loki rig): an
            // arrow body carrying only the key (a delete) without this header fails
            // in the BE with "insert_range_from … out of bound (data.size() = 0)".
            ("columns".to_string(), columns.to_string()),
            ("Authorization".to_string(), self.cfg.basic_auth()),
            // MANDATORY, and measured: without it Doris answers
            // {"status":"FAILED","msg":"There is no 100-continue header"}.
            // curl omits it for small bodies, so leaving it out makes SMALL
            // windows fail while large ones succeed — a size-dependent bug that
            // would look like data corruption.
            ("Expect".to_string(), "100-continue".to_string()),
            ("label".to_string(), label.to_string()),
            (
                "format".to_string(),
                self.cfg.format.header_value().to_string(),
            ),
            ("merge_type".to_string(), merge_type.to_string()),
        ];
        if self.cfg.format == LoadFormat::Json {
            headers.push(("strip_outer_array".to_string(), "true".to_string()));
        }
        // A dropped row is the worst failure a change feed can have, so do not
        // let Doris silently filter any: any unparseable row fails the load.
        headers.push(("max_filter_ratio".to_string(), "0".to_string()));

        let mut hops = 0_u8;
        let resp = loop {
            let req = StreamLoadRequest {
                method: "PUT",
                url: url.clone(),
                headers: headers.clone(),
                body: body.clone(),
            };
            let resp = self.transport.send(&req).await.map_err(|e| {
                if hops > 0 {
                    BifrostError::Other(format!(
                        "doris: the FE redirected this load to {url} and the request did not \
                         complete — this is a NETWORK REACHABILITY failure. Measured \
                         behaviour: an unreachable BE does not refuse the connection, it \
                         HANGS until the timeout ({}s connect / {}s request), while the FE \
                         keeps answering /api/health with 200. The loader must reach the FE \
                         ({}) AND every BE it can be redirected to on :8040. Underlying \
                         error: {e}",
                        self.cfg.connect_timeout.as_secs(),
                        self.cfg.request_timeout.as_secs(),
                        self.cfg.fe_url
                    ))
                } else {
                    e
                }
            })?;
            if !resp.is_redirect() {
                break resp;
            }
            hops += 1;
            if hops > 4 {
                return Err(BifrostError::Other(format!(
                    "doris: more than 4 redirects starting at {}",
                    self.cfg.stream_load_url(base)
                )));
            }
            url = resp.location.clone().expect("is_redirect checked location");
            tracing::debug!(
                target: "bifrost::doris",
                event = "stream_load_redirect",
                to = %url,
                "FE redirected the load to a backend",
            );
        };

        if resp.status != 200 {
            return Err(BifrostError::Other(format!(
                "doris: stream load {label} returned HTTP {} — {}",
                resp.status, resp.body
            )));
        }
        parse_load_result(label, &resp.body)
    }

    /// Ship one side of a window (upserts or deletes) as one or more labelled
    /// chunks. `chunk0` is the label index to start at, so upserts and deletes
    /// never share a label.
    async fn ship(
        &self,
        window: &TableWindow<'_>,
        batches: &[RecordBatch],
        merge_type: &str,
        chunk0: usize,
        labels: &mut Vec<String>,
    ) -> Result<(usize, usize)> {
        let bodies = encode_chunks(batches, self.cfg.max_chunk_bytes, self.cfg.format)?;
        let columns = batches
            .first()
            .map(|b| b.schema().fields().iter().map(|f| f.name().as_str()).collect::<Vec<_>>().join(","))
            .unwrap_or_default();
        let columns = columns.as_str();
        // Every chunk carries its own derived label, so chunks are independent
        // loads: they go through the ONE async I/O engine with up to
        // `max_parallel_loads` in flight, results re-sequenced into chunk order.
        // A failed chunk fails the window; the ones that landed are refused by
        // label on the replay, so order of completion never matters.
        let jobs = bodies.into_iter().enumerate().map(|(i, body)| {
            let label = stream_load_label(
                &self.cfg.db,
                &self.cfg.table,
                window.snapshots.from,
                window.snapshots.to,
                chunk0 + i,
            );
            async move {
                let (rows, dup) = self.put_chunk(&label, merge_type, columns, bytes::Bytes::from(body)).await?;
                Ok::<_, anyhow::Error>((label, rows, dup))
            }
        });
        let done = znippy_zoomies::gatling::io::run_ordered(jobs, self.cfg.max_parallel_loads.max(1), |x| x)
            .await
            .map_err(|e| BifrostError::Other(format!("{e:#}")))?;
        let mut loaded = 0;
        let mut duplicate = 0;
        for (label, rows, was_dup) in done {
            if was_dup {
                duplicate += rows;
            } else {
                loaded += rows;
            }
            labels.push(label);
        }
        Ok((loaded, duplicate))
    }
}

#[async_trait]
impl<T: StreamLoadTransport> TableSink for DorisSink<T> {
    /// Refuse to start on a schema the destination cannot accept.
    ///
    /// The check is "every source column exists in the destination". A
    /// destination column the source does not send is fine (Doris fills its
    /// default); a source column the destination does not have is a load that
    /// would fail row by row, halfway through.
    async fn probe(&mut self, schema: &Schema) -> Result<()> {
        let remote = self.remote_columns().await?;
        let missing: Vec<&str> = schema
            .fields()
            .iter()
            .map(|f| f.name().as_str())
            .filter(|n| !remote.contains(*n))
            .collect();
        let ok = missing.is_empty();
        crate::testmatrix::emit(
            crate::testmatrix::COMPONENT,
            "doris_sink::schema_probe",
            ok,
            &format!(
                "{}.{}: {} source column(s), {} missing remotely",
                self.cfg.db,
                self.cfg.table,
                schema.fields().len(),
                missing.len()
            ),
        );
        if !ok {
            return Err(BifrostError::Other(format!(
                "doris: {}.{} is missing column(s) the source sends: {} — refusing to \
                 start. Discovering this mid-stream, after a partial load, is worse.",
                self.cfg.db,
                self.cfg.table,
                missing.join(", ")
            )));
        }
        self.probed = true;
        Ok(())
    }

    async fn apply(&mut self, window: &TableWindow<'_>, token: AckToken) -> Result<LoadAck> {
        // Append mode means the source is append-only. If it is not, say so.
        if self.cfg.mode == LoadMode::Append && window.delete_rows() > 0 {
            crate::testmatrix::emit(
                crate::testmatrix::COMPONENT,
                "doris_sink::append_mode_refuses_deletes",
                false,
                &format!(
                    "{}: {} delete row(s) in Append mode",
                    window.table,
                    window.delete_rows()
                ),
            );
            return Err(BifrostError::Other(format!(
                "doris: {} delete row(s) arrived for {}.{} in Append mode — the source is \
                 not append-only after all. Refusing rather than silently diverging from \
                 the lake. Use a unique-key table with LoadMode::Upsert, or fix the writer.",
                window.delete_rows(),
                self.cfg.db,
                self.cfg.table
            )));
        }

        let mut labels = Vec::new();
        // Deletes first, then upserts: a delete followed by a re-insert of the
        // same key inside one window then nets to "present", which is the
        // common case (a `Delete` snapshot followed by an `Append` snapshot).
        //
        // Stated rather than hidden: a window carries no per-row ordering, so
        // the reverse sequence (insert then delete of one key, inside a single
        // window) nets to "present" here and should net to "absent". Keep
        // windows short, or feed a `ChangelogBatch::collapsed()` changelog,
        // which resolves the pair on skade's side where the ordering is known.
        let (del_loaded, del_dup) = if window.delete_rows() > 0 {
            self.ship(window, window.deletes, "DELETE", 0, &mut labels)
                .await?
        } else {
            (0, 0)
        };
        let chunk0 = labels.len();
        let (up_loaded, up_dup) = if window.upsert_rows() > 0 {
            self.ship(window, window.upserts, "APPEND", chunk0, &mut labels)
                .await?
        } else {
            (0, 0)
        };

        tracing::info!(
            target: "bifrost::doris",
            event = "stream_load_window",
            table = %window.table,
            dest = %format!("{}.{}", self.cfg.db, self.cfg.table),
            from = ?window.snapshots.from,
            to = window.snapshots.to,
            chunks = labels.len(),
            loaded = up_loaded,
            deleted = del_loaded,
            duplicate = up_dup + del_dup,
            "stream load complete",
        );

        Ok(token.ack(window.snapshots.to, labels).with_counts(
            up_loaded,
            del_loaded,
            up_dup + del_dup,
        ))
    }
}

// ── Stream Load result parsing ──────────────────────────────────────────────

/// Interpret a Stream Load response body.
///
/// `Label Already Exists` is **not** an error: it is the replay path working as
/// designed. The window was already loaded under this exact label, so the
/// correct action is to carry on and let the cursor advance.
///
/// Doris's own wording for the label guarantee is *At-Most-Once*, not
/// exactly-once; combined with a cursor that only advances on success (see
/// [`crate::sink::table::deliver_window`]) that is what yields effectively-once.
pub fn parse_load_result(label: &str, body: &str) -> Result<(usize, bool)> {
    let v: serde_json::Value = serde_json::from_str(body).map_err(|e| {
        BifrostError::Other(format!(
            "doris: stream load {label} returned non-JSON: {e} — {body}"
        ))
    })?;
    let status = v
        .get("Status")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("");
    let msg = v
        .get("Message")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("");
    let loaded = v
        .get("NumberLoadedRows")
        .and_then(serde_json::Value::as_u64)
        .unwrap_or(0) as usize;

    match status {
        "Success" | "Publish Timeout" => {
            let filtered = v
                .get("NumberFilteredRows")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(0);
            if filtered > 0 {
                return Err(BifrostError::Other(format!(
                    "doris: stream load {label} filtered {filtered} row(s) — refusing. A \
                     dropped row is the worst failure a change feed can have. {}",
                    v.get("ErrorURL")
                        .and_then(serde_json::Value::as_str)
                        .unwrap_or("")
                )));
            }
            Ok((loaded, false))
        }
        // The replay path. Do NOT add a second dedup layer on top of this.
        "Label Already Exists" => {
            tracing::info!(
                target: "bifrost::doris",
                event = "duplicate_label_rejected",
                label = %label,
                "replay rejected by label — window already loaded",
            );
            Ok((0, true))
        }
        other => Err(BifrostError::Other(format!(
            "doris: stream load {label} failed with status `{other}` — {msg}. {}",
            v.get("ErrorURL")
                .and_then(serde_json::Value::as_str)
                .unwrap_or("")
        ))),
    }
}

// ── encoding ────────────────────────────────────────────────────────────────

/// Encode `batches` as JSON arrays, split so no single body exceeds `budget`.
///
/// A batch that alone exceeds the budget is halved recursively rather than
/// guessed at from an average row width, so a table with a few enormous rows
/// still produces bodies under the cap.
fn encode_chunks(batches: &[RecordBatch], budget: usize, format: LoadFormat) -> Result<Vec<Vec<u8>>> {
    let mut out: Vec<Vec<u8>> = Vec::new();
    let mut pending: Vec<RecordBatch> = Vec::new();
    let mut pending_bytes = 0usize;

    for batch in batches {
        if batch.num_rows() == 0 {
            continue;
        }
        let encoded = encode(format, std::slice::from_ref(batch))?;
        if encoded.len() > budget {
            // Flush what we have, then split this one.
            if !pending.is_empty() {
                out.push(encode(format, &pending)?);
                pending.clear();
                pending_bytes = 0;
            }
            split_encode(batch, budget, format, &mut out)?;
            continue;
        }
        if pending_bytes + encoded.len() > budget && !pending.is_empty() {
            out.push(encode(format, &pending)?);
            pending.clear();
            pending_bytes = 0;
        }
        pending_bytes += encoded.len();
        pending.push(batch.clone());
    }
    if !pending.is_empty() {
        out.push(encode(format, &pending)?);
    }
    Ok(out)
}

/// Halve a too-large batch until each half encodes under `budget`. A
/// single-row batch that still exceeds it is emitted anyway — refusing would
/// drop data, and Doris's own cap is far above any sane single row.
fn split_encode(batch: &RecordBatch, budget: usize, format: LoadFormat, out: &mut Vec<Vec<u8>>) -> Result<()> {
    let encoded = encode(format, std::slice::from_ref(batch))?;
    if encoded.len() <= budget || batch.num_rows() <= 1 {
        out.push(encoded);
        return Ok(());
    }
    let half = batch.num_rows() / 2;
    split_encode(&batch.slice(0, half), budget, format, out)?;
    split_encode(&batch.slice(half, batch.num_rows() - half), budget, format, out)
}

/// One Stream Load body in `format`.
fn encode(format: LoadFormat, batches: &[RecordBatch]) -> Result<Vec<u8>> {
    match format {
        LoadFormat::Json => encode_json(batches),
        LoadFormat::Arrow => encode_arrow(batches),
    }
}

/// Arrow → one Arrow IPC **stream** (schema message, then the batches), which
/// Stream Load reads with `format: arrow`. No text round trip: the columns go
/// as they are.
fn encode_arrow(batches: &[RecordBatch]) -> Result<Vec<u8>> {
    let Some(first) = batches.first() else { return Ok(Vec::new()) };
    let bytes: usize = batches.iter().map(RecordBatch::get_array_memory_size).sum();
    let mut writer = arrow::ipc::writer::StreamWriter::try_new(Vec::with_capacity(bytes), &first.schema())?;
    for b in batches {
        writer.write(b)?;
    }
    writer.finish()?;
    Ok(writer.into_inner()?)
}

/// Arrow → a single JSON array, which Stream Load reads with
/// `strip_outer_array: true`.
fn encode_json(batches: &[RecordBatch]) -> Result<Vec<u8>> {
    let mut writer = arrow::json::ArrayWriter::new(Vec::new());
    for b in batches {
        writer.write(b)?;
    }
    writer.finish()?;
    Ok(writer.into_inner())
}

/// Standard base64, for the HTTP Basic header. Twelve lines is cheaper than a
/// dependency, and this is the only place the crate needs it.
fn base64_std(input: &[u8]) -> String {
    const T: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
    for c in input.chunks(3) {
        let b = [c[0], *c.get(1).unwrap_or(&0), *c.get(2).unwrap_or(&0)];
        let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]);
        out.push(T[(n >> 18 & 63) as usize] as char);
        out.push(T[(n >> 12 & 63) as usize] as char);
        out.push(if c.len() > 1 {
            T[(n >> 6 & 63) as usize] as char
        } else {
            '='
        });
        out.push(if c.len() > 2 {
            T[(n & 63) as usize] as char
        } else {
            '='
        });
    }
    out
}

// ── the real transport ──────────────────────────────────────────────────────

/// The live HTTP transport.
///
/// Automatic redirect following is **disabled**: [`DorisSink`] follows the
/// FE→BE 307 itself so that an unreachable backend is reported as a network
/// failure instead of as the authentication error it resembles. (It is also how
/// `curl --location-trusted` differs from plain `curl` here — a client that
/// re-authenticates on redirect appears to fail auth.)
pub struct ReqwestTransport {
    client: reqwest::Client,
}

impl ReqwestTransport {
    /// Build the client for `cfg`. The timeouts are not decoration: an
    /// unreachable BE after the 307 **hangs** rather than refusing (measured on
    /// Doris 4.0.8), so without them one routing mistake stalls the stream
    /// forever with no error.
    pub fn new(cfg: &DorisConfig) -> Result<Self> {
        let client = reqwest::Client::builder()
            .redirect(reqwest::redirect::Policy::none())
            .connect_timeout(cfg.connect_timeout)
            .timeout(cfg.request_timeout)
            .build()
            .map_err(|e| BifrostError::Other(format!("doris: http client: {e}")))?;
        Ok(Self { client })
    }
}

#[async_trait]
impl StreamLoadTransport for ReqwestTransport {
    async fn send(&self, req: &StreamLoadRequest) -> Result<StreamLoadResponse> {
        let method = reqwest::Method::from_bytes(req.method.as_bytes())
            .map_err(|e| BifrostError::Other(format!("doris: bad method: {e}")))?;
        let mut rb = self.client.request(method, &req.url);
        for (k, v) in &req.headers {
            rb = rb.header(k.as_str(), v.as_str());
        }
        if !req.body.is_empty() {
            rb = rb.body(req.body.clone()); // Bytes: a refcount, not a copy
        }
        let resp = rb
            .send()
            .await
            .map_err(|e| BifrostError::Other(format!("doris: {} {}: {e}", req.method, req.url)))?;
        let status = resp.status().as_u16();
        let location = resp
            .headers()
            .get(reqwest::header::LOCATION)
            .and_then(|v| v.to_str().ok())
            .map(str::to_owned);
        let body = resp
            .text()
            .await
            .map_err(|e| BifrostError::Other(format!("doris: reading response body: {e}")))?;
        Ok(StreamLoadResponse {
            status,
            location,
            body,
        })
    }
}