use anyhow::Context;
use crates_io::db;
use crates_io::schema::background_jobs;
use crates_io::schema::{crates, default_versions, versions};
use crates_io::worker::jobs::BuildCrateZip;
use crates_io_worker::BackgroundJob;
use diesel::prelude::*;
use diesel_async::{AsyncPgConnection, RunQueryDsl};
use indicatif::{ProgressBar, ProgressStyle};

const CHUNK_SIZE: usize = 100;

#[derive(clap::Parser, Debug, Eq, PartialEq)]
#[clap(
    name = "build-crate-zips",
    about = "Queue background jobs to build the zip source archive for uploaded crate files.",
    group(clap::ArgGroup::new("mode").required(true))
)]
pub struct Options {
    /// Backfill *all* versions that are missing a zip source archive.
    #[clap(long, group = "mode")]
    backfill: bool,

    /// Crate specifications to build (format: `crate@version` or just `crate`)
    #[clap(group = "mode")]
    crates: Vec<String>,
}

pub async fn run(opts: Options) -> anyhow::Result<()> {
    let conn = db::oneoff_connection().await;
    let mut conn = conn.context("Failed to connect to the database")?;

    let results = load_versions(&opts, &conn).await;
    let results = results.context("Failed to load versions")?;
    if results.is_empty() {
        println!("No matching versions found.");
        return Ok(());
    }

    println!("Found {} matching versions", results.len());
    if opts.backfill {
        let default_count = results.iter().filter(|(_, c)| *c).count();
        println!("  {default_count} default versions (priority -20)");

        let regular_count = results.len() - default_count;
        println!("  {regular_count} regular versions (priority -50)");
    }

    let pb_style = ProgressStyle::with_template("{bar:60} ({pos}/{len}, ETA {eta})")?;
    let pb = ProgressBar::new(results.len() as u64).with_style(pb_style);

    let mut queued_count = 0;
    let mut error_count = 0;

    for batch in results.chunks(CHUNK_SIZE) {
        let jobs = batch
            .iter()
            .map(|(version_id, is_default_version)| {
                let priority = if *is_default_version { -20 } else { -50 };
                NewBackgroundJob::new(*version_id, priority)
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        let num_jobs = jobs.len();

        let result = diesel::insert_into(background_jobs::table)
            .values(&jobs)
            .execute(&mut conn)
            .await;

        pb.inc(num_jobs as u64);

        if let Err(err) = result {
            error_count += num_jobs;
            pb.suspend(|| eprintln!("Failed to queue background jobs: {err}"));
        } else {
            queued_count += num_jobs;
        }
    }

    pb.finish_with_message("Done");

    println!("Successfully queued {queued_count} jobs");
    if error_count > 0 {
        println!("Failed to queue {error_count} jobs");
    }

    Ok(())
}

async fn load_versions(
    opts: &Options,
    mut conn: &AsyncPgConnection,
) -> QueryResult<Vec<(i32, bool)>> {
    let mut query = versions::table
        .inner_join(crates::table)
        .left_join(default_versions::table.on(default_versions::version_id.eq(versions::id)))
        .select((
            versions::id,
            default_versions::crate_id.nullable().is_not_null(),
        ))
        .into_boxed();

    if opts.backfill {
        query = query.filter(
            versions::zip_sha256
                .is_null()
                .or(versions::zip_json_sha256.is_null()),
        )
    } else {
        for crate_spec in &opts.crates {
            let (krate, version) = parse_crate_spec(crate_spec);

            query = match version {
                Some(ver) => query.or_filter(crates::name.eq(krate).and(versions::num.eq(ver))),
                None => query.or_filter(crates::name.eq(krate)),
            };
        }
    }

    query.load(&mut conn).await
}

/// Parses crate specification in the format "crate@version" or just "crate"
fn parse_crate_spec(spec: &str) -> (&str, Option<&str>) {
    if let Some((name, ver)) = spec.split_once('@') {
        (name, Some(ver))
    } else {
        (spec, None)
    }
}

/// Represents a new background job to be inserted into the database
#[derive(Debug, diesel::Insertable)]
#[diesel(table_name = background_jobs)]
struct NewBackgroundJob {
    job_type: &'static str,
    data: serde_json::Value,
    priority: i16,
}

impl NewBackgroundJob {
    /// Creates a new [`BuildCrateZip`] background job with the specified priority
    fn new(version_id: i32, priority: i16) -> anyhow::Result<Self> {
        let job = BuildCrateZip::new(version_id);
        let data = serde_json::to_value(&job).context("Failed to serialize job data")?;

        Ok(Self {
            job_type: BuildCrateZip::JOB_NAME,
            data,
            priority,
        })
    }
}
