use crate::args::Generate;
use crate::{args::shadow::PKG_VERSION, config::toml::ConfigToml};
use anyhow::Context;
use concepts::{ComponentType, ExecutionId};
use schemars::schema_for;
use std::{
    borrow::Cow,
    fs::File,
    io::{BufWriter, Write as _, stdout},
    path::PathBuf,
};
use utils::{wasm_tools::WasmComponent, wit};

impl Generate {
    pub(crate) async fn run(self) -> Result<(), anyhow::Error> {
        match self {
            Generate::ConfigSchema { output } => generate_toml_schema(output),
            Generate::Extensions {
                component_type,
                input_wit_directory,
                output_directory,
                force,
            } => {
                generate_exported_extension_wits(
                    input_wit_directory,
                    output_directory,
                    component_type,
                    force,
                )
                .await
            }
            Generate::WitSupport {
                component_type,
                output_directory,
            } => generate_support_wits(component_type, output_directory).await,
            Generate::ExecutionId => {
                println!("{}", ExecutionId::generate());
                Ok(())
            }
        }
    }
}

pub(crate) fn generate_toml_schema(output: Option<PathBuf>) -> Result<(), anyhow::Error> {
    let schema = schema_for!(ConfigToml);
    if let Some(output) = output {
        // Save to a file
        let mut writer = BufWriter::new(File::create(&output)?);
        serde_json::to_writer_pretty(&mut writer, &schema)?;
        writer.write_all(b"\n")?;
        writer.flush()?; // Do not swallow errors
    } else {
        serde_json::to_writer_pretty(stdout().lock(), &schema)?;
    }
    Ok(())
}

const HEADER: &str = "// Generated by Obelisk";

pub(crate) async fn generate_exported_extension_wits(
    input_wit_directory: PathBuf,
    output_directory: PathBuf,
    component_type: ComponentType,
    force: bool,
) -> Result<(), anyhow::Error> {
    let wasm_component = WasmComponent::new_from_wit_folder(&input_wit_directory, component_type)?;
    let pkgs_to_wits = wasm_component.exported_extension_wits()?;
    for (pkg_fqn, new_content) in pkgs_to_wits {
        let pkg_folder = output_directory.join(pkg_fqn.to_string().replace(':', "_"));
        let wit_file = pkg_folder.join(format!("{}.wit", pkg_fqn.package_name));

        let old_content = tokio::fs::read_to_string(&wit_file)
            .await
            .unwrap_or_default();

        let old_content = if force {
            None
        } else {
            Some(strip_header(&old_content))
        };
        if old_content.as_ref() != Some(&new_content) {
            let new_content = format!("{HEADER} {PKG_VERSION}\n{new_content}");
            tokio::fs::create_dir_all(&pkg_folder)
                .await
                .with_context(|| format!("cannot write {pkg_folder:?}"))?;
            tokio::fs::write(&wit_file, new_content.as_bytes())
                .await
                .with_context(|| format!("cannot write {wit_file:?}"))?;
            println!("{wit_file:?} created or updated");
        } else {
            println!("{wit_file:?} is up to date");
        }
    }
    Ok(())
}

fn strip_header(old_content: &str) -> String {
    let old_content = match old_content.strip_prefix(HEADER) {
        Some(wit) => {
            if let Some((_, wit)) = wit.split_once('\n') {
                Cow::Borrowed(wit)
            } else {
                Cow::Borrowed(wit)
            }
        }
        None => Cow::Borrowed(old_content),
    };
    let old_content = match old_content.strip_prefix(&format!("/{HEADER}")) {
        // Bug in wasm_tools is turning // into ///
        Some(wit) => {
            if let Some((_, wit)) = wit.split_once('\n') {
                Cow::Borrowed(wit)
            } else {
                Cow::Borrowed(wit)
            }
        }
        None => old_content,
    };
    old_content.into_owned()
}

pub(crate) async fn generate_support_wits(
    component_type: ComponentType,
    output_directory: PathBuf,
) -> Result<(), anyhow::Error> {
    let files = match component_type {
        ComponentType::ActivityWasm => {
            vec![
                wit::WIT_OBELISK_ACTIVITY_PACKAGE_PROCESS,
                wit::WIT_OBELISK_LOG_PACKAGE,
            ]
        }
        ComponentType::ActivityStub | ComponentType::ActivityExternal => vec![],
        ComponentType::Workflow => vec![
            wit::WIT_OBELISK_TYPES_PACKAGE,
            wit::WIT_OBELISK_WORKFLOW_PACKAGE,
            wit::WIT_OBELISK_LOG_PACKAGE,
        ],
        ComponentType::WebhookEndpoint => {
            vec![
                wit::WIT_OBELISK_TYPES_PACKAGE, // Needed for -schedule ext functions.
                wit::WIT_OBELISK_LOG_PACKAGE,
            ]
        }
    };
    for [folder, filename, content] in files {
        let output_directory = output_directory.join(folder);
        let path = output_directory.join(filename);
        if let Ok(actual) = tokio::fs::read_to_string(&path).await
            && actual == content
        {
            println!("{path:?} is up to date");
        } else {
            tokio::fs::create_dir_all(&output_directory)
                .await
                .with_context(|| format!("cannot write {output_directory:?}"))?;
            tokio::fs::write(&path, content)
                .await
                .with_context(|| format!("cannot write {path:?}"))?;
            println!("{path:?} created or updated");
        }
    }
    Ok(())
}
