Aether/crates/aether-data/adapters/sqlite/src/candidates.rs

use std::collections::{BTreeMap, BTreeSet};

use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqliteConnection};

use aether_data_contracts::repository::candidates::{
    request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket,
    RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
    StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use aether_data_contracts::DataLayerError;
use aether_data_query::{push_in, WhereClause};

use crate::error::SqlResultExt;
use crate::SqlitePool;

const CANDIDATE_COLUMNS: &str = r#"
SELECT
  id,
  request_id,
  user_id,
  api_key_id,
  username,
  api_key_name,
  candidate_index,
  retry_index,
  provider_id,
  endpoint_id,
  key_id,
  status,
  skip_reason,
  is_cached,
  status_code,
  error_type,
  error_message,
  latency_ms,
  concurrent_requests,
  extra_data,
  required_capabilities,
  created_at AS created_at_unix_ms,
  started_at AS started_at_unix_ms,
  finished_at AS finished_at_unix_ms
FROM request_candidates
"#;

#[derive(Debug, Clone)]
pub struct SqliteRequestCandidateRepository {
    pool: SqlitePool,
}

impl SqliteRequestCandidateRepository {
    pub fn new(pool: SqlitePool) -> Self {
        Self { pool }
    }
}

#[async_trait]
impl RequestCandidateReadRepository for SqliteRequestCandidateRepository {
    async fn list_by_request_id(
        &self,
        request_id: &str,
    ) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
        let rows = sqlx::query(&format!(
            "{CANDIDATE_COLUMNS} WHERE request_id = ? ORDER BY candidate_index ASC, retry_index ASC, created_at ASC"
        ))
        .bind(request_id)
        .fetch_all(&self.pool)
        .await
        .map_sql_err()?;
        rows.iter().map(map_candidate_row).collect()
    }

    async fn list_attempted_by_request_id(
        &self,
        request_id: &str,
    ) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
        let rows = sqlx::query(&format!(
            "{CANDIDATE_COLUMNS} WHERE request_id = ? \
             AND (status IN ('streaming', 'success', 'failed', 'cancelled') \
             OR (status = 'pending' AND started_at IS NOT NULL)) \
             ORDER BY candidate_index ASC, retry_index ASC, created_at ASC"
        ))
        .bind(request_id)
        .fetch_all(&self.pool)
        .await
        .map_sql_err()?;
        rows.iter().map(map_candidate_row).collect()
    }

    async fn list_recent(
        &self,
        limit: usize,
    ) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
        if limit == 0 {
            return Ok(Vec::new());
        }
        let rows = sqlx::query(&format!(
            "{CANDIDATE_COLUMNS} ORDER BY created_at DESC LIMIT ?"
        ))
        .bind(limit_i64(limit, "recent request candidate limit")?)
        .fetch_all(&self.pool)
        .await
        .map_sql_err()?;
        rows.iter().map(map_candidate_row).collect()
    }

    async fn list_by_provider_id(
        &self,
        provider_id: &str,
        limit: usize,
    ) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
        if limit == 0 {
            return Ok(Vec::new());
        }
        let rows = sqlx::query(&format!(
            "{CANDIDATE_COLUMNS} WHERE provider_id = ? ORDER BY created_at DESC LIMIT ?"
        ))
        .bind(provider_id)
        .bind(limit_i64(limit, "provider request candidate limit")?)
        .fetch_all(&self.pool)
        .await
        .map_sql_err()?;
        rows.iter().map(map_candidate_row).collect()
    }

    async fn list_finalized_by_endpoint_ids_since(
        &self,
        endpoint_ids: &[String],
        since_unix_secs: u64,
        limit: usize,
    ) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
        if endpoint_ids.is_empty() || limit == 0 {
            return Ok(Vec::new());
        }
        let mut builder = QueryBuilder::<Sqlite>::new(CANDIDATE_COLUMNS);
        let mut where_clause = WhereClause::new();
        push_in(&mut builder, &mut where_clause, "endpoint_id", endpoint_ids);
        builder
            .push(" AND created_at >= ")
            .push_bind(unix_secs_to_ms_i64(since_unix_secs)?)
            .push(" AND status IN ('success', 'failed', 'skipped')")
            .push(" ORDER BY created_at DESC LIMIT ")
            .push_bind(limit_i64(limit, "finalized request candidate limit")?);
        let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
        rows.iter().map(map_candidate_row).collect()
    }

    async fn count_finalized_statuses_by_endpoint_ids_since(
        &self,
        endpoint_ids: &[String],
        since_unix_secs: u64,
    ) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
        if endpoint_ids.is_empty() {
            return Ok(Vec::new());
        }
        let mut builder = QueryBuilder::<Sqlite>::new(
            "SELECT endpoint_id, status, COUNT(id) AS count FROM request_candidates",
        );
        let mut where_clause = WhereClause::new();
        push_in(&mut builder, &mut where_clause, "endpoint_id", endpoint_ids);
        builder
            .push(" AND created_at >= ")
            .push_bind(unix_secs_to_ms_i64(since_unix_secs)?)
            .push(" AND status IN ('success', 'failed', 'skipped')")
            .push(" GROUP BY endpoint_id, status");
        let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
        rows.iter()
            .map(|row| {
                Ok(PublicHealthStatusCount {
                    endpoint_id: row.try_get("endpoint_id").map_sql_err()?,
                    status: RequestCandidateStatus::from_database(
                        row.try_get::<String, _>("status").map_sql_err()?.as_str(),
                    )?,
                    count: u64::try_from(row.try_get::<i64, _>("count").map_sql_err()?).map_err(
                        |_| {
                            DataLayerError::UnexpectedValue(
                                "public health status count out of range".to_string(),
                            )
                        },
                    )?,
                })
            })
            .collect()
    }

    async fn aggregate_finalized_timeline_by_endpoint_ids_since(
        &self,
        endpoint_ids: &[String],
        since_unix_secs: u64,
        until_unix_secs: u64,
        segments: u32,
    ) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
        if endpoint_ids.is_empty() || segments == 0 || until_unix_secs < since_unix_secs {
            return Ok(Vec::new());
        }
        let since_ms = unix_secs_to_ms_i64(since_unix_secs)?;
        let until_ms = unix_secs_to_ms_i64(until_unix_secs)?;
        let mut builder = QueryBuilder::<Sqlite>::new(CANDIDATE_COLUMNS);
        let mut where_clause = WhereClause::new();
        push_in(&mut builder, &mut where_clause, "endpoint_id", endpoint_ids);
        builder
            .push(" AND created_at >= ")
            .push_bind(since_ms)
            .push(" AND created_at <= ")
            .push_bind(until_ms)
            .push(" AND status IN ('success', 'failed', 'skipped')");
        let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
        aggregate_timeline(
            rows.iter()
                .map(map_candidate_row)
                .collect::<Result<Vec<_>, _>>()?,
            since_unix_secs,
            until_unix_secs,
            segments,
        )
    }
}

#[async_trait]
impl RequestCandidateWriteRepository for SqliteRequestCandidateRepository {
    async fn upsert(
        &self,
        candidate: UpsertRequestCandidateRecord,
    ) -> Result<StoredRequestCandidate, DataLayerError> {
        candidate.validate()?;
        let mut tx = self.pool.begin().await.map_sql_err()?;
        match upsert_candidate_in_transaction(&mut tx, candidate).await {
            Ok(candidate) => {
                tx.commit().await.map_sql_err()?;
                Ok(candidate)
            }
            Err(err) => {
                tx.rollback().await.map_sql_err()?;
                Err(err)
            }
        }
    }

    async fn upsert_many(
        &self,
        candidates: Vec<UpsertRequestCandidateRecord>,
    ) -> Result<usize, DataLayerError> {
        if candidates.is_empty() {
            return Ok(0);
        }
        for candidate in &candidates {
            candidate.validate()?;
        }

        let mut tx = self.pool.begin().await.map_sql_err()?;
        let result: Result<usize, DataLayerError> = async {
            let mut persisted = 0usize;
            for candidate in candidates {
                upsert_candidate_in_transaction(&mut tx, candidate).await?;
                persisted = persisted.saturating_add(1);
            }
            Ok(persisted)
        }
        .await;
        match result {
            Ok(persisted) => {
                tx.commit().await.map_sql_err()?;
                Ok(persisted)
            }
            Err(err) => {
                tx.rollback().await.map_sql_err()?;
                Err(err)
            }
        }
    }

    async fn delete_created_before(
        &self,
        created_before_unix_secs: u64,
        limit: usize,
    ) -> Result<usize, DataLayerError> {
        if limit == 0 {
            return Ok(0);
        }
        let rows_affected = sqlx::query(
            r#"
DELETE FROM request_candidates
WHERE id IN (
  SELECT id
  FROM request_candidates
  WHERE created_at < ?
  ORDER BY created_at ASC, id ASC
  LIMIT ?
)
"#,
        )
        .bind(unix_secs_to_ms_i64(created_before_unix_secs)?)
        .bind(limit_i64(limit, "request candidate delete limit")?)
        .execute(&self.pool)
        .await
        .map_sql_err()?
        .rows_affected();
        Ok(usize::try_from(rows_affected).unwrap_or_default())
    }
}

async fn upsert_candidate_in_transaction(
    tx: &mut sqlx::Transaction<'_, Sqlite>,
    candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, DataLayerError> {
    // Write first so SQLite upgrades the deferred transaction before the Rust
    // merge reads the latest state for this unique key.
    let insert_candidate = merge_candidate(candidate.clone(), None)?;
    insert_candidate_if_absent(tx, &insert_candidate).await?;
    let existing = find_by_unique(
        tx,
        &candidate.request_id,
        candidate.candidate_index,
        candidate.retry_index,
    )
    .await?
    .ok_or_else(|| {
        DataLayerError::UnexpectedValue(
            "request candidate row was not locked after insert-if-absent".to_string(),
        )
    })?;
    let merged = merge_candidate(candidate, Some(existing))?;
    upsert_merged_candidate(tx, &merged).await?;
    find_by_unique(
        tx,
        &merged.request_id,
        merged.candidate_index,
        merged.retry_index,
    )
    .await?
    .ok_or_else(|| {
        DataLayerError::UnexpectedValue(
            "request candidate row disappeared after atomic upsert".to_string(),
        )
    })
}

async fn insert_candidate_if_absent(
    connection: &mut SqliteConnection,
    candidate: &StoredRequestCandidate,
) -> Result<(), DataLayerError> {
    sqlx::query(
        r#"
INSERT INTO request_candidates (
  id, request_id, candidate_index, retry_index, status, created_at
)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT(request_id, candidate_index, retry_index) DO NOTHING
"#,
    )
    .bind(&candidate.id)
    .bind(&candidate.request_id)
    .bind(to_i32(candidate.candidate_index)?)
    .bind(to_i32(candidate.retry_index)?)
    .bind(status_to_database(candidate.status))
    .bind(u64_to_i64(
        candidate.created_at_unix_ms,
        "request candidate created_at",
    )?)
    .execute(connection)
    .await
    .map_sql_err()?;
    Ok(())
}

async fn find_by_unique(
    connection: &mut SqliteConnection,
    request_id: &str,
    candidate_index: u32,
    retry_index: u32,
) -> Result<Option<StoredRequestCandidate>, DataLayerError> {
    let row = sqlx::query(&format!(
        "{CANDIDATE_COLUMNS} WHERE request_id = ? AND candidate_index = ? AND retry_index = ? LIMIT 1"
    ))
    .bind(request_id)
    .bind(to_i32(candidate_index)?)
    .bind(to_i32(retry_index)?)
    .fetch_optional(connection)
    .await
    .map_sql_err()?;
    row.as_ref().map(map_candidate_row).transpose()
}

async fn upsert_merged_candidate(
    connection: &mut SqliteConnection,
    candidate: &StoredRequestCandidate,
) -> Result<(), DataLayerError> {
    sqlx::query(
        r#"
INSERT INTO request_candidates (
  id, request_id, user_id, api_key_id, username, api_key_name,
  candidate_index, retry_index, provider_id, endpoint_id, key_id, status,
  skip_reason, is_cached, status_code, error_type, error_message, latency_ms,
  concurrent_requests, extra_data, required_capabilities, created_at, started_at, finished_at
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(request_id, candidate_index, retry_index) DO UPDATE SET
  user_id = excluded.user_id,
  api_key_id = excluded.api_key_id,
  username = excluded.username,
  api_key_name = excluded.api_key_name,
  provider_id = excluded.provider_id,
  endpoint_id = excluded.endpoint_id,
  key_id = excluded.key_id,
  status = CASE
    WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
      AND excluded.status IN ('available', 'unused', 'pending', 'streaming')
      THEN request_candidates.status
    WHEN request_candidates.status = 'pending'
      AND excluded.status IN ('available', 'unused')
      THEN request_candidates.status
    WHEN request_candidates.status = 'streaming'
      AND excluded.status IN ('available', 'unused', 'pending')
      THEN request_candidates.status
    ELSE excluded.status
  END,
  skip_reason = excluded.skip_reason,
  is_cached = excluded.is_cached,
  status_code = CASE
    WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
      AND excluded.status IN ('available', 'unused', 'pending', 'streaming')
      THEN request_candidates.status_code
    WHEN request_candidates.status = 'pending'
      AND excluded.status IN ('available', 'unused')
      THEN request_candidates.status_code
    WHEN request_candidates.status = 'streaming'
      AND excluded.status IN ('available', 'unused', 'pending')
      THEN request_candidates.status_code
    ELSE COALESCE(excluded.status_code, request_candidates.status_code)
  END,
  error_type = CASE
    WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
      AND excluded.status IN ('available', 'unused', 'pending', 'streaming')
      THEN request_candidates.error_type
    WHEN request_candidates.status = 'pending'
      AND excluded.status IN ('available', 'unused')
      THEN request_candidates.error_type
    WHEN request_candidates.status = 'streaming'
      AND excluded.status IN ('available', 'unused', 'pending')
      THEN request_candidates.error_type
    ELSE COALESCE(excluded.error_type, request_candidates.error_type)
  END,
  error_message = CASE
    WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
      AND excluded.status IN ('available', 'unused', 'pending', 'streaming')
      THEN request_candidates.error_message
    WHEN request_candidates.status = 'pending'
      AND excluded.status IN ('available', 'unused')
      THEN request_candidates.error_message
    WHEN request_candidates.status = 'streaming'
      AND excluded.status IN ('available', 'unused', 'pending')
      THEN request_candidates.error_message
    ELSE COALESCE(excluded.error_message, request_candidates.error_message)
  END,
  latency_ms = CASE
    WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
      AND excluded.status IN ('available', 'unused', 'pending', 'streaming')
      THEN request_candidates.latency_ms
    WHEN request_candidates.status = 'pending'
      AND excluded.status IN ('available', 'unused')
      THEN request_candidates.latency_ms
    WHEN request_candidates.status = 'streaming'
      AND excluded.status IN ('available', 'unused', 'pending')
      THEN request_candidates.latency_ms
    ELSE COALESCE(excluded.latency_ms, request_candidates.latency_ms)
  END,
  concurrent_requests = excluded.concurrent_requests,
  extra_data = excluded.extra_data,
  required_capabilities = excluded.required_capabilities,
  created_at = excluded.created_at,
  started_at = excluded.started_at,
  finished_at = CASE
    WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
      AND excluded.status IN ('available', 'unused', 'pending', 'streaming')
      THEN request_candidates.finished_at
    WHEN request_candidates.status = 'pending'
      AND excluded.status IN ('available', 'unused')
      THEN request_candidates.finished_at
    WHEN request_candidates.status = 'streaming'
      AND excluded.status IN ('available', 'unused', 'pending')
      THEN request_candidates.finished_at
    ELSE COALESCE(excluded.finished_at, request_candidates.finished_at)
  END
"#,
    )
    .bind(&candidate.id)
    .bind(&candidate.request_id)
    .bind(&candidate.user_id)
    .bind(&candidate.api_key_id)
    .bind(&candidate.username)
    .bind(&candidate.api_key_name)
    .bind(to_i32(candidate.candidate_index)?)
    .bind(to_i32(candidate.retry_index)?)
    .bind(&candidate.provider_id)
    .bind(&candidate.endpoint_id)
    .bind(&candidate.key_id)
    .bind(status_to_database(candidate.status))
    .bind(&candidate.skip_reason)
    .bind(candidate.is_cached)
    .bind(candidate.status_code.map(i32::from))
    .bind(&candidate.error_type)
    .bind(&candidate.error_message)
    .bind(candidate.latency_ms.map(to_i32_u64).transpose()?)
    .bind(candidate.concurrent_requests.map(to_i32).transpose()?)
    .bind(json_to_string(&candidate.extra_data)?)
    .bind(json_to_string(&candidate.required_capabilities)?)
    .bind(u64_to_i64(
        candidate.created_at_unix_ms,
        "request candidate created_at",
    )?)
    .bind(optional_u64_to_i64(
        candidate.started_at_unix_ms,
        "request candidate started_at",
    )?)
    .bind(optional_u64_to_i64(
        candidate.finished_at_unix_ms,
        "request candidate finished_at",
    )?)
    .execute(connection)
    .await
    .map_sql_err()?;
    Ok(())
}

fn merge_candidate(
    candidate: UpsertRequestCandidateRecord,
    existing: Option<StoredRequestCandidate>,
) -> Result<StoredRequestCandidate, DataLayerError> {
    let preserve_existing_lifecycle = existing.as_ref().is_some_and(|value| {
        request_candidate_lifecycle_would_regress(value.status, candidate.status)
    });
    let merged_status = if preserve_existing_lifecycle {
        existing
            .as_ref()
            .map(|value| value.status)
            .unwrap_or(candidate.status)
    } else {
        candidate.status
    };
    let created_at_unix_ms = candidate
        .created_at_unix_ms
        .filter(|value| *value > 1000)
        .or_else(|| {
            existing
                .as_ref()
                .map(|value| value.created_at_unix_ms)
                .filter(|value| *value > 1000)
        })
        .or(candidate.started_at_unix_ms)
        .or(candidate.finished_at_unix_ms)
        .unwrap_or_else(current_unix_ms);
    let id = existing
        .as_ref()
        .map(|value| value.id.clone())
        .unwrap_or(candidate.id);
    let extra_data = merge_json_objects(
        existing.as_ref().and_then(|value| value.extra_data.clone()),
        candidate.extra_data,
    );
    StoredRequestCandidate::new(
        id,
        candidate.request_id,
        candidate
            .user_id
            .or_else(|| existing.as_ref().and_then(|value| value.user_id.clone())),
        candidate
            .api_key_id
            .or_else(|| existing.as_ref().and_then(|value| value.api_key_id.clone())),
        candidate
            .username
            .or_else(|| existing.as_ref().and_then(|value| value.username.clone())),
        candidate.api_key_name.or_else(|| {
            existing
                .as_ref()
                .and_then(|value| value.api_key_name.clone())
        }),
        to_i32(candidate.candidate_index)?,
        to_i32(candidate.retry_index)?,
        candidate.provider_id.or_else(|| {
            existing
                .as_ref()
                .and_then(|value| value.provider_id.clone())
        }),
        candidate.endpoint_id.or_else(|| {
            existing
                .as_ref()
                .and_then(|value| value.endpoint_id.clone())
        }),
        candidate
            .key_id
            .or_else(|| existing.as_ref().and_then(|value| value.key_id.clone())),
        merged_status,
        candidate.skip_reason.or_else(|| {
            existing
                .as_ref()
                .and_then(|value| value.skip_reason.clone())
        }),
        candidate
            .is_cached
            .unwrap_or_else(|| existing.as_ref().is_some_and(|value| value.is_cached)),
        if preserve_existing_lifecycle {
            existing
                .as_ref()
                .and_then(|value| value.status_code.map(i32::from))
        } else {
            candidate.status_code.map(i32::from).or_else(|| {
                existing
                    .as_ref()
                    .and_then(|value| value.status_code.map(i32::from))
            })
        },
        if preserve_existing_lifecycle {
            existing.as_ref().and_then(|value| value.error_type.clone())
        } else {
            candidate
                .error_type
                .or_else(|| existing.as_ref().and_then(|value| value.error_type.clone()))
        },
        if preserve_existing_lifecycle {
            existing
                .as_ref()
                .and_then(|value| value.error_message.clone())
        } else {
            candidate.error_message.or_else(|| {
                existing
                    .as_ref()
                    .and_then(|value| value.error_message.clone())
            })
        },
        if preserve_existing_lifecycle {
            match existing.as_ref().and_then(|value| value.latency_ms) {
                Some(value) => Some(to_i32_u64(value)?),
                None => None,
            }
        } else {
            candidate.latency_ms.map(to_i32_u64).transpose()?.or(
                match existing.as_ref().and_then(|value| value.latency_ms) {
                    Some(value) => Some(to_i32_u64(value)?),
                    None => None,
                },
            )
        },
        candidate.concurrent_requests.map(to_i32).transpose()?.or(
            match existing
                .as_ref()
                .and_then(|value| value.concurrent_requests)
            {
                Some(value) => Some(to_i32(value)?),
                None => None,
            },
        ),
        extra_data,
        candidate.required_capabilities.or_else(|| {
            existing
                .as_ref()
                .and_then(|value| value.required_capabilities.clone())
        }),
        u64_to_i64(created_at_unix_ms, "request candidate created_at")?,
        candidate
            .started_at_unix_ms
            .or_else(|| existing.as_ref().and_then(|value| value.started_at_unix_ms))
            .map(|value| u64_to_i64(value, "request candidate started_at"))
            .transpose()?,
        if preserve_existing_lifecycle {
            existing
                .as_ref()
                .and_then(|value| value.finished_at_unix_ms)
        } else {
            candidate.finished_at_unix_ms.or_else(|| {
                existing
                    .as_ref()
                    .and_then(|value| value.finished_at_unix_ms)
            })
        }
        .map(|value| u64_to_i64(value, "request candidate finished_at"))
        .transpose()?,
    )
}

fn aggregate_timeline(
    candidates: Vec<StoredRequestCandidate>,
    since_unix_secs: u64,
    until_unix_secs: u64,
    segments: u32,
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
    let endpoint_ids = candidates
        .iter()
        .filter_map(|candidate| candidate.endpoint_id.clone())
        .collect::<BTreeSet<_>>();
    let span_ms = until_unix_secs
        .saturating_sub(since_unix_secs)
        .saturating_mul(1000)
        .max(1);
    let since_ms = since_unix_secs.saturating_mul(1000);
    let mut buckets = BTreeMap::<(String, u32), PublicHealthTimelineBucket>::new();
    for candidate in candidates {
        let Some(endpoint_id) = candidate.endpoint_id.clone() else {
            continue;
        };
        let offset = candidate.created_at_unix_ms.saturating_sub(since_ms);
        let segment_idx = ((offset.saturating_mul(u64::from(segments))) / span_ms)
            .min(u64::from(segments.saturating_sub(1))) as u32;
        let bucket = buckets.entry((endpoint_id.clone(), segment_idx)).or_insert(
            PublicHealthTimelineBucket {
                endpoint_id,
                segment_idx,
                total_count: 0,
                success_count: 0,
                failed_count: 0,
                min_created_at_unix_ms: Some(candidate.created_at_unix_ms),
                max_created_at_unix_ms: Some(candidate.created_at_unix_ms),
            },
        );
        bucket.total_count += 1;
        if candidate.status == RequestCandidateStatus::Success {
            bucket.success_count += 1;
        }
        if candidate.status == RequestCandidateStatus::Failed {
            bucket.failed_count += 1;
        }
        bucket.min_created_at_unix_ms = bucket
            .min_created_at_unix_ms
            .map(|value| value.min(candidate.created_at_unix_ms));
        bucket.max_created_at_unix_ms = bucket
            .max_created_at_unix_ms
            .map(|value| value.max(candidate.created_at_unix_ms));
    }
    for endpoint_id in endpoint_ids {
        for segment_idx in 0..segments {
            buckets.entry((endpoint_id.clone(), segment_idx)).or_insert(
                PublicHealthTimelineBucket {
                    endpoint_id: endpoint_id.clone(),
                    segment_idx,
                    total_count: 0,
                    success_count: 0,
                    failed_count: 0,
                    min_created_at_unix_ms: None,
                    max_created_at_unix_ms: None,
                },
            );
        }
    }
    Ok(buckets.into_values().collect())
}

fn map_candidate_row(row: &SqliteRow) -> Result<StoredRequestCandidate, DataLayerError> {
    StoredRequestCandidate::new(
        row.try_get("id").map_sql_err()?,
        row.try_get("request_id").map_sql_err()?,
        row.try_get("user_id").map_sql_err()?,
        row.try_get("api_key_id").map_sql_err()?,
        row.try_get("username").map_sql_err()?,
        row.try_get("api_key_name").map_sql_err()?,
        row.try_get("candidate_index").map_sql_err()?,
        row.try_get("retry_index").map_sql_err()?,
        row.try_get("provider_id").map_sql_err()?,
        row.try_get("endpoint_id").map_sql_err()?,
        row.try_get("key_id").map_sql_err()?,
        RequestCandidateStatus::from_database(
            row.try_get::<String, _>("status").map_sql_err()?.as_str(),
        )?,
        row.try_get("skip_reason").map_sql_err()?,
        row.try_get("is_cached").map_sql_err()?,
        row.try_get("status_code").map_sql_err()?,
        row.try_get("error_type").map_sql_err()?,
        row.try_get("error_message").map_sql_err()?,
        row.try_get("latency_ms").map_sql_err()?,
        row.try_get("concurrent_requests").map_sql_err()?,
        parse_json(row.try_get("extra_data").ok().flatten())?,
        parse_json(row.try_get("required_capabilities").ok().flatten())?,
        row.try_get("created_at_unix_ms").map_sql_err()?,
        row.try_get("started_at_unix_ms").map_sql_err()?,
        row.try_get("finished_at_unix_ms").map_sql_err()?,
    )
}

fn parse_json(value: Option<String>) -> Result<Option<serde_json::Value>, DataLayerError> {
    value
        .filter(|value| !value.trim().is_empty())
        .map(|value| {
            serde_json::from_str(&value).map_err(|err| {
                DataLayerError::UnexpectedValue(format!(
                    "request_candidates JSON field is invalid: {err}"
                ))
            })
        })
        .transpose()
}

fn json_to_string(value: &Option<serde_json::Value>) -> Result<Option<String>, DataLayerError> {
    value
        .as_ref()
        .map(|value| {
            serde_json::to_string(value).map_err(|err| {
                DataLayerError::UnexpectedValue(format!(
                    "request_candidates JSON field is unserializable: {err}"
                ))
            })
        })
        .transpose()
}

fn merge_json_objects(
    existing: Option<serde_json::Value>,
    overlay: Option<serde_json::Value>,
) -> Option<serde_json::Value> {
    match (existing, overlay) {
        (
            Some(serde_json::Value::Object(mut existing_object)),
            Some(serde_json::Value::Object(overlay_object)),
        ) => {
            existing_object.extend(overlay_object);
            Some(serde_json::Value::Object(existing_object))
        }
        (_existing, Some(overlay)) => Some(overlay),
        (existing, None) => existing,
    }
}

fn status_to_database(status: RequestCandidateStatus) -> &'static str {
    match status {
        RequestCandidateStatus::Available => "available",
        RequestCandidateStatus::Unused => "unused",
        RequestCandidateStatus::Pending => "pending",
        RequestCandidateStatus::Streaming => "streaming",
        RequestCandidateStatus::Success => "success",
        RequestCandidateStatus::Failed => "failed",
        RequestCandidateStatus::Cancelled => "cancelled",
        RequestCandidateStatus::Skipped => "skipped",
    }
}

fn current_unix_ms() -> u64 {
    chrono::Utc::now().timestamp_millis().max(0) as u64
}

fn unix_secs_to_ms_i64(value: u64) -> Result<i64, DataLayerError> {
    let value = value.checked_mul(1000).ok_or_else(|| {
        DataLayerError::UnexpectedValue("request candidate timestamp overflow".to_string())
    })?;
    i64::try_from(value).map_err(|_| {
        DataLayerError::UnexpectedValue("request candidate timestamp overflow".to_string())
    })
}

fn limit_i64(value: usize, name: &str) -> Result<i64, DataLayerError> {
    i64::try_from(value)
        .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid {name}: {value}")))
}

fn to_i32(value: u32) -> Result<i32, DataLayerError> {
    i32::try_from(value).map_err(|_| {
        DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}"))
    })
}

fn to_i32_u64(value: u64) -> Result<i32, DataLayerError> {
    i32::try_from(value).map_err(|_| {
        DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}"))
    })
}

fn u64_to_i64(value: u64, name: &str) -> Result<i64, DataLayerError> {
    i64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{name} overflow")))
}

fn optional_u64_to_i64(value: Option<u64>, name: &str) -> Result<Option<i64>, DataLayerError> {
    value.map(|value| u64_to_i64(value, name)).transpose()
}

#[cfg(test)]
mod tests {
    use super::SqliteRequestCandidateRepository;
    use crate::run_migrations;
    use aether_data_contracts::repository::candidates::{
        RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
        UpsertRequestCandidateRecord,
    };
    use serde_json::json;

    #[tokio::test]
    async fn sqlite_repository_writes_and_reads_request_candidates() {
        let pool = sqlx::sqlite::SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("sqlite pool should connect");
        run_migrations(&pool)
            .await
            .expect("sqlite migrations should run");

        let repository = SqliteRequestCandidateRepository::new(pool);
        let created = repository
            .upsert(sample_upsert(
                "candidate-1",
                RequestCandidateStatus::Pending,
                Some(json!({"a": 1})),
                1_000_000,
            ))
            .await
            .expect("candidate should insert");
        assert_eq!(created.request_id, "request-1");

        let updated = repository
            .upsert(sample_upsert(
                "candidate-replacement",
                RequestCandidateStatus::Success,
                Some(json!({"b": 2})),
                1_000_500,
            ))
            .await
            .expect("candidate should update");
        assert_eq!(updated.id, "candidate-1");
        assert_eq!(updated.extra_data, Some(json!({"a": 1, "b": 2})));

        let late_streaming = repository
            .upsert(sample_upsert(
                "candidate-late-streaming",
                RequestCandidateStatus::Streaming,
                Some(json!({"late": true})),
                1_000_250,
            ))
            .await
            .expect("late streaming candidate should not regress terminal status");
        assert_eq!(late_streaming.id, "candidate-1");
        assert_eq!(late_streaming.status, RequestCandidateStatus::Success);
        assert_eq!(late_streaming.finished_at_unix_ms, Some(1_000_502));
        assert_eq!(
            late_streaming.extra_data,
            Some(json!({"a": 1, "b": 2, "late": true}))
        );

        assert_eq!(
            repository
                .list_by_request_id("request-1")
                .await
                .expect("request list should load")
                .len(),
            1
        );
        assert_eq!(
            repository
                .count_finalized_statuses_by_endpoint_ids_since(&["endpoint-1".to_string()], 900)
                .await
                .expect("status counts should load")[0]
                .count,
            1
        );
        assert_eq!(
            repository
                .aggregate_finalized_timeline_by_endpoint_ids_since(
                    &["endpoint-1".to_string()],
                    900,
                    1200,
                    3,
                )
                .await
                .expect("timeline should load")
                .len(),
            3
        );
        assert_eq!(
            repository
                .delete_created_before(2_000, 10)
                .await
                .expect("old candidates should delete"),
            1
        );
    }

    #[tokio::test]
    async fn sqlite_concurrent_upserts_merge_without_lifecycle_regression() {
        let database_path = std::env::temp_dir().join(format!(
            "aether-candidate-concurrency-{}.sqlite",
            uuid::Uuid::new_v4()
        ));
        let options = sqlx::sqlite::SqliteConnectOptions::new()
            .filename(&database_path)
            .create_if_missing(true)
            .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
            .busy_timeout(std::time::Duration::from_secs(30));
        let pool = sqlx::sqlite::SqlitePoolOptions::new()
            .max_connections(12)
            .connect_with(options)
            .await
            .expect("sqlite pool should connect");
        run_migrations(&pool)
            .await
            .expect("sqlite migrations should run");
        let repository = SqliteRequestCandidateRepository::new(pool.clone());
        let request_id = format!("candidate-concurrency-{}", uuid::Uuid::new_v4());

        let mut initial = sample_upsert(
            "initial",
            RequestCandidateStatus::Pending,
            Some(json!({"initial": true})),
            3_000_000,
        );
        initial.request_id = request_id.clone();
        repository
            .upsert(initial)
            .await
            .expect("initial candidate should insert");

        const WRITERS: usize = 10;
        let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(WRITERS));
        let mut tasks = Vec::new();
        for writer in 0..WRITERS {
            let repository = repository.clone();
            let request_id = request_id.clone();
            let barrier = barrier.clone();
            tasks.push(tokio::spawn(async move {
                let status = if writer == 0 {
                    RequestCandidateStatus::Success
                } else {
                    RequestCandidateStatus::Streaming
                };
                let mut extra_data = serde_json::Map::new();
                extra_data.insert(format!("writer_{writer}"), json!(writer));
                let mut candidate = sample_upsert(
                    format!("writer-{writer}").as_str(),
                    status,
                    Some(serde_json::Value::Object(extra_data)),
                    3_100_000 + u64::try_from(writer).expect("writer index should fit") * 10,
                );
                candidate.request_id = request_id;
                if writer != 0 {
                    candidate.latency_ms = Some(9_000 + writer as u64);
                    candidate.finished_at_unix_ms = Some(9_000_000 + writer as u64);
                }
                barrier.wait().await;
                repository.upsert(candidate).await
            }));
        }
        for task in tasks {
            task.await
                .expect("candidate writer should join")
                .expect("candidate writer should persist");
        }

        let candidates = repository
            .list_by_request_id(&request_id)
            .await
            .expect("concurrent candidate should load");
        assert_eq!(candidates.len(), 1);
        let candidate = &candidates[0];
        assert_eq!(candidate.id, "initial");
        assert_eq!(candidate.status, RequestCandidateStatus::Success);
        assert_eq!(candidate.latency_ms, Some(123));
        assert_eq!(candidate.finished_at_unix_ms, Some(3_100_002));
        let extra_data = candidate
            .extra_data
            .as_ref()
            .and_then(serde_json::Value::as_object)
            .expect("merged extra data should be an object");
        assert_eq!(extra_data.get("initial"), Some(&json!(true)));
        for writer in 0..WRITERS {
            assert_eq!(
                extra_data.get(format!("writer_{writer}").as_str()),
                Some(&json!(writer))
            );
        }

        drop(repository);
        pool.close().await;
        let _ = std::fs::remove_file(&database_path);
        let _ = std::fs::remove_file(format!("{}-wal", database_path.display()));
        let _ = std::fs::remove_file(format!("{}-shm", database_path.display()));
    }

    #[tokio::test]
    async fn sqlite_batch_upsert_preserves_order_and_rolls_back_on_error() {
        let pool = sqlx::sqlite::SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("sqlite pool should connect");
        run_migrations(&pool)
            .await
            .expect("sqlite migrations should run");
        let repository = SqliteRequestCandidateRepository::new(pool);
        let request_id = "request-batch";

        let mut pending = sample_upsert(
            "batch-first",
            RequestCandidateStatus::Pending,
            Some(json!({"pending": true})),
            4_000_000,
        );
        pending.request_id = request_id.to_string();
        let mut streaming = sample_upsert(
            "batch-second",
            RequestCandidateStatus::Streaming,
            Some(json!({"streaming": true})),
            4_000_100,
        );
        streaming.request_id = request_id.to_string();
        streaming.is_cached = None;
        let mut success = sample_upsert(
            "batch-third",
            RequestCandidateStatus::Success,
            Some(json!({"success": true})),
            4_000_200,
        );
        success.request_id = request_id.to_string();
        success.is_cached = Some(true);
        let mut late_pending = sample_upsert(
            "batch-fourth",
            RequestCandidateStatus::Pending,
            Some(json!({"late": true})),
            4_000_300,
        );
        late_pending.request_id = request_id.to_string();
        late_pending.is_cached = None;
        late_pending.latency_ms = Some(9_999);
        late_pending.finished_at_unix_ms = Some(9_999_999);

        assert_eq!(
            repository
                .upsert_many(vec![pending, streaming, success, late_pending])
                .await
                .expect("ordered batch should persist"),
            4
        );
        let candidates = repository
            .list_by_request_id(request_id)
            .await
            .expect("batch candidate should load");
        assert_eq!(candidates.len(), 1);
        let candidate = &candidates[0];
        assert_eq!(candidate.id, "batch-first");
        assert_eq!(candidate.status, RequestCandidateStatus::Success);
        assert!(candidate.is_cached);
        assert_eq!(candidate.latency_ms, Some(123));
        assert_eq!(candidate.finished_at_unix_ms, Some(4_000_202));
        assert_eq!(
            candidate.extra_data,
            Some(json!({
                "pending": true,
                "streaming": true,
                "success": true,
                "late": true
            }))
        );

        let rollback_request_id = "request-batch-rollback";
        let mut valid = sample_upsert(
            "rollback-valid",
            RequestCandidateStatus::Pending,
            None,
            5_000_000,
        );
        valid.request_id = rollback_request_id.to_string();
        let mut invalid = sample_upsert(
            "rollback-invalid",
            RequestCandidateStatus::Success,
            None,
            5_000_100,
        );
        invalid.request_id = rollback_request_id.to_string();
        invalid.candidate_index = 1;
        invalid.latency_ms = Some(u64::MAX);

        repository
            .upsert_many(vec![valid, invalid])
            .await
            .expect_err("invalid later row should roll back the batch");
        assert!(repository
            .list_by_request_id(rollback_request_id)
            .await
            .expect("rolled-back batch should be readable")
            .is_empty());
    }

    fn sample_upsert(
        id: &str,
        status: RequestCandidateStatus,
        extra_data: Option<serde_json::Value>,
        created_at_unix_ms: u64,
    ) -> UpsertRequestCandidateRecord {
        UpsertRequestCandidateRecord {
            id: id.to_string(),
            request_id: "request-1".to_string(),
            user_id: Some("user-1".to_string()),
            api_key_id: Some("key-1".to_string()),
            username: Some("user".to_string()),
            api_key_name: Some("Key".to_string()),
            candidate_index: 0,
            retry_index: 0,
            provider_id: Some("provider-1".to_string()),
            endpoint_id: Some("endpoint-1".to_string()),
            key_id: Some("provider-key-1".to_string()),
            status,
            skip_reason: None,
            is_cached: Some(false),
            status_code: Some(200),
            error_type: None,
            error_message: None,
            latency_ms: Some(123),
            concurrent_requests: Some(2),
            extra_data,
            required_capabilities: Some(json!({"streaming": true})),
            created_at_unix_ms: Some(created_at_unix_ms),
            started_at_unix_ms: Some(created_at_unix_ms + 1),
            finished_at_unix_ms: Some(created_at_unix_ms + 2),
        }
    }
}