use std::io::Read;

use azure_core::credentials::TokenCredential;
use azure_identity::{AzureCliCredential, AzureCliCredentialOptions};
use clap::{Parser, crate_name};
use log::{info, trace};
use serde::Deserialize;
use serde_json::json;
use url::Url;

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Cli {
    /// The operation this program was run with.
    operation: Operation,

    /// Logging level. Set to one of
    /// `error`, `warn`, `info`, `debug`, or `trace`.
    #[clap(long, env = "ACR_LOGIN_LOG_LEVEL", default_value = "warn")]
    log_level: log::Level,

    // Azure Tenant ID. Must be the tenant where the ACR is
    // located.
    #[clap(long, env = "AZURE_TENANT_ID")]
    azure_tenant_id: String,
}

#[derive(Debug, Clone, clap::ValueEnum)]
enum Operation {
    Store,
    Get,
    Erase,
    List,
}

// The standard ACR domain for Azure
const ACR_DOMAIN: &str = ".azurecr.io";

// The standard ACR username
const ACR_USERNAME: &str = "00000000-0000-0000-0000-000000000000";

// Result with boxed Result
type StdResult<T> = Result<T, Box<dyn std::error::Error>>;

// ACR refresh token structure
#[derive(Deserialize)]
struct AcrRefreshToken {
    refresh_token: String,
}

#[tokio::main]
async fn main() -> StdResult<()> {
    let cli = Cli::parse();

    // Set up the stderr logger
    stderrlog::new()
        .module(module_path!())
        .verbosity(cli.log_level)
        .init()?;

    match cli.operation {
        // For gets, get any credentials for the given registry.
        Operation::Get => {
            let creds = get_docker_credential(&cli).await?;

            // Write the output to stdout
            serde_json::to_writer(std::io::stdout(), &creds)?;
        }

        // For other operations, do nothing.
        _ => info!("{} is a read-only provider", crate_name!()),
    }

    Ok(())
}

async fn get_docker_credential(cli: &Cli) -> StdResult<serde_json::Value> {
    // Expecting the registry of the ACR as input from stdin.
    let mut registry = String::new();
    std::io::stdin().read_to_string(&mut registry)?;

    // Trim the registry of whitespace.
    let registry = registry.trim();

    // If the url doesn't end with ".azurecr.io" then
    // we don't know how to handle this.
    if !registry.ends_with(ACR_DOMAIN) {
        // Normal operation when used as a credStore.
        info!("{} not handling registry: {}", crate_name!(), registry);

        // Return an empty JSON dictionary
        return Ok(json!({}));
    }

    // Need to connect to the repository's OAuth endpoint to exchange the token we just got.
    let url = Url::parse(&format!("https://{registry}/oauth2/exchange"))?;

    // Obtain an Azure access token scoped to management.azure.com for this
    // specific tenant. Using AzureCliCredential with tenant_id ensures
    // `az account get-access-token --tenant` is called, so the active
    // subscription does not need to match the ACR's tenant.
    let tenant = &cli.azure_tenant_id;
    let credential = AzureCliCredential::new(Some(AzureCliCredentialOptions {
        tenant_id: Some(tenant.clone()),
        ..Default::default()
    }))?;
    let token_response = credential
        .get_token(&["https://management.azure.com/.default"], None)
        .await?;

    // Set up the parameters for the post to the OAuth endpoint as per
    // https://github.com/Azure/acr/blob/main/docs/AAD-OAuth.md#calling-post-oauth2exchange-to-get-an-acr-refresh-token
    let grant_type = String::from("access_token");
    let params = [
        ("grant_type", grant_type.as_str()),
        ("service", registry),
        ("tenant", tenant.as_str()),
        ("access_token", token_response.token.secret()),
    ];
    trace!("Params: {params:?}");

    let form_body = url::form_urlencoded::Serializer::new(String::new())
        .extend_pairs(params)
        .finish();

    // Send the request to the endpoint in order to get the ACR refresh token.
    let client = reqwest::Client::new();
    let response = client
        .post(url)
        .header(
            reqwest::header::CONTENT_TYPE,
            "application/x-www-form-urlencoded",
        )
        .body(form_body)
        .send()
        .await?
        .error_for_status()?
        .json::<AcrRefreshToken>()
        .await?;

    // At this point we have a refresh token, which can be provided to
    // Docker.
    let creds = json!({
        "Username": ACR_USERNAME,
        "Secret": response.refresh_token
    });
    trace!("Credentials: {creds:?}");

    // If we get here we succeeded!
    Ok(creds)
}
