use crate::selector::{CompileTestCase, RuntimeTestCase};
use crate::{
    ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkJobConclusion, BenchmarkRequest,
    BenchmarkRequestIndex, BenchmarkRequestInsertResult, BenchmarkRequestStatus,
    BenchmarkRequestWithErrors, CodegenBackend, CollectorConfig, CompileBenchmark,
    PendingBenchmarkRequests, QueueRuntimeConfig, Target,
};
use crate::{CollectionId, Index, Profile, Scenario};
use hashbrown::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};

pub mod postgres;
pub mod sqlite;

#[derive(Debug)]
pub enum JobEnqueueResult {
    JobCreated(u32),
    JobExistedOrParentNotFound,
    Other(anyhow::Error),
}

#[async_trait::async_trait]
pub trait Connection: Send + Sync {
    async fn maybe_create_indices(&mut self);
    async fn transaction(&mut self) -> Box<dyn Transaction + '_>;

    async fn load_index(&mut self) -> Index;

    /// Returns true if the given database backend supports the job queue system.
    fn supports_job_queue(&self) -> bool;

    /// None means that the caller doesn't know; it should be left alone if
    /// known or set to false if unknown.
    async fn record_compile_benchmark(
        &self,
        krate: &str,
        supports_stable: Option<bool>,
        category: String,
    );
    async fn get_compile_benchmarks(&self) -> Vec<CompileBenchmark>;

    async fn get_repos(&self) -> Vec<String>;

    async fn get_benchmarked_commits(
        &self,
        repo: Option<&str>,
    ) -> Vec<(String, String, Vec<String>)>;

    async fn artifact_by_name(&self, artifact: &str) -> Option<ArtifactId>;

    /// One collection corresponds to all gathered metrics for a single iteration of a test case.
    async fn collection_id(&self, version: &str) -> CollectionId;
    async fn artifact_id(&self, artifact: &ArtifactId, tag: &str) -> ArtifactIdNumber;

    #[allow(clippy::too_many_arguments)]
    async fn record_statistic(
        &self,
        collection: CollectionId,
        artifact: ArtifactIdNumber,
        benchmark: &str,
        profile: Profile,
        scenario: Scenario,
        backend: CodegenBackend,
        target: Target,
        metric: &str,
        value: f64,
    );
    #[allow(clippy::too_many_arguments)]
    async fn record_runtime_statistic(
        &self,
        collection: CollectionId,
        artifact: ArtifactIdNumber,
        benchmark: &str,
        metric: &str,
        target: Target,
        value: f64,
    );

    /// Record a runtime statistic with a JSON value (for structured/complex metrics).
    /// The numeric `value` column is set to 0 as a placeholder; the actual data is in `json_value`.
    #[allow(clippy::too_many_arguments)]
    async fn record_runtime_json_statistic(
        &self,
        collection: CollectionId,
        artifact: ArtifactIdNumber,
        benchmark: &str,
        metric: &str,
        target: Target,
        json_value: &str,
    );
    async fn record_error(
        &self,
        artifact: ArtifactIdNumber,
        context: &str,
        message: &str,
        job_id: Option<u32>,
    );
    async fn record_rustc_crate(
        &self,
        collection: CollectionId,
        artifact: ArtifactIdNumber,
        krate: &str,
        value: Duration,
    );

    /// Records the size of an artifact component (like `librustc_driver.so` or `libLLVM.so`) in
    /// bytes.
    async fn record_artifact_size(&self, artifact: ArtifactIdNumber, component: &str, size: u64);

    /// Returns the sizes of individual components of a single artifact.
    async fn get_artifact_size(&self, aid: ArtifactIdNumber) -> HashMap<String, u64>;

    /// Returns vector of bootstrap build times for the given artifacts. The kth
    /// element is the minimum build time for the kth artifact in `aids`, across
    /// all collections for the artifact, or none if there is no bootstrap data
    /// for that artifact (for example, because the rustc benchmark wasn't
    /// executed for that artifact).
    async fn get_bootstrap(&self, aids: &[ArtifactIdNumber]) -> Vec<Option<Duration>>;
    /// Returns map from rustc crate name to vector of build times for that crate
    /// for the given artifacts. Within a crate's corresponding vector, the kth
    /// element is the minimum build time for the kth artifact in `aids`, across
    /// all collections for the artifact, or none if there is no data for that
    /// artifact / crate combination (for example, because that rustc crate
    /// wasn't present when building rustc with that artifact, or because the
    /// rustc benchmark wasn't executed for that artifact). A crate will not be
    /// included as a key in the map unless at least one artifact in `aids` has a
    /// build time for it.
    async fn get_bootstrap_by_crate(
        &self,
        aids: &[ArtifactIdNumber],
    ) -> HashMap<String, Vec<Option<Duration>>>;
    /// Returns compile-time benchmark statistics (aggregated as min value).
    async fn get_compile_pstats(
        &self,
        pstat_series_row_ids: &[u32],
        artifact_row_id: &[Option<ArtifactIdNumber>],
    ) -> Vec<Vec<Option<f64>>>;

    /// Returns all sample values for each (series, artifact) combination for compile-time benchmarks.
    /// Unlike `get_compile_pstats` which returns only the min value, this returns all values
    /// for bootstrap analysis.
    async fn get_compile_pstats_all_samples(
        &self,
        pstat_series_row_ids: &[u32],
        artifact_row_id: &[Option<ArtifactIdNumber>],
    ) -> Vec<Vec<Vec<f64>>>;

    /// Returns runtime benchmark statistics (aggregated as min value).
    async fn get_runtime_pstats(
        &self,
        runtime_pstat_series_row_ids: &[u32],
        artifact_row_id: &[Option<ArtifactIdNumber>],
    ) -> Vec<Vec<Option<f64>>>;

    /// Returns all sample values for each (series, artifact) combination.
    /// Unlike `get_runtime_pstats` which returns only the min value, this returns all values
    /// for bootstrap analysis.
    async fn get_runtime_pstats_all_samples(
        &self,
        runtime_pstat_series_row_ids: &[u32],
        artifact_row_id: &[Option<ArtifactIdNumber>],
    ) -> Vec<Vec<Vec<f64>>>;

    /// Returns JSON values for runtime pstats.
    /// For each (series, artifact) combination, returns the first non-null json_value found.
    /// Outer vec: one per series. Inner vec: one per artifact.
    async fn get_runtime_json_pstats(
        &self,
        runtime_pstat_series_row_ids: &[u32],
        artifact_row_id: &[Option<ArtifactIdNumber>],
    ) -> Vec<Vec<Option<String>>>;

    async fn get_error(&self, artifact_row_id: ArtifactIdNumber) -> HashMap<String, String>;

    /// Returns error messages for a specific job, as a list of (context, message) pairs.
    async fn get_errors_for_job(&self, job_id: u32) -> Vec<(String, String)>;

    /// Returns the set of runtime metrics that have data for the given artifact IDs.
    async fn get_runtime_metrics_for_artifacts(&self, aids: &[ArtifactIdNumber]) -> Vec<String>;

    /// Returns the SHA of the parent of the given SHA commit, if available.
    async fn parent_of(&self, sha: &str) -> Option<String>;

    /// Returns the PR associated with an artifact with the given SHA, if available.
    async fn pr_of(&self, sha: &str) -> Option<u32>;

    /// Returns master commits grouped by repo, loaded from the `benchmark_request` table.
    /// Each repo key maps to a list of commits sorted newest-first.
    async fn list_master_commits_by_repo(&self) -> HashMap<String, Vec<crate::MasterCommitInfo>>;

    /// Returns the collection ids corresponding to the query. Usually just one.
    ///
    /// Currently only supported by postgres (sqlite does not store self-profile
    /// results in the raw format).
    async fn list_self_profile(
        &self,
        aid: ArtifactId,
        crate_: &str,
        profile: &str,
        cache: &str,
    ) -> Vec<(ArtifactIdNumber, CollectionId)>;

    /// Removes all data associated with the given artifact.
    async fn purge_artifact(&self, aid: &ArtifactId);

    /// Add an item to the `benchmark_requests`, if the `benchmark_request`
    /// exists an Error will be returned.
    /// We require the caller to pass an index, to ensure that it is always kept up-to-date.
    async fn insert_benchmark_request(
        &self,
        benchmark_request: &BenchmarkRequest,
    ) -> anyhow::Result<BenchmarkRequestInsertResult>;

    /// Load all known benchmark request SHAs and all completed benchmark requests.
    async fn load_benchmark_request_index(&self) -> anyhow::Result<BenchmarkRequestIndex>;

    /// Load all pending benchmark requests, i.e. those that have artifacts ready, but haven't
    /// been completed yet. Pending statuses are `ArtifactsReady` and `InProgress`.
    /// Also returns their parents, so that we can quickly check which requests are ready for being
    /// enqueued.
    async fn load_pending_benchmark_requests(&self) -> anyhow::Result<PendingBenchmarkRequests>;

    /// Update the status of a `benchmark_request` with the given `tag`.
    /// If no such request exists in the DB, returns an error.
    async fn update_benchmark_request_status(
        &self,
        tag: &str,
        status: BenchmarkRequestStatus,
    ) -> anyhow::Result<()>;

    /// Reset a completed benchmark request back to ArtifactsReady so it can be
    /// re-run. Deletes completed/failed jobs and associated artifact measurement data.
    /// In-progress jobs are left untouched (a collector may still be executing them).
    /// Returns true if the request was actually reset, false if it was not found
    /// or not completed.
    async fn reset_benchmark_request_for_rerun(&mut self, tag: &str) -> anyhow::Result<bool>;

    /// Add a benchmark job to the job queue and returns its ID, if it was not
    /// already in the DB previously.
    async fn enqueue_benchmark_job(
        &self,
        request_tag: &str,
        benchmark_group: &str,
        runtime_config: &QueueRuntimeConfig,
        is_optional: bool,
        tag: &str,
    ) -> JobEnqueueResult;

    /// Returns a set of compile-time benchmark test cases that were already computed for the
    /// given artifact.
    /// Note that for efficiency reasons, the function only checks if we have at least a single
    /// result for a given test case. It does not check if *all* test results from all test
    /// iterations were finished.
    /// Therefore, the result is an over-approximation.
    async fn get_compile_test_cases_with_measurements(
        &self,
        artifact_row_id: &ArtifactIdNumber,
    ) -> anyhow::Result<HashSet<CompileTestCase>>;

    /// Returns a set of runtime benchmark names that already have measurements
    /// for the given artifact in the runtime_pstat table.
    async fn get_runtime_benchmarks_with_measurements(
        &self,
        artifact_row_id: &ArtifactIdNumber,
    ) -> anyhow::Result<HashSet<RuntimeTestCase>>;

    /// Add the confiuguration for a collector
    async fn add_collector_config(
        &self,
        collector_name: &str,
        is_active: bool,
    ) -> anyhow::Result<CollectorConfig>;

    /// Call this function when a job queue collector starts.
    /// It ensures that a collector with the given name exists, updates its commit SHA and heartbeat
    /// and returns its collector config.
    async fn start_collector(
        &self,
        collector_name: &str,
        commit_sha: &str,
    ) -> anyhow::Result<Option<CollectorConfig>>;

    /// Dequeues a single job for the given collector.
    /// The collector name is also used as the queue tag.
    /// Also returns detailed information about the compiler artifact that should be benchmarked
    /// in the job.
    async fn dequeue_benchmark_job(
        &self,
        collector_name: &str,
    ) -> anyhow::Result<Option<(BenchmarkJob, ArtifactId)>>;

    /// Try and mark the benchmark_request as completed. Will return `true` if
    /// it has been marked as completed else `false` meaning there was no change
    async fn maybe_mark_benchmark_request_as_completed(&self, tag: &str) -> anyhow::Result<bool>;

    /// Mark the job as completed. Sets the status to 'failed' or 'success'
    /// depending on the enum's completed state being a success
    async fn mark_benchmark_job_as_completed(
        &self,
        id: u32,
        conclusion: BenchmarkJobConclusion,
    ) -> anyhow::Result<()>;

    /// Return the last `count` completed benchmark requests, along with all errors associated with
    /// them.
    ///
    /// The requests will be ordered from most recently to least recently completed.
    async fn get_last_n_completed_benchmark_requests(
        &self,
        count: u64,
    ) -> anyhow::Result<Vec<BenchmarkRequestWithErrors>>;

    /// Return jobs of all requests that are currently in progress, and the jobs of their parents.
    /// The keys of the hashmap contain the request tags.
    async fn get_jobs_of_in_progress_benchmark_requests(
        &self,
    ) -> anyhow::Result<HashMap<String, Vec<BenchmarkJob>>>;

    /// Return all jobs associated with the given benchmark request tags.
    async fn get_jobs_of_benchmark_requests(
        &self,
        tags: &[String],
    ) -> anyhow::Result<Vec<BenchmarkJob>>;

    /// Get all of the configuration for all of the collectors
    async fn get_collector_configs(&self) -> anyhow::Result<Vec<CollectorConfig>>;

    /// Updates the last known heartbeat of a collector to the current time.
    async fn update_collector_heartbeat(&self, collector_name: &str) -> anyhow::Result<()>;
}

#[async_trait::async_trait]
pub trait Transaction: Send + Sync {
    fn conn(&mut self) -> &mut dyn Connection;
    fn conn_ref(&self) -> &dyn Connection;

    async fn commit(self: Box<Self>) -> Result<(), anyhow::Error>;
    async fn finish(self: Box<Self>) -> Result<(), anyhow::Error>;
}

#[async_trait::async_trait]
pub trait ConnectionManager {
    type Connection;
    async fn open(&self) -> Self::Connection;
    async fn is_valid(&self, c: &mut Self::Connection) -> bool;
}

pub struct ConnectionPool<M: ConnectionManager> {
    connections: Arc<Mutex<Vec<M::Connection>>>,
    permits: Arc<Semaphore>,
    manager: M,
}

pub struct ManagedConnection<T> {
    conn: Option<T>,
    connections: Arc<Mutex<Vec<T>>>,
    #[allow(unused)]
    permit: OwnedSemaphorePermit,
}

impl<T> std::ops::Deref for ManagedConnection<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        self.conn.as_ref().unwrap()
    }
}
impl<T> std::ops::DerefMut for ManagedConnection<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.conn.as_mut().unwrap()
    }
}

impl<T> Drop for ManagedConnection<T> {
    fn drop(&mut self) {
        let conn = self.conn.take().unwrap();
        self.connections
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .push(conn);
    }
}

impl<T, M> ConnectionPool<M>
where
    T: Send,
    M: ConnectionManager<Connection = T>,
{
    fn new(manager: M) -> Self {
        ConnectionPool {
            connections: Arc::new(Mutex::new(Vec::with_capacity(16))),
            permits: Arc::new(Semaphore::new(16)),
            manager,
        }
    }

    pub fn raw(&mut self) -> &mut M {
        &mut self.manager
    }

    async fn get(&self) -> ManagedConnection<T> {
        let permit = self.permits.clone().acquire_owned().await.unwrap();
        let conn = {
            let mut slots = self.connections.lock().unwrap_or_else(|e| e.into_inner());
            slots.pop()
        };
        if let Some(mut c) = conn {
            if self.manager.is_valid(&mut c).await {
                return ManagedConnection {
                    conn: Some(c),
                    permit,
                    connections: self.connections.clone(),
                };
            }
        }

        let conn = self.manager.open().await;
        ManagedConnection {
            conn: Some(conn),
            connections: self.connections.clone(),
            permit,
        }
    }
}

pub enum Pool {
    Sqlite(ConnectionPool<sqlite::Sqlite>),
    Postgres(ConnectionPool<postgres::Postgres>),
}

impl Pool {
    pub async fn connection(&self) -> Box<dyn Connection> {
        match self {
            Pool::Sqlite(p) => Box::new(sqlite::SqliteConnection::new(p.get().await)),
            Pool::Postgres(p) => Box::new(p.get().await),
        }
    }

    pub fn open(uri: &str) -> Pool {
        if uri.starts_with("postgres") {
            Pool::Postgres(ConnectionPool::new(postgres::Postgres::new(uri.into())))
        } else {
            Pool::Sqlite(ConnectionPool::new(sqlite::Sqlite::new(uri.into())))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::metric::Metric;
    use crate::tests::builder::{job, CollectorBuilder, RequestBuilder};
    use crate::tests::run_postgres_test;
    use crate::{tests::run_db_test, Commit, CommitType, Date};
    use chrono::Utc;
    use std::collections::BTreeSet;
    use std::str::FromStr;

    fn create_commit(commit_sha: &str, time: chrono::DateTime<Utc>, r#type: CommitType) -> Commit {
        Commit {
            sha: commit_sha.into(),
            date: Date(time),
            r#type,
        }
    }

    impl JobEnqueueResult {
        pub fn unwrap(self) -> u32 {
            match self {
                JobEnqueueResult::JobCreated(id) => id,
                error => panic!("Unexpected job enqueue result: {error:?}"),
            }
        }
    }

    #[tokio::test]
    async fn pstat_returns_empty_vector_when_empty() {
        run_db_test(|ctx| async {
            // This is essentially testing the database testing framework is
            // wired up correctly. Though makes sense that there should be
            // an empty vector returned if there are no pstats.
            let result = ctx.db().get_compile_pstats(&[], &[]).await;
            let expected: Vec<Vec<Option<f64>>> = vec![];

            assert_eq!(result, expected);
            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn artifact_storage() {
        run_db_test(|ctx| async {
            let db = ctx.db();
            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();

            let artifact_one = ArtifactId::Commit {
                repo: "rust".to_string(),
                commit: create_commit("abc", time, CommitType::Master),
            };
            let artifact_two = ArtifactId::Tag {
                repo: "rust".to_string(),
                tag: "nightly-2025-05-14".to_string(),
            };

            let artifact_one_id_number = db.artifact_id(&artifact_one, crate::DEFAULT_TAG).await;
            let artifact_two_id_number = db.artifact_id(&artifact_two, crate::DEFAULT_TAG).await;

            // We cannot arbitrarily add random sizes to the artifact size
            // table, as there is a constraint that the artifact must actually
            // exist before attaching something to it.

            // Artifact one inserts
            db.record_artifact_size(artifact_one_id_number, "llvm.so", 32)
                .await;
            db.record_artifact_size(artifact_one_id_number, "llvm.a", 64)
                .await;

            // Artifact two inserts
            db.record_artifact_size(artifact_two_id_number, "another-llvm.a", 128)
                .await;

            let result_one = db.get_artifact_size(artifact_one_id_number).await;
            let result_two = db.get_artifact_size(artifact_two_id_number).await;

            // artifact one
            assert_eq!(Some(32u64), result_one.get("llvm.so").copied());
            assert_eq!(Some(64u64), result_one.get("llvm.a").copied());
            assert_eq!(None, result_one.get("another-llvm.a").copied());

            // artifact two
            assert_eq!(Some(128), result_two.get("another-llvm.a").copied());
            Ok(ctx)
        })
        .await;
    }

    // Check that we can't have multiple requests with the same SHA
    #[tokio::test]
    async fn multiple_requests_same_sha() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            db.insert_benchmark_request(&BenchmarkRequest::create_master(
                "a-sha-1",
                "parent-sha-1",
                42,
                Utc::now(),
                "rust",
            ))
            .await
            .unwrap();

            let result = db
                .insert_benchmark_request(&BenchmarkRequest::create_release(
                    "a-sha-1",
                    Utc::now(),
                    "rust",
                ))
                .await;

            assert!(result.is_ok());
            match result.unwrap() {
                BenchmarkRequestInsertResult::NothingInserted
                | BenchmarkRequestInsertResult::Conflict { .. } => {}
                other => panic!("Expected NothingInserted or Conflict, got {other:?}"),
            }

            Ok(ctx)
        })
        .await;
    }

    // Check that we can't have multiple non-completed try requests on the same PR
    #[tokio::test]
    async fn multiple_non_completed_try_requests() {
        run_postgres_test(|mut ctx| async {
            // Insert a try build
            ctx.db()
                .insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                    42,
                    "",
                    "sha-1",
                    "sha-parent-1",
                    Utc::now(),
                ))
                .await
                .unwrap();

            // Then finish it
            ctx.complete_request("sha-1").await;

            // Reset the completed request to allow re-insertion
            let reset = ctx
                .db_mut()
                .reset_benchmark_request_for_rerun("sha-1")
                .await
                .unwrap();
            assert!(reset, "should have reset completed request");

            // Insert a try build for the same PR again — succeeds after reset
            ctx.insert_try_request(42).await;

            // But this should conflict, as we can't have two queued requests at once
            let result = ctx
                .db()
                .insert_benchmark_request(&BenchmarkRequest::create_try_without_artifacts(42, ""))
                .await
                .unwrap();

            assert!(matches!(
                result,
                BenchmarkRequestInsertResult::Conflict { .. }
            ));
            Ok(ctx)
        })
        .await;
    }

    // Check that we can't have multiple master requests on the same PR
    #[tokio::test]
    async fn multiple_master_requests_same_pr() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            db.insert_benchmark_request(&BenchmarkRequest::create_master(
                "a-sha-1",
                "parent-sha-1",
                42,
                Utc::now(),
                "rust",
            ))
            .await
            .unwrap();

            let result = db
                .insert_benchmark_request(&BenchmarkRequest::create_master(
                    "a-sha-2",
                    "parent-sha-2",
                    42,
                    Utc::now(),
                    "rust",
                ))
                .await
                .unwrap();

            assert!(matches!(
                result,
                BenchmarkRequestInsertResult::NothingInserted
                    | BenchmarkRequestInsertResult::Conflict { .. }
            ));

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn load_pending_benchmark_requests() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            // ArtifactsReady
            let req_a = ctx.insert_master_request("sha-1", "parent-sha-1", 42).await;
            // ArtifactsReady
            let req_b = ctx.insert_release_request("1.80.0").await;
            // WaitingForArtifacts
            ctx.insert_try_request(50).await;
            // InProgress
            let req_d = ctx.insert_master_request("sha-2", "parent-sha-2", 51).await;
            // Completed
            ctx.insert_release_request("1.79.0").await;

            ctx.complete_request("1.79.0").await;
            ctx.insert_master_request("parent-sha-1", "grandparent-sha-0", 100)
                .await;
            ctx.complete_request("parent-sha-1").await;
            ctx.insert_master_request("parent-sha-2", "grandparent-sha-1", 101)
                .await;
            ctx.complete_request("parent-sha-2").await;

            db.update_benchmark_request_status("sha-2", BenchmarkRequestStatus::InProgress)
                .await
                .unwrap();

            let pending = db.load_pending_benchmark_requests().await.unwrap();
            let requests = pending.requests;

            assert_eq!(requests.len(), 3);
            for req in &[req_a, req_b, req_d] {
                assert!(requests.iter().any(|r| r.tag() == req.tag()));
            }

            assert_eq!(
                pending
                    .completed_parent_tags
                    .into_iter()
                    .collect::<BTreeSet<_>>()
                    .into_iter()
                    .collect::<Vec<_>>(),
                vec!["parent-sha-1".to_string(), "parent-sha-2".to_string()]
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn enqueue_benchmark_job() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
            let benchmark_request =
                BenchmarkRequest::create_master("sha-1", "parent-sha-1", 42, time, "rust");

            // Insert the request so we don't violate the foreign key
            db.insert_benchmark_request(&benchmark_request)
                .await
                .unwrap();

            // Now we can insert the job
            let result = db
                .enqueue_benchmark_job(
                    benchmark_request.tag().unwrap(),
                    "bufreader",
                    &QueueRuntimeConfig::default(),
                    false,
                    crate::DEFAULT_TAG,
                )
                .await;
            match result {
                JobEnqueueResult::JobCreated(_) => {}
                error => panic!("Invalid result: {error:?}"),
            }

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn get_compile_test_cases_with_data() {
        run_db_test(|ctx| async {
            let db = ctx.db();

            let collection = db.collection_id("test").await;
            let artifact = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: "rust".to_string(),
                        commit: create_commit("abcdef", Utc::now(), CommitType::Try),
                    },
                    crate::DEFAULT_TAG,
                )
                .await;
            db.record_compile_benchmark("benchmark", None, "primary".to_string())
                .await;

            db.record_statistic(
                collection,
                artifact,
                "benchmark",
                Profile::Check,
                Scenario::IncrementalFresh,
                CodegenBackend::Llvm,
                Target::X86_64UnknownLinuxGnu,
                Metric::CacheMisses.as_str(),
                1.0,
            )
            .await;

            assert_eq!(
                db.get_compile_test_cases_with_measurements(&artifact)
                    .await
                    .unwrap(),
                HashSet::from([CompileTestCase {
                    benchmark: "benchmark".into(),
                    profile: Profile::Check,
                    scenario: Scenario::IncrementalFresh,
                    backend: CodegenBackend::Llvm,
                    target: Target::X86_64UnknownLinuxGnu,
                }])
            );

            let artifact2 = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: "rust".to_string(),
                        commit: create_commit("abcdef2", Utc::now(), CommitType::Try),
                    },
                    crate::DEFAULT_TAG,
                )
                .await;
            assert!(db
                .get_compile_test_cases_with_measurements(&artifact2)
                .await
                .unwrap()
                .is_empty());
            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn get_collector_config_error_if_not_exist() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            let collector_config_result = db.start_collector("collector-1", "foo").await.unwrap();

            assert!(collector_config_result.is_none());

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn add_collector_config() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            let mut inserted_config = db.add_collector_config("collector-1", true).await.unwrap();

            let config = db
                .start_collector("collector-1", "foo")
                .await
                .unwrap()
                .expect("collector config not found");

            inserted_config.commit_sha = Some("foo".to_string());
            inserted_config.last_heartbeat_at = config.last_heartbeat_at;

            assert_eq!(inserted_config, config);
            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn dequeue_benchmark_job_empty_queue() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            let benchmark_job_result = db.dequeue_benchmark_job("collector-1").await;

            assert!(benchmark_job_result.is_ok());
            assert!(benchmark_job_result.unwrap().is_none());

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn dequeue_benchmark_job() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();
            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();

            let collector_config = db.add_collector_config("collector-1", true).await.unwrap();

            let benchmark_request =
                BenchmarkRequest::create_master("sha-1", "parent-sha-1", 42, time, "rust");

            // Insert the request so we don't violate the foreign key
            db.insert_benchmark_request(&benchmark_request)
                .await
                .unwrap();

            // Now we can insert the job
            match db
                .enqueue_benchmark_job(
                    benchmark_request.tag().unwrap(),
                    "bufreader",
                    &QueueRuntimeConfig::default(),
                    false,
                    collector_config.name(),
                )
                .await
            {
                JobEnqueueResult::JobCreated(_) => {}
                error => panic!("Invalid result: {error:?}"),
            };

            let (benchmark_job, artifact_id) = db
                .dequeue_benchmark_job(collector_config.name())
                .await
                .unwrap()
                .unwrap();

            // Ensure the properties of the job match both the request and the
            // collector configuration
            assert_eq!(
                benchmark_job.request_tag(),
                benchmark_request.tag().unwrap()
            );
            assert_eq!(benchmark_job.benchmark_group(), "bufreader");
            assert_eq!(
                benchmark_job.collector_name().unwrap(),
                collector_config.name(),
            );

            assert_eq!(
                artifact_id,
                ArtifactId::Commit {
                    repo: "rust".to_string(),
                    commit: Commit {
                        sha: "sha-1".to_string(),
                        date: Date(time),
                        r#type: CommitType::Master,
                    },
                }
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn mark_request_as_complete_empty() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();
            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();

            let insert_result = db.add_collector_config("collector-1", true).await;
            assert!(insert_result.is_ok());

            let benchmark_request =
                BenchmarkRequest::create_master("sha-1", "parent-sha-1", 42, time, "rust");
            db.insert_benchmark_request(&benchmark_request)
                .await
                .unwrap();
            assert!(db
                .maybe_mark_benchmark_request_as_completed("sha-1")
                .await
                .unwrap());
            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn mark_request_as_complete() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();
            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
            let tag = "sha-1";
            let collector_name = "collector-1";

            let insert_result = db.add_collector_config(collector_name, true).await;
            assert!(insert_result.is_ok());

            /* Create the request */
            let benchmark_request = BenchmarkRequest::create_release(tag, time, "rust");
            db.insert_benchmark_request(&benchmark_request)
                .await
                .unwrap();

            /* Create job for the request */
            db.enqueue_benchmark_job(
                benchmark_request.tag().unwrap(),
                "bufreader",
                &QueueRuntimeConfig::default(),
                false,
                collector_name,
            )
            .await
            .unwrap();

            let (job, _) = db
                .dequeue_benchmark_job(collector_name)
                .await
                .unwrap()
                .unwrap();

            assert_eq!(job.request_tag(), benchmark_request.tag().unwrap());

            /* Make the job take some amount of time */
            std::thread::sleep(Duration::from_millis(1000));

            /* Mark the job as complete */
            db.mark_benchmark_job_as_completed(job.id(), BenchmarkJobConclusion::Success)
                .await
                .unwrap();

            db.maybe_mark_benchmark_request_as_completed(tag)
                .await
                .unwrap();

            /* From the status page view we can see that the duration has been
             * updated. Albeit that it will be a very short duration. */
            let completed = db.get_last_n_completed_benchmark_requests(1).await.unwrap();
            let req = &completed
                .iter()
                .find(|it| it.request.tag() == Some(tag))
                .unwrap()
                .request;

            assert!(matches!(
                req.status(),
                BenchmarkRequestStatus::Completed { .. }
            ));
            let BenchmarkRequestStatus::Completed { duration, .. } = req.status() else {
                unreachable!();
            };
            assert!(duration >= Duration::from_millis(1000));

            let completed_index = db.load_benchmark_request_index().await.unwrap();
            assert!(completed_index.contains_tag("sha-1"));

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn get_collector_configs() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            let collector_name_one = "collector-1";
            db.add_collector_config(collector_name_one, true)
                .await
                .unwrap();

            let collector_name_two = "collector-2";
            db.add_collector_config(collector_name_two, true)
                .await
                .unwrap();

            let collector_configs = db.get_collector_configs().await;
            assert!(collector_configs.is_ok());
            let collector_configs = collector_configs.unwrap();

            assert_eq!(collector_configs[0].name(), collector_name_one);
            assert!(collector_configs[0].is_active());

            assert_eq!(collector_configs[1].name(), collector_name_two);
            assert!(collector_configs[1].is_active());

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn get_last_completed_requests() {
        run_postgres_test(|ctx| async {
            let mut requests = vec![];
            let db = ctx.db();

            let collector = ctx.add_collector(Default::default()).await;

            // Create several completed requests
            for id in 1..=3 {
                // Make some space between completions
                tokio::time::sleep(Duration::from_millis(100)).await;

                requests.push(
                    RequestBuilder::master(db, &format!("sha{id}"), &format!("sha{}", id - 1), id)
                        .await
                        .add_job(db, job())
                        .await
                        .complete(db, &collector)
                        .await,
                );
            }

            // Create an additional non-completed request
            ctx.insert_master_request("foo", "bar", 1000).await;

            // Request 1 will have artifact with errors
            let aid1 = ctx.upsert_master_artifact("sha1").await;
            db.record_error(aid1, "crate1", "error1", None).await;
            db.record_error(aid1, "crate2", "error2", None).await;

            // Request 2 will have artifact without errors
            let _aid2 = ctx.upsert_master_artifact("sha2").await;

            // Request 3 will have no artifact (shouldn't happen in practice, but...)

            let reqs = db.get_last_n_completed_benchmark_requests(5).await.unwrap();
            assert_eq!(reqs.len(), 3);

            let expected = [
                ("sha3", HashMap::new()),
                ("sha2", HashMap::new()),
                (
                    "sha1",
                    HashMap::from([
                        ("crate1".to_string(), "error1".to_string()),
                        ("crate2".to_string(), "error2".to_string()),
                    ]),
                ),
            ];
            for ((sha, errors), req) in expected.into_iter().zip(reqs) {
                assert_eq!(
                    req.request.tag().unwrap(),
                    sha,
                    "Request {req:?} does not have expected sha {sha}"
                );
                assert_eq!(
                    req.errors, errors,
                    "Request {req:?} does not have expected errors {errors:?}"
                );
            }

            let reqs = db.get_last_n_completed_benchmark_requests(1).await.unwrap();
            assert_eq!(reqs.len(), 1);
            assert_eq!(reqs[0].request.tag().unwrap(), "sha3");

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn get_in_progress_jobs() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            let collector = ctx.add_collector(Default::default()).await;

            // Artifacts ready request, should be ignored
            RequestBuilder::master(db, "foo", "bar", 1000).await;

            // Create a completed parent with jobs
            let completed = RequestBuilder::master(db, "sha4-parent", "sha0", 1001)
                .await
                .add_jobs(
                    db,
                    &[job().benchmark_group("doc"), job().benchmark_group("opt")],
                )
                .await
                .complete(db, &collector)
                .await;

            // In progress request without a parent
            let req1 = RequestBuilder::master(db, "sha1", "sha0", 1)
                .await
                .set_in_progress(db)
                .await;

            // In progress request with a parent that has no jobs
            let req2 = RequestBuilder::master(db, "sha2", "sha1", 2)
                .await
                .add_jobs(
                    db,
                    &[
                        job().benchmark_group("check"),
                        job().benchmark_group("debug"),
                    ],
                )
                .await
                .set_in_progress(db)
                .await;

            // In progress request with a parent that has jobs
            let req3 = RequestBuilder::master(db, "sha3", "sha2", 3)
                .await
                .add_jobs(
                    db,
                    &[job().benchmark_group("doc"), job().benchmark_group("opt")],
                )
                .await
                .set_in_progress(db)
                .await;

            // In progress request with a parent that has jobs, but is completed
            let req4 = RequestBuilder::master(db, "sha4", completed.tag(), 4)
                .await
                .add_jobs(
                    db,
                    &[job().benchmark_group("doc"), job().benchmark_group("check")],
                )
                .await
                .set_in_progress(db)
                .await;

            let mut reqs = db
                .get_jobs_of_in_progress_benchmark_requests()
                .await
                .unwrap();

            // Check that all jobs are unique
            let mut job_ids = HashSet::new();
            for job in reqs.values().flatten() {
                assert!(job_ids.insert(job.id));
            }

            // Check that all jobs were returned
            assert!(!reqs.contains_key(req1.tag()));
            req2.assert_has_exact_jobs(&reqs.remove(req2.tag()).unwrap());
            req3.assert_has_exact_jobs(&reqs.remove(req3.tag()).unwrap());
            req4.assert_has_exact_jobs(&reqs.remove(req4.tag()).unwrap());
            completed.assert_has_exact_jobs(&reqs.remove(completed.tag()).unwrap());
            assert!(reqs.is_empty());

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn purge_artifact() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            ctx.upsert_master_artifact("foo").await;
            ctx.insert_master_request("foo", "bar", 1).await;
            db.enqueue_benchmark_job(
                "foo",
                "bufreader",
                &QueueRuntimeConfig::default(),
                false,
                crate::DEFAULT_TAG,
            )
            .await
            .unwrap();
            db.purge_artifact(&ArtifactId::Tag {
                repo: "rust".to_string(),
                tag: "foo".to_string(),
            })
            .await;

            assert!(!db
                .load_benchmark_request_index()
                .await
                .unwrap()
                .contains_tag("foo"));

            let collector = ctx.add_collector(CollectorBuilder::default()).await;
            assert!(db
                .dequeue_benchmark_job(collector.name())
                .await
                .unwrap()
                .is_none());

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn optional_jobs_should_not_block_benchmark_request_from_being_completed() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();
            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
            let tag = "sha-1";
            let collector_name = "collector-1";

            let insert_result = db.add_collector_config(collector_name, true).await;
            assert!(insert_result.is_ok());

            /* Create the request */
            let benchmark_request = BenchmarkRequest::create_release(tag, time, "rust");
            db.insert_benchmark_request(&benchmark_request)
                .await
                .unwrap();

            /* Create required job for the request */
            db.enqueue_benchmark_job(
                benchmark_request.tag().unwrap(),
                "bufreader",
                &QueueRuntimeConfig::default(),
                false,
                collector_name,
            )
            .await
            .unwrap();

            /* Create optional job for the request */
            db.enqueue_benchmark_job(
                benchmark_request.tag().unwrap(),
                "sort",
                &QueueRuntimeConfig::default(),
                true,
                collector_name,
            )
            .await
            .unwrap();

            let (job, _) = db
                .dequeue_benchmark_job(collector_name)
                .await
                .unwrap()
                .unwrap();

            assert_eq!(job.request_tag(), benchmark_request.tag().unwrap());

            /* Make the job take some amount of time */
            tokio::time::sleep(Duration::from_millis(1000)).await;

            /* Mark the job as complete */
            db.mark_benchmark_job_as_completed(job.id(), BenchmarkJobConclusion::Success)
                .await
                .unwrap();

            db.maybe_mark_benchmark_request_as_completed(tag)
                .await
                .unwrap();

            /* From the status page view we can see that the duration has been
             * updated. Albeit that it will be a very short duration. */
            let completed = db.get_last_n_completed_benchmark_requests(1).await.unwrap();
            let req = &completed
                .iter()
                .find(|it| it.request.tag() == Some(tag))
                .unwrap()
                .request;

            assert!(matches!(
                req.status(),
                BenchmarkRequestStatus::Completed { .. }
            ));
            let BenchmarkRequestStatus::Completed { duration, .. } = req.status() else {
                unreachable!();
            };
            assert!(duration >= Duration::from_millis(1000));

            let completed_index = db.load_benchmark_request_index().await.unwrap();
            assert!(completed_index.contains_tag("sha-1"));

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn runtime_statistic_storage_and_retrieval() {
        run_db_test(|mut ctx| async {
            let collection = ctx.db().collection_id("test").await;
            let artifact = ctx
                .db()
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: "rust".to_string(),
                        commit: create_commit("abc123", Utc::now(), CommitType::Try),
                    },
                    crate::DEFAULT_TAG,
                )
                .await;

            ctx.db()
                .record_runtime_statistic(
                    collection,
                    artifact,
                    "my_bench",
                    "wall-time",
                    Target::X86_64UnknownLinuxGnu,
                    42.5,
                )
                .await;

            let index = crate::Index::load(ctx.db_mut()).await;
            let series: Vec<u32> = index
                .runtime_statistic_descriptions()
                .filter(|((bench, _, metric), _)| {
                    bench.as_str() == "my_bench" && metric.as_str() == "wall-time"
                })
                .map(|(_, id)| id)
                .collect();

            assert!(!series.is_empty());
            let values = ctx
                .db()
                .get_runtime_pstats(&series, &[Some(artifact)])
                .await;
            assert_eq!(values.len(), 1);
            assert_eq!(values[0].len(), 1);
            assert!((values[0][0].unwrap() - 42.5).abs() < f64::EPSILON);

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn runtime_json_statistic_storage_and_retrieval() {
        run_db_test(|mut ctx| async {
            let collection = ctx.db().collection_id("test").await;
            let artifact = ctx
                .db()
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: "rust".to_string(),
                        commit: create_commit("json123", Utc::now(), CommitType::Try),
                    },
                    crate::DEFAULT_TAG,
                )
                .await;

            let json_data = r#"{"qps":1234,"p99_ms":12.5}"#;
            ctx.db()
                .record_runtime_json_statistic(
                    collection,
                    artifact,
                    "vllm_bench",
                    "perf_detail",
                    Target::X86_64UnknownLinuxGnu,
                    json_data,
                )
                .await;

            let index = crate::Index::load(ctx.db_mut()).await;
            let series: Vec<u32> = index
                .runtime_statistic_descriptions()
                .filter(|((bench, _, metric), _)| {
                    bench.as_str() == "vllm_bench" && metric.as_str() == "perf_detail"
                })
                .map(|(_, id)| id)
                .collect();

            assert!(!series.is_empty());
            let json_values = ctx
                .db()
                .get_runtime_json_pstats(&series, &[Some(artifact)])
                .await;
            assert_eq!(json_values.len(), 1);
            assert_eq!(json_values[0].len(), 1);
            let retrieved = json_values[0][0].as_ref().unwrap();
            assert_eq!(retrieved, json_data);

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn runtime_json_pstats_empty_returns_none() {
        run_db_test(|ctx| async {
            let db = ctx.db();
            let artifact = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: "rust".to_string(),
                        commit: create_commit("empty123", Utc::now(), CommitType::Try),
                    },
                    crate::DEFAULT_TAG,
                )
                .await;

            // No series exist, should return empty
            let result = db.get_runtime_json_pstats(&[], &[Some(artifact)]).await;
            assert!(result.is_empty());

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn runtime_numeric_and_json_coexist() {
        run_db_test(|mut ctx| async {
            let collection = ctx.db().collection_id("test").await;
            let artifact = ctx
                .db()
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: "rust".to_string(),
                        commit: create_commit("coexist", Utc::now(), CommitType::Try),
                    },
                    crate::DEFAULT_TAG,
                )
                .await;

            // Record numeric metric
            ctx.db()
                .record_runtime_statistic(
                    collection,
                    artifact,
                    "bench1",
                    "wall-time",
                    Target::X86_64UnknownLinuxGnu,
                    100.0,
                )
                .await;

            // Record JSON metric for same benchmark
            ctx.db()
                .record_runtime_json_statistic(
                    collection,
                    artifact,
                    "bench1",
                    "detail",
                    Target::X86_64UnknownLinuxGnu,
                    r#"{"extra":"data"}"#,
                )
                .await;

            let index = crate::Index::load(ctx.db_mut()).await;

            // Check numeric
            let num_series: Vec<u32> = index
                .runtime_statistic_descriptions()
                .filter(|((b, _, m), _)| b.as_str() == "bench1" && m.as_str() == "wall-time")
                .map(|(_, id)| id)
                .collect();
            let num_values = ctx
                .db()
                .get_runtime_pstats(&num_series, &[Some(artifact)])
                .await;
            assert_eq!(num_values[0][0].unwrap(), 100.0);

            // Check JSON
            let json_series: Vec<u32> = index
                .runtime_statistic_descriptions()
                .filter(|((b, _, m), _)| b.as_str() == "bench1" && m.as_str() == "detail")
                .map(|(_, id)| id)
                .collect();
            let json_values = ctx
                .db()
                .get_runtime_json_pstats(&json_series, &[Some(artifact)])
                .await;
            assert_eq!(json_values[0][0].as_ref().unwrap(), r#"{"extra":"data"}"#);

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn index_repo_isolation() {
        run_db_test(|mut ctx| async move {
            let db = ctx.db();
            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();

            let repo_rust = "github/rust-lang/rust";
            let repo_daft = "gitcode/ywxtcwh/Daft";

            let sha_shared = "abc123";
            let sha_rust_only = "rust_only";
            let sha_daft_only = "daft_only";

            let tag_default = crate::DEFAULT_TAG;
            let tag_collector = "Kunpeng 920B";

            // Insert artifacts for rust repo
            let _rust_shared = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: repo_rust.to_string(),
                        commit: create_commit(sha_shared, time, CommitType::Master),
                    },
                    tag_default,
                )
                .await;
            let _rust_only = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: repo_rust.to_string(),
                        commit: create_commit(sha_rust_only, time, CommitType::Master),
                    },
                    tag_default,
                )
                .await;

            // Insert artifacts for daft repo (same SHA as rust + unique ones)
            let _daft_shared = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: repo_daft.to_string(),
                        commit: create_commit(sha_shared, time, CommitType::Master),
                    },
                    tag_default,
                )
                .await;
            let _daft_collector = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: repo_daft.to_string(),
                        commit: create_commit(sha_daft_only, time, CommitType::Master),
                    },
                    tag_collector,
                )
                .await;

            // Load index
            let index = ctx.db_mut().load_index().await;

            // commits_for_repo should only return commits for the specified repo
            let rust_commits = index.commits_for_repo(repo_rust);
            let daft_commits = index.commits_for_repo(repo_daft);

            assert_eq!(rust_commits.len(), 2, "rust should have 2 commits");
            assert_eq!(daft_commits.len(), 2, "daft should have 2 commits");

            let rust_shas: Vec<&str> = rust_commits.iter().map(|c| c.sha.as_str()).collect();
            let daft_shas: Vec<&str> = daft_commits.iter().map(|c| c.sha.as_str()).collect();

            assert!(rust_shas.contains(&sha_shared));
            assert!(rust_shas.contains(&sha_rust_only));
            assert!(!rust_shas.contains(&sha_daft_only));

            assert!(daft_shas.contains(&sha_shared));
            assert!(daft_shas.contains(&sha_daft_only));
            assert!(!daft_shas.contains(&sha_rust_only));

            // commits() returns all commits (deduplicated by SHA)
            let all_commits = index.commits();
            assert_eq!(all_commits.len(), 3, "all commits deduplicated");

            // lookup_artifact_with_tag should find by (repo, sha, tag)
            let rust_shared_aid = ArtifactId::Commit {
                repo: repo_rust.to_string(),
                commit: create_commit(sha_shared, time, CommitType::Master),
            };
            let daft_shared_aid = ArtifactId::Commit {
                repo: repo_daft.to_string(),
                commit: create_commit(sha_shared, time, CommitType::Master),
            };

            assert!(
                index
                    .lookup_artifact_with_tag(&rust_shared_aid, tag_default.parse().unwrap())
                    .is_some(),
                "should find rust/shared/default"
            );
            assert!(
                index
                    .lookup_artifact_with_tag(&daft_shared_aid, tag_default.parse().unwrap())
                    .is_some(),
                "should find daft/shared/default"
            );

            // Same SHA but different repo should return different IDs
            let rust_id = index
                .lookup_artifact_with_tag(&rust_shared_aid, tag_default.parse().unwrap())
                .unwrap();
            let daft_id = index
                .lookup_artifact_with_tag(&daft_shared_aid, tag_default.parse().unwrap())
                .unwrap();
            assert_ne!(
                rust_id, daft_id,
                "same SHA in different repos must have different artifact IDs"
            );

            // tags_for_commit should work with repo
            let tags = index.tags_for_commit_in_repo(
                repo_daft,
                &create_commit(sha_daft_only, time, CommitType::Master),
            );
            assert_eq!(tags.len(), 1);
            assert_eq!(tags[0].as_str(), tag_collector);

            // tags_for_repo should return only tags for that repo
            let rust_tags = index.tags_for_repo(repo_rust);
            let daft_tags = index.tags_for_repo(repo_daft);

            assert!(rust_tags.iter().any(|t| t.as_str() == tag_default));
            assert!(daft_tags.iter().any(|t| t.as_str() == tag_default));
            assert!(daft_tags.iter().any(|t| t.as_str() == tag_collector));
            assert!(!rust_tags.iter().any(|t| t.as_str() == tag_collector));

            // artifact_id_for_commit should be repo-scoped
            let found = index.artifact_id_for_commit(sha_daft_only, repo_daft);
            assert!(found.is_some(), "should find daft_only in daft repo");

            let not_found = index.artifact_id_for_commit(sha_daft_only, repo_rust);
            assert!(
                not_found.is_none(),
                "should NOT find daft_only in rust repo"
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn index_tag_isolation() {
        run_db_test(|mut ctx| async move {
            let db = ctx.db();
            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();

            let repo = "github/rust-lang/rust";
            let sha = "abc123";
            let tag_default = crate::DEFAULT_TAG;
            let tag_collector_a = "collector-a";
            let tag_collector_b = "collector-b";

            // Same commit, same repo, different tags
            let _aid_default = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: repo.to_string(),
                        commit: create_commit(sha, time, CommitType::Master),
                    },
                    tag_default,
                )
                .await;
            let _aid_a = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: repo.to_string(),
                        commit: create_commit(sha, time, CommitType::Master),
                    },
                    tag_collector_a,
                )
                .await;
            let _aid_b = db
                .artifact_id(
                    &ArtifactId::Commit {
                        repo: repo.to_string(),
                        commit: create_commit(sha, time, CommitType::Master),
                    },
                    tag_collector_b,
                )
                .await;

            let index = ctx.db_mut().load_index().await;

            // commits_for_repo should deduplicate by commit SHA (1 commit, not 3)
            let commits = index.commits_for_repo(repo);
            assert_eq!(commits.len(), 1);

            // tags_for_commit_in_repo should return all 3 tags
            let tags =
                index.tags_for_commit_in_repo(repo, &create_commit(sha, time, CommitType::Master));
            assert_eq!(tags.len(), 3);

            let tag_strs: Vec<&str> = tags.iter().map(|t| t.as_str()).collect();
            assert!(tag_strs.contains(&tag_default));
            assert!(tag_strs.contains(&tag_collector_a));
            assert!(tag_strs.contains(&tag_collector_b));

            // lookup_artifact_with_tag should distinguish by tag
            let aid = ArtifactId::Commit {
                repo: repo.to_string(),
                commit: create_commit(sha, time, CommitType::Master),
            };

            let id_default = index
                .lookup_artifact_with_tag(&aid, tag_default.parse().unwrap())
                .unwrap();
            let id_a = index
                .lookup_artifact_with_tag(&aid, tag_collector_a.parse().unwrap())
                .unwrap();
            let id_b = index
                .lookup_artifact_with_tag(&aid, tag_collector_b.parse().unwrap())
                .unwrap();

            assert_ne!(id_default, id_a);
            assert_ne!(id_default, id_b);
            assert_ne!(id_a, id_b);

            // Non-existent tag should return None
            let missing = index.lookup_artifact_with_tag(&aid, "nonexistent".parse().unwrap());
            assert!(missing.is_none());

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn reset_completed_request_clears_artifact_and_measurements() {
        run_postgres_test(|mut ctx| async {
            let collector = ctx.add_collector(Default::default()).await;

            let sha = "reset-sha-1";
            let req = {
                let db = ctx.db();
                let req = RequestBuilder::master(db, sha, "parent-1", 100)
                    .await
                    .add_job(db, job())
                    .await
                    .complete(db, &collector)
                    .await;

                let artifact = ctx.upsert_master_artifact(sha).await;
                let collection = db.collection_id("test").await;
                db.record_runtime_statistic(
                    collection,
                    artifact,
                    "bench-1",
                    "wall-time",
                    Target::X86_64UnknownLinuxGnu,
                    42.0,
                )
                .await;

                let measured = db
                    .get_runtime_benchmarks_with_measurements(&artifact)
                    .await
                    .unwrap();
                assert_eq!(measured.len(), 1, "should have measurement before reset");

                req
            };

            let reset = ctx
                .db_mut()
                .reset_benchmark_request_for_rerun(req.tag())
                .await
                .unwrap();
            assert!(reset, "should have reset a completed request");

            let artifact = ctx.upsert_master_artifact(sha).await;
            let measured_after = ctx
                .db()
                .get_runtime_benchmarks_with_measurements(&artifact)
                .await
                .unwrap();
            assert!(
                measured_after.is_empty(),
                "measurements should be cleared after reset"
            );

            let reinserted = ctx
                .db()
                .insert_benchmark_request(&BenchmarkRequest::create_master(
                    sha,
                    "parent-1",
                    100,
                    Utc::now(),
                    "rust",
                ))
                .await
                .unwrap();
            assert!(
                matches!(reinserted, BenchmarkRequestInsertResult::Inserted),
                "should be able to re-insert after request was deleted"
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn reset_non_completed_request_is_noop() {
        run_postgres_test(|mut ctx| async {
            {
                let db = ctx.db();
                // Create a try request with a known SHA but in ArtifactsReady
                // (non-completed) status.
                let request = BenchmarkRequest::create_try_with_artifacts(
                    200,
                    "rust",
                    "non-completed-sha-200",
                    "parent-200",
                    Utc::now(),
                );
                db.insert_benchmark_request(&request).await.unwrap();
            }

            let reset = ctx
                .db_mut()
                .reset_benchmark_request_for_rerun("non-completed-sha-200")
                .await
                .unwrap();
            assert!(
                !reset,
                "should not reset an ArtifactsReady (non-completed) request"
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn cross_repo_same_pr_no_conflict() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            // Insert a try request for repo A with pr=42
            db.insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                42,
                "owner/repo-a",
                "sha-a-1",
                "parent-a-1",
                Utc::now(),
            ))
            .await
            .unwrap();

            // Insert a try request for repo B with the same pr=42 — should succeed
            // because (pr, commit_type, repo) is different
            let result = db
                .insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                    42,
                    "owner/repo-b",
                    "sha-b-1",
                    "parent-b-1",
                    Utc::now(),
                ))
                .await
                .unwrap();

            assert!(
                matches!(result, BenchmarkRequestInsertResult::Inserted),
                "different repos with same PR should not conflict, got {:?}",
                result
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn same_repo_same_pr_conflicts() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            // Insert a try request for repo A with pr=42
            db.insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                42,
                "owner/repo-a",
                "sha-a-1",
                "parent-a-1",
                Utc::now(),
            ))
            .await
            .unwrap();

            // Try to insert another try request for the same repo with same pr=42
            let result = db
                .insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                    42,
                    "owner/repo-a",
                    "sha-a-2",
                    "parent-a-2",
                    Utc::now(),
                ))
                .await
                .unwrap();

            assert!(
                matches!(result, BenchmarkRequestInsertResult::Conflict { .. }),
                "same repo with same PR should conflict, got {:?}",
                result
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn null_tag_conflict_detection() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            // Insert a try request without artifacts (tag=NULL)
            db.insert_benchmark_request(&BenchmarkRequest::create_try_without_artifacts(42, ""))
                .await
                .unwrap();

            // Try to insert another try request for same repo/pr without artifacts
            let result = db
                .insert_benchmark_request(&BenchmarkRequest::create_try_without_artifacts(42, ""))
                .await
                .unwrap();

            assert!(
                matches!(result, BenchmarkRequestInsertResult::Conflict { .. }),
                "same repo/PR with NULL tag should conflict, got {:?}",
                result
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn same_tag_always_conflicts() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            // Insert and complete a try request with tag "sha-1"
            db.insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                42,
                "owner/repo-a",
                "sha-1",
                "parent-1",
                Utc::now(),
            ))
            .await
            .unwrap();
            ctx.complete_request("sha-1").await;

            // Same tag should conflict even for a completed request
            let result = db
                .insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                    99,
                    "owner/repo-b",
                    "sha-1",
                    "parent-1b",
                    Utc::now(),
                ))
                .await
                .unwrap();

            assert!(
                matches!(result, BenchmarkRequestInsertResult::Conflict { .. }),
                "same tag should always conflict, got {:?}",
                result
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn completed_request_does_not_conflict() {
        run_postgres_test(|ctx| async {
            let db = ctx.db();

            db.insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                42,
                "owner/repo-a",
                "sha-42",
                "parent-42",
                Utc::now(),
            ))
            .await
            .unwrap();
            ctx.complete_request("sha-42").await;

            // Completed request with same (pr, commit_type, repo) should NOT
            // conflict — a new request with a different tag can be inserted.
            let result = db
                .insert_benchmark_request(&BenchmarkRequest::create_try_with_artifacts(
                    42,
                    "owner/repo-a",
                    "sha-42-new",
                    "parent-42-new",
                    Utc::now(),
                ))
                .await
                .unwrap();

            assert!(
                matches!(result, BenchmarkRequestInsertResult::Inserted),
                "completed request should not conflict, got {:?}",
                result
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn reset_nonexistent_request_is_noop() {
        run_postgres_test(|mut ctx| async {
            let reset = ctx
                .db_mut()
                .reset_benchmark_request_for_rerun("does-not-exist")
                .await
                .unwrap();
            assert!(!reset, "should not reset a nonexistent request");

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn reset_clears_all_jobs_and_resets_status() {
        run_postgres_test(|mut ctx| async {
            let collector = ctx.add_collector(Default::default()).await;

            let sha = "reset-jobs-sha";
            let req = {
                let db = ctx.db();
                RequestBuilder::master(db, sha, "parent-3", 300)
                    .await
                    .add_jobs(
                        db,
                        &[
                            job().benchmark_group("group-a"),
                            job().benchmark_group("group-b"),
                        ],
                    )
                    .await
                    .complete(db, &collector)
                    .await
            };

            ctx.db_mut()
                .reset_benchmark_request_for_rerun(req.tag())
                .await
                .unwrap();

            let reinserted = ctx
                .db()
                .insert_benchmark_request(&BenchmarkRequest::create_master(
                    sha,
                    "parent-3",
                    300,
                    Utc::now(),
                    "rust",
                ))
                .await
                .unwrap();
            assert!(
                matches!(reinserted, BenchmarkRequestInsertResult::Inserted),
                "should be able to re-insert after request was deleted"
            );

            let pending = ctx.db().load_pending_benchmark_requests().await.unwrap();
            let found = pending.requests.iter().find(|r| r.tag() == Some(sha));
            assert!(
                found.is_some(),
                "re-inserted request should appear in pending"
            );
            assert!(
                matches!(
                    found.unwrap().status(),
                    BenchmarkRequestStatus::ArtifactsReady
                ),
                "re-inserted request should be ArtifactsReady"
            );

            Ok(ctx)
        })
        .await;
    }

    #[tokio::test]
    async fn reset_in_progress_request_is_noop() {
        run_postgres_test(|mut ctx| async {
            let sha = "inprogress-sha";
            {
                let db = ctx.db();
                RequestBuilder::master(db, sha, "parent-ip", 400)
                    .await
                    .add_jobs(
                        db,
                        &[
                            job().benchmark_group("group-a"),
                            job().benchmark_group("group-b"),
                        ],
                    )
                    .await
                    .set_in_progress(db)
                    .await;
            }

            let reset = ctx
                .db_mut()
                .reset_benchmark_request_for_rerun(sha)
                .await
                .unwrap();
            assert!(!reset, "should not reset an in-progress request");

            Ok(ctx)
        })
        .await;
    }
}