#![doc = include_str!("../README.md")]
#![feature(uefi_std)]

use crate::context::{RootContext, SproutContext};
use crate::entries::BootableEntry;
use crate::options::SproutOptions;
use crate::options::parser::OptionsRepresentable;
use crate::phases::phase;
use anyhow::{Context, Result};
use log::info;
use std::collections::BTreeMap;
use std::ops::Deref;
use uefi::proto::device_path::LoadedImageDevicePath;
use uefi::proto::device_path::text::{AllowShortcuts, DisplayOnly};

/// actions: Code that can be configured and executed by Sprout.
pub mod actions;

/// config: Sprout configuration mechanism.
pub mod config;

/// context: Stored values that can be cheaply forked and cloned.
pub mod context;

/// drivers: EFI drivers to load and provide extra functionality.
pub mod drivers;

/// entries: Boot menu entries that have a title and can execute actions.
pub mod entries;

/// extractors: Runtime code that can extract values into the Sprout context.
pub mod extractors;

/// generators: Runtime code that can generate entries with specific values.
pub mod generators;

/// phases: Hooks into specific parts of the boot process.
pub mod phases;

/// setup: Code that initializes the UEFI environment for Sprout.
pub mod setup;

/// options: Parse the options of the Sprout executable.
pub mod options;

/// utils: Utility functions that are used by other parts of Sprout.
pub mod utils;

/// The main entrypoint of sprout.
/// It is possible this function will not return if actions that are executed
/// exit boot services or do not return control to sprout.
fn main() -> Result<()> {
    // Initialize the basic UEFI environment.
    setup::init()?;

    // Parse the options to the sprout executable.
    let options = SproutOptions::parse().context("unable to parse options")?;

    // Load the configuration of sprout.
    // At this point, the configuration has been validated and the specified
    // version is checked to ensure compatibility.
    let config = config::loader::load(&options)?;

    // Load the root context.
    // This is done in a block to ensure the release of the LoadedImageDevicePath protocol.
    let mut root = {
        let current_image_device_path_protocol = uefi::boot::open_protocol_exclusive::<
            LoadedImageDevicePath,
        >(uefi::boot::image_handle())
        .context("unable to get loaded image device path")?;
        let loaded_image_path = current_image_device_path_protocol.deref().to_boxed();
        info!(
            "loaded image path: {}",
            loaded_image_path.to_string(DisplayOnly(false), AllowShortcuts(false))?
        );
        RootContext::new(loaded_image_path, options)
    };

    // Insert the configuration actions into the root context.
    root.actions_mut().extend(config.actions.clone());

    // Create a new sprout context with the root context.
    let mut context = SproutContext::new(root);

    // Insert the configuration values into the sprout context.
    context.insert(&config.values);

    // Freeze the sprout context so it can be shared and cheaply cloned.
    let context = context.freeze();

    // Execute the early phase.
    phase(context.clone(), &config.phases.early).context("unable to execute early phase")?;

    // Load all configured drivers.
    drivers::load(context.clone(), &config.drivers).context("unable to load drivers")?;

    // Run all the extractors declared in the configuration.
    let mut extracted = BTreeMap::new();
    for (name, extractor) in &config.extractors {
        let value = extractors::extract(context.clone(), extractor)
            .context(format!("unable to extract value {}", name))?;
        info!("extracted value {}: {}", name, value);
        extracted.insert(name.clone(), value);
    }
    let mut context = context.fork();
    // Insert the extracted values into the sprout context.
    context.insert(&extracted);
    let context = context.freeze();

    // Execute the startup phase.
    phase(context.clone(), &config.phases.startup).context("unable to execute startup phase")?;

    let mut entries = Vec::new();

    // Insert all the static entries from the configuration into the entry list.
    for (name, entry) in config.entries {
        // Associate the main context with the static entry.
        entries.push(BootableEntry::new(
            name,
            entry.title.clone(),
            context.clone(),
            entry,
        ));
    }

    // Run all the generators declared in the configuration.
    for (name, generator) in config.generators {
        let context = context.fork().freeze();

        // We will prefix all entries with [name]-.
        let prefix = format!("{}-", name);

        // Add all the entries generated by the generator to the entry list.
        // The generator specifies the context associated with the entry.
        for mut entry in generators::generate(context.clone(), &generator)? {
            entry.prepend_name_prefix(&prefix);
            entries.push(entry);
        }
    }

    for entry in &mut entries {
        let mut context = entry.context().fork();
        // Insert the values from the entry configuration into the
        // sprout context to use with the entry itself.
        context.insert(&entry.declaration().values);
        let context = context
            .finalize()
            .context("unable to finalize context")?
            .freeze();
        // Provide the new context to the bootable entry.
        entry.swap_context(context);
        // Restamp the title with any values.
        entry.restamp_title();
    }

    // TODO(azenla): Implement boot menu here.
    // For now, we just print all of the entries.
    info!("entries:");
    for (index, entry) in entries.iter().enumerate() {
        let title = entry.context().stamp(&entry.declaration().title);
        info!("  entry {} [{}]: {}", index, entry.name(), title);
    }

    // Execute the late phase.
    phase(context.clone(), &config.phases.late).context("unable to execute late phase")?;

    // If --boot is specified, or defaults.entry is specified, use that to find the entry to boot.
    let boot = context
        .root()
        .options()
        .boot
        .as_ref()
        .or(config.defaults.entry.as_ref());

    // Use the boot option if possible, otherwise pick the first entry.
    let entry = if let Some(ref boot) = boot {
        entries
            .iter()
            .enumerate()
            .find(|(index, entry)| {
                entry.name() == boot.as_str()
                    || entry.title() == boot.as_str()
                    || index.to_string() == boot.as_str()
            })
            .context(format!("unable to find entry: {boot}"))?
            .1 // select the bootable entry.
    } else {
        entries.first().context("no entries found")?
    };

    // Execute all the actions for the selected entry.
    for action in &entry.declaration().actions {
        let action = entry.context().stamp(action);
        actions::execute(entry.context().clone(), &action)
            .context(format!("unable to execute action '{}'", action))?;
    }

    // Sprout doesn't necessarily guarantee anything was booted.
    // If we reach here, we will exit back to whoever called us.
    Ok(())
}