// Copyright (C) 2025 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Script to automate routine 3rd party Rust crate updates.
//
// Usage:
//
// * Check out a clean copy of main-without-vendor, dedicated to
//   crate updates. See go/repo-init/main-without-vendor.
//
// * `cargo run -- <path to repo for crate updates>`
//   For example, `cargo run -- ~/src/main-for-crate-updates/`
//
// The crate updater can also be run on the cron, but requires
// a workstation or cloudtop with a valid credential from running gcert.
// Therefore, it is suggested that the crontab entry be configured
// to run shortly after the end of your regular work day.
//
// Example crontab entry:
//
// 0 2 * * * cargo run --manifest-path=$HOME/src/main-without-vendor/development/tools/crate-updater/Cargo.toml -- $HOME/src/main-for-crate-updates &> $HOME/crate-updater-`date +"\%Y-\%m-\%d"`.log

use std::{
    collections::{BTreeMap, HashMap},
    env,
    path::{Path, PathBuf},
    process::{Command, ExitStatus, Output},
    str::from_utf8,
};

use anyhow::{bail, Result};
use chrono::Datelike;
use clap::Parser;
use crate_updater::UpdatesTried;
use rand::seq::SliceRandom;
use rand::thread_rng;
use regex::Regex;
use serde::Deserialize;
#[derive(Parser)]
struct Cli {
    /// Absolute path to a repo checkout of main-without-vendor.
    /// It is strongly recommended that you use a source tree dedicated to
    /// running this updater.
    android_root: PathBuf,
    /// Send generated email to rotation mailing list
    #[arg(long, default_value_t = false)]
    rotation: bool,
}

pub trait SuccessOrError {
    fn success_or_error(self) -> Result<Self>
    where
        Self: std::marker::Sized;
}
impl SuccessOrError for ExitStatus {
    fn success_or_error(self) -> Result<Self> {
        if !self.success() {
            let exit_code =
                self.code().map(|code| format!("{code}")).unwrap_or("(unknown)".to_string());
            bail!("Process failed with exit code {exit_code}");
        }
        Ok(self)
    }
}
impl SuccessOrError for Output {
    fn success_or_error(self) -> Result<Self> {
        (&self).success_or_error()?;
        Ok(self)
    }
}
impl SuccessOrError for &Output {
    fn success_or_error(self) -> Result<Self> {
        if !self.status.success() {
            let exit_code =
                self.status.code().map(|code| format!("{code}")).unwrap_or("(unknown)".to_string());
            bail!(
                "Process failed with exit code {}\nstdout:\n{}\nstderr:\n{}",
                exit_code,
                from_utf8(&self.stdout)?,
                from_utf8(&self.stderr)?
            );
        }
        Ok(self)
    }
}

pub trait RunAndStreamOutput {
    fn run_and_stream_output(&mut self) -> Result<ExitStatus>;
}
impl RunAndStreamOutput for Command {
    fn run_and_stream_output(&mut self) -> Result<ExitStatus> {
        self.spawn()?.wait()?.success_or_error()
    }
}

fn cleanup_and_sync_monorepo(monorepo_path: &Path) -> Result<()> {
    Command::new("git")
        .args(["restore", "--staged", "."])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;
    Command::new("git")
        .args(["restore", "."])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;

    Command::new("git")
        .args(["clean", "-f", "-d"])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;

    Command::new("git")
        .args(["checkout", "goog/main"])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;

    let output = Command::new("git")
        .args(["status", "--porcelain", "."])
        .current_dir(monorepo_path)
        .output()?
        .success_or_error()?;
    if !output.stdout.is_empty() {
        bail!("Monorepo {} has uncommitted changes", monorepo_path.display());
    }

    Command::new("repo").args(["sync", "."]).current_dir(monorepo_path).run_and_stream_output()?;

    Ok(())
}

#[derive(Deserialize, Default, Debug)]
struct UpdateSuggestions {
    updates: Vec<UpdateSuggestion>,
}

#[derive(Deserialize, Default, Debug)]
struct UpdateSuggestion {
    name: String,
    version: String,
}

fn smartsync(monorepo_path: &Path, bid: &str) -> Result<ExitStatus> {
    Command::new("/google/data/ro/projects/android/smartsync_repo")
        .args(["sync", "-j32", "--retry-fetches=2", "-t", bid])
        .current_dir(monorepo_path)
        .run_and_stream_output()
}

fn sync_to_green(monorepo_path: &Path) -> Result<()> {
    Command::new("prodcertstatus").run_and_stream_output()?;
    Command::new("/google/data/ro/projects/android/smartsync_login").run_and_stream_output()?;

    let output = Command::new("/google/data/ro/projects/android/ab")
        .args([
            "lkgb",
            "--branch=git_main-without-vendor",
            "--target=aosp_arm64-trunk_staging-userdebug",
            "--raw",
            "--custom_raw_format={o[buildId]}",
        ])
        .output()?
        .success_or_error()?;
    let bid = from_utf8(&output.stdout)?.trim();
    println!("bid = {bid}");

    // Sometimes the repo sync fails, particularly if there there updates
    // to large repos like kernel prebuilts. An immediate re-try usually succeeds.
    if smartsync(monorepo_path, bid).is_err() {
        smartsync(monorepo_path, bid)?;
    }

    // Even though we sync the rest of the repository to a green build,
    // we sync the monorepo to tip-of-tree, which reduces merge conflicts
    // and duplicate update CLs.
    Command::new("repo").args(["sync", "."]).current_dir(monorepo_path).run_and_stream_output()?;

    Ok(())
}

fn get_suggestions(monorepo_path: &Path) -> Result<Vec<UpdateSuggestion>> {
    // TODO: Improve update suggestion algorithm.
    let mut suggestions = Vec::new();
    for compatibility in ["ignore", "loose", "strict"] {
        let output = Command::new(monorepo_path.join("crate_tool"))
            .args([
                "suggest-updates",
                "--json",
                "--patches",
                "--semver-compatibility",
                compatibility,
            ])
            .current_dir(monorepo_path)
            .output()?
            .success_or_error()?;
        let json: UpdateSuggestions = serde_json::from_slice(&output.stdout)?;
        suggestions.extend(json.updates);
    }

    // Return suggestions in random order. This reduces merge conflicts and ensures
    // all crates eventually get tried, even if something goes wrong and the program
    // terminates prematurely.
    let mut rng = thread_rng();
    suggestions.shuffle(&mut rng);

    Ok(suggestions)
}

fn try_update(
    android_root: &Path,
    monorepo_path: &Path,
    crate_name: &str,
    version: &str,
) -> Result<()> {
    println!("Trying to update {crate_name} to {version}");

    Command::new(monorepo_path.join("crate_tool"))
        .args(["update", crate_name, version])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;

    if Command::new("git")
        .args(["diff", "--exit-code", "pseudo_crate/Cargo.toml"])
        .current_dir(monorepo_path)
        .run_and_stream_output()
        .is_ok()
    {
        bail!("Crate {crate_name} was already updated");
    }

    Command::new("/usr/bin/bash")
        .args([
            "-c",
            format!(
                "source {}/build/envsetup.sh && lunch aosp_cf_x86_64_phone-trunk_staging-eng && mm",
                android_root.display()
            )
            .as_str(),
        ])
        .env_remove("OUT_DIR")
        .current_dir(monorepo_path)
        .run_and_stream_output()?;

    let now = chrono::Utc::now();
    Command::new("repo")
        .args([
            "start",
            format!(
                "automatic-crate-update-{}-{}-{}-{}-{}",
                crate_name,
                version,
                now.year(),
                now.month(),
                now.day()
            )
            .as_str(),
        ])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;
    Command::new("git").args(["add", "."]).current_dir(monorepo_path).run_and_stream_output()?;
    Command::new("git")
        .args([
            "commit",
            "-m",
            format!("Update {crate_name} to {version}\n\nTest: m rust").as_str(),
        ])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;
    Command::new("repo")
        .args([
            "upload",
            "-c",
            ".",
            "-y",
            "-o",
            "banned-words~skip",
            "-o",
            "nokeycheck",
            "--label",
            "Presubmit-Ready+1",
            "--label",
            "Autosubmit+1",
            "-o",
            "t=automatic-crate-updates",
        ])
        .current_dir(monorepo_path)
        .run_and_stream_output()?;

    Ok(())
}

fn send_email(body: String, rotation: bool) -> Result<()> {
    println!("Sending email");
    let username = env::var("USER").or_else(|_err| env::var("LOGNAME"))?;
    let mut args: Vec<&str> = vec!["--subject", "Automated crate updates"];
    if rotation {
        args.extend([
            "--to",
            "crate-update-oncall@rotations.google.com",
            "--cc",
            "crate-update-rotation@google.com",
        ])
    } else {
        args.extend(["--to", username.as_str()])
    };
    args.extend(["--inline_body", &body]);
    Command::new("/google/bin/releases/gws-sre/files/sendgmr/sendgmr")
        .args(args)
        .run_and_stream_output()?;

    Ok(())
}

fn get_stalled_crates() -> Result<HashMap<String, String>> {
    let mut stalled_crates = HashMap::new();
    let bugged_path = Path::new("/usr/bin/bugged");
    if !bugged_path.exists() {
        panic!(
            "/usr/bin/bugged is not installed. Please go install bugged. Instructions at go/bugged"
        );
    }
    let output =
        Command::new(bugged_path).args(["search", "hotlistid:7610355", "status:Open"]).output()?;
    let output_string = String::from_utf8(output.stdout).unwrap();

    let re = Regex::new(r"(\d+).+Stalled crate update:\s*([\w_-]+)\s*:\s*(\w+)").unwrap();
    for (i, line) in re.captures_iter(&output_string).enumerate() {
        let name = line.get(2).expect("NAME_NOT_FOUND").as_str();
        let ticket = line.get(1).expect("TICKET_NOT_FOUND").as_str();
        let reason = line.get(3).expect("REASON_NOT_FOUND").as_str();
        stalled_crates.insert(name.to_string(), ticket.to_string());
        println!("{:?}. {:?} ticket:{:?} reason:{:?}", i, name, ticket, reason);
    }
    return Ok(stalled_crates);
}

fn main() -> Result<()> {
    let args = Cli::parse();
    if !args.android_root.is_absolute() {
        bail!("Must be an absolute path: {}", args.android_root.display());
    }
    if !args.android_root.is_dir() {
        bail!("Does not exist, or is not a directory: {}", args.android_root.display());
    }

    let stalled_crates = get_stalled_crates().unwrap_or_default();

    let monorepo_path = args.android_root.join("external/rust/android-crates-io");

    cleanup_and_sync_monorepo(&monorepo_path)?;

    sync_to_green(&monorepo_path)?;

    Command::new("/usr/bin/bash")
        .args([
            "-c",
            "source build/envsetup.sh && lunch aosp_cf_x86_64_phone-trunk_staging-eng && m cargo_embargo",
        ])
        .env_remove("OUT_DIR")
        .current_dir(&args.android_root)
        .run_and_stream_output()?;

    let mut updates_tried = UpdatesTried::read(&monorepo_path)?;
    let mut updates_tried_string = Vec::new();

    let mut targeted_crates = BTreeMap::new();
    let mut cl_crates = BTreeMap::new();
    for suggestion in get_suggestions(&monorepo_path)? {
        let crate_name = suggestion.name.as_str();
        let version = suggestion.version.as_str();

        if !args.rotation && updates_tried.contains(crate_name, version) {
            println!("Skipping {crate_name} (already attempted recently)");
            continue;
        }
        if args.rotation && stalled_crates.contains_key(crate_name) {
            println!("Skipping {crate_name} (stalled crate)");
            continue;
        }
        cleanup_and_sync_monorepo(&monorepo_path)?;
        let res = try_update(&args.android_root, &monorepo_path, crate_name, version)
            .inspect_err(|e| println!("Update failed: {}", e));
        updates_tried_string.push(format!(
            "{} {}{}",
            suggestion.name,
            suggestion.version,
            if res.is_ok() {
                cl_crates.insert(suggestion.name.clone(), suggestion.version.clone());
                " passed local testing. A CL was generated"
            } else {
                targeted_crates.insert(suggestion.name.clone(), suggestion.version.clone());
                "failed."
            }
        ));
        updates_tried.record(suggestion.name, suggestion.version, res.is_ok())?;
    }
    cleanup_and_sync_monorepo(&monorepo_path)?;

    let mut cl_crates_string = Vec::new();
    for (i, (name, version)) in cl_crates.iter().enumerate() {
        cl_crates_string.push(format!("{}. {} {}", i + 1, name, version));
        targeted_crates.remove(name);
    }

    let mut targeted_crates_string = Vec::new();
    for (i, (name, version)) in targeted_crates.iter().enumerate() {
        targeted_crates_string.push(format!("{}. {} {}", i + 1, name, version));
    }
    let line = "-----------------------------------------------------\n";
    let body = "Please try manually updating the following crates\n".to_owned()
        + line
        + targeted_crates_string.join("\n").as_str()
        + "\n\n\n\nA CL was generated for the following crates\n"
        + line
        + cl_crates_string.join("\n").as_str()
        + "\n\n\n\nNumber of stalled crates: \n"
        + line
        + stalled_crates.len().to_string().as_str()
        + "\n"
        + "\n\n\n\nHere are the results of running the automatic updater\n"
        + line
        + updates_tried_string.join("\n").as_str();

    send_email(body, args.rotation)?;

    Ok(())
}