/*
 * Copyright The Monogon Project Authors.
 * SPDX-License-Identifier: Apache-2.0
 */

#![no_main]
#![no_std]

extern crate alloc;

use alloc::vec::Vec;
use core::result::Result;
use core::fmt;
use prost::Message;
use uefi::fs::FileSystem;
use uefi::proto::device_path::build::media::FilePath;
use uefi::proto::device_path::build::DevicePathBuilder;
use uefi::proto::device_path::{DeviceSubType, DeviceType, LoadedImageDevicePath};
use uefi::table::boot;
use uefi::{prelude::*, CStr16};
use uefi_services::println;

use abloader_proto::metropolis::node::abloader::spec::*;

const A_LOADER_PATH: &CStr16 = cstr16!("\\EFI\\metropolis\\boot-a.efi");
const B_LOADER_PATH: &CStr16 = cstr16!("\\EFI\\metropolis\\boot-b.efi");

const LOADER_STATE_PATH: &CStr16 = cstr16!("\\EFI\\metropolis\\loader_state.pb");

enum ValidSlot {
    A,
    B,
}

impl ValidSlot {
    // other returns B if the value is A and A if the value is B.
    fn other(&self) -> Self {
        match self {
            ValidSlot::A => ValidSlot::B,
            ValidSlot::B => ValidSlot::A,
        }
    }
    // path returns the path to the slot's EFI payload.
    fn path(&self) -> &'static CStr16 {
        match self {
            ValidSlot::A => A_LOADER_PATH,
            ValidSlot::B => B_LOADER_PATH,
        }
    }
}

enum ReadLoaderStateError {
    FSReadError(uefi::fs::Error),
    DecodeError(prost::DecodeError),
}

impl fmt::Display for ReadLoaderStateError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
           ReadLoaderStateError::FSReadError(e) => write!(f, "while reading state file: {}", e),
           ReadLoaderStateError::DecodeError(e) => write!(f, "while decoding state file contents: {}", e),
        }
    }
}

fn read_loader_state(fs: &mut FileSystem) -> Result<AbLoaderData, ReadLoaderStateError> {
    let state_raw = fs.read(&LOADER_STATE_PATH).map_err(|e| ReadLoaderStateError::FSReadError(e))?;
    AbLoaderData::decode(state_raw.as_slice()).map_err(|e| ReadLoaderStateError::DecodeError(e))
}

fn load_slot_image(slot: &ValidSlot, boot_services: &BootServices) -> uefi::Result<Handle> {
    let mut storage = Vec::new();

    // Build the path to the slot payload. This takes the path to the loader
    // itself, strips off the file path and following element(s) and appends
    // the path to the correct slot payload.
    let new_image_path = {
        let loaded_image_device_path = boot_services
            .open_protocol_exclusive::<LoadedImageDevicePath>(boot_services.image_handle())?;

        let mut builder = DevicePathBuilder::with_vec(&mut storage);

        for node in loaded_image_device_path.node_iter() {
            if node.full_type() == (DeviceType::MEDIA, DeviceSubType::MEDIA_FILE_PATH) {
                break;
            }

            builder = builder.push(&node).unwrap();
        }

        builder = builder
            .push(&FilePath {
                path_name: slot.path(),
            })
            .unwrap();

        builder.finalize().unwrap()
    };

    boot_services
        .load_image(
            boot_services.image_handle(),
            boot::LoadImageSource::FromDevicePath {
                device_path: new_image_path,
                from_boot_manager: false,
            },
        )
}

#[entry]
fn main(_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
    uefi_services::init(&mut system_table).unwrap();

    let boot_services = system_table.boot_services();

    let boot_slot_raw = {
        let mut esp_fs = boot_services
            .get_image_file_system(boot_services.image_handle())
            .expect("image filesystem not available");

        let mut loader_data = match read_loader_state(&mut esp_fs) {
            Ok(d) => d,
            Err(e) => {
                println!("Unable to load A/B loader state, using default slot A: {}", e);
                AbLoaderData {
                    active_slot: Slot::A.into(),
                    next_slot: Slot::None.into(),
                }
            }
        };

        // If next_slot is set, use it as slot to boot but clear it in the
        // state file as the next boot should not use it again. If it should
        // be permanently activated, it is the OS's job to put it into
        if loader_data.next_slot != Slot::None.into() {
            let next_slot = loader_data.next_slot;
            loader_data.next_slot = Slot::None.into();
            let new_loader_data = loader_data.encode_to_vec();
            esp_fs
                .write(&LOADER_STATE_PATH, new_loader_data)
                .expect("failed to write back abdata");
            next_slot
        } else {
            loader_data.active_slot
        }
    };

    let boot_slot = match Slot::try_from(boot_slot_raw) {
        Ok(Slot::A) => ValidSlot::A,
        Ok(Slot::B) => ValidSlot::B,
        _ => {
            println!("Invalid slot ({}) active, falling back to A", boot_slot_raw);
            ValidSlot::A
        }
    };

    let payload_image = match load_slot_image(&boot_slot, boot_services) {
        Ok(img) => img,
        Err(e) => {
            println!("Error loading intended slot, falling back to other slot: {}", e);
            match load_slot_image(&boot_slot.other(), boot_services) {
                Ok(img) => img,
                Err(e) => {
                    panic!("Loading from both slots failed, second slot error: {}", e);
                },
            }
        }
    };

    // Boot the payload.
    boot_services
        .start_image(payload_image)
        .expect("failed to start payload");
    Status::SUCCESS
}