#![no_std]
#![no_main]
#![feature(alloc_error_handler)]

extern crate alloc;

use panic_itm as _;

use alloc::boxed::Box;
use alloc_cortex_m::CortexMHeap;
use core::alloc::Layout;
use core::cell::RefCell;
use cortex_m::asm::bootload;
use cortex_m_rt::entry;
use embedded_hal::digital::v2::InputPin;
use stm32f1xx_hal_bxcan::delay::Delay;
use stm32f1xx_hal_bxcan::flash::{FlashSize, SectorSize};
use stm32f1xx_hal_bxcan::pac::{CorePeripherals, Peripherals};
use stm32f1xx_hal_bxcan::prelude::*;

use ross_eeprom::DeviceInfo;
use ross_logger::{log_debug, log_error, log_info, log_warning, LogLevel, Logger};
use ross_protocol::convert_packet::ConvertPacket;
use ross_protocol::event::bootloader::BootloaderHelloEvent;
use ross_protocol::event::general::{AckEvent, DataEvent};
use ross_protocol::event::programmer::{
    ProgrammerHelloEvent, ProgrammerStartConfigUpgradeEvent, ProgrammerStartFirmwareUpgradeEvent,
};

use crate::config::*;
use crate::config_upgrader::{ConfigUpgrader, ConfigUpgraderError};
use crate::firmware_upgrader::{FirmwareUpgrader, FirmwareUpgraderError};
use crate::helper::can_helper::setup_can_protocol;
use crate::helper::eeprom_helper::setup_eeprom;
use crate::helper::heap_helper::allocate_heap;
use crate::helper::type_helper::CanProtocol;

mod config;
mod config_upgrader;
mod firmware_upgrader;
mod helper;

const PROGRAM_ADDRESS: u32 = 0x0801_0000;

#[global_allocator]
static ALLOCATOR: CortexMHeap = CortexMHeap::empty();

#[entry]
fn main() -> ! {
    let dp = Peripherals::take().unwrap();
    let cp = CorePeripherals::take().unwrap();

    let mut flash = dp.FLASH.constrain();
    let mut rcc = dp.RCC.constrain();

    let clocks = rcc
        .cfgr
        .use_hse(8.mhz())
        .sysclk(72.mhz())
        .hclk(72.mhz())
        .pclk1(36.mhz())
        .pclk2(72.mhz())
        .freeze(&mut flash.acr);

    let config = Config {
        heap_config: Some(HeapConfig { size: 4096 }),
        eeprom_config: Some(EepromConfig {
            bitrate: 400_000,
            device_info_address: 0,
        }),
        can_config: Some(CanConfig {
            bitrate: 50_000,
            tseg1: 13,
            tseg2: 2,
            sjw: 1,
        }),
    };

    allocate_heap(&config.unwrap_heap_config(), &ALLOCATOR);

    let logger = RefCell::new(Logger::new(LogLevel::Info, cp.ITM));

    log_info!(
        logger,
        "Bootloader v{} initialized.",
        env!("CARGO_PKG_VERSION")
    );

    let mut gpioa = dp.GPIOA.split(&mut rcc.apb2);

    let upgrade_input = gpioa.pa1.into_pull_down_input(&mut gpioa.crl);

    // If no firmware upgrade is requested, proceed with bootloading the program
    if upgrade_input.is_low().unwrap() {
        log_info!(logger, "Booting firmware.");
        boot();
    }

    log_info!(logger, "Entering upgrade mode.");

    let mut gpiob = dp.GPIOB.split(&mut rcc.apb2);
    let mut afio = dp.AFIO.constrain(&mut rcc.apb2);

    let eeprom = RefCell::new(setup_eeprom(
        config.unwrap_eeprom_config(),
        clocks,
        gpiob.pb6.into_alternate_open_drain(&mut gpiob.crl),
        gpiob.pb7.into_alternate_open_drain(&mut gpiob.crl),
        dp.I2C1,
        &mut rcc.apb1,
        &mut afio.mapr,
    ));

    let device_info = eeprom.borrow_mut().read_device_info().unwrap();

    log_info!(
        logger,
        "Loaded device information from EEPROM ({:?}).",
        device_info
    );

    let protocol = RefCell::new(setup_can_protocol(
        config.unwrap_can_config(),
        &device_info,
        clocks,
        gpioa.pa11.into_floating_input(&mut gpioa.crh),
        gpioa.pa12.into_alternate_push_pull(&mut gpioa.crh),
        dp.CAN1,
        dp.USB,
        &mut rcc.apb1,
        &mut afio.mapr,
    ));

    let mut flash_writer = flash.writer(SectorSize::Sz1K, FlashSize::Sz128K);
    // TODO: change to true
    flash_writer.change_verification(false);

    let firmware_upgrader = RefCell::new(FirmwareUpgrader::new(flash_writer));

    let mut delay = Delay::new(cp.SYST, clocks);

    let config_upgrader = RefCell::new(ConfigUpgrader::new(&eeprom));

    protocol
        .borrow_mut()
        .add_packet_handler(
            Box::new(|packet, protocol| {
                if let Ok(event) = ProgrammerHelloEvent::try_from_packet(&packet) {
                    log_debug!(logger, "Received `programmer_hello_event` ({:?}).", event);
                    transmit_bootloader_hello_event(protocol, &device_info, &event, &logger);
                }
            }),
            false,
        )
        .unwrap();

    protocol
        .borrow_mut()
        .add_packet_handler(
            Box::new(|packet, protocol| {
                if let Ok(event) = ProgrammerStartFirmwareUpgradeEvent::try_from_packet(&packet) {
                    log_debug!(
                        logger,
                        "Received `programmer_start_firmware_upgrade_event` ({:?}).",
                        event
                    );

                    if let Err(err) = firmware_upgrader.borrow_mut().start_upgrade(&event) {
                        log_error!(
                            logger,
                            "Failed to start firmware upgrade with error ({:?}).",
                            err
                        );
                    } else {
                        log_info!(logger, "Started firmware upgrade.");

                        transmit_ack_event(
                            protocol,
                            &device_info,
                            event.programmer_address,
                            &logger,
                        );
                    }
                }
            }),
            false,
        )
        .unwrap();

    protocol
        .borrow_mut()
        .add_packet_handler(
            Box::new(|packet, protocol| {
                if let Ok(event) = ProgrammerStartConfigUpgradeEvent::try_from_packet(&packet) {
                    log_debug!(
                        logger,
                        "Received `programmer_start_config_upgrade_event` ({:?}).",
                        event
                    );

                    if let Err(err) = config_upgrader.borrow_mut().start_upgrade(&event) {
                        log_error!(
                            logger,
                            "Failed to start config upgrade with error ({:?}).",
                            err
                        );
                    } else {
                        log_info!(logger, "Started config upgrade.");

                        transmit_ack_event(
                            protocol,
                            &device_info,
                            event.programmer_address,
                            &logger,
                        );
                    }
                }
            }),
            false,
        )
        .unwrap();

    protocol
        .borrow_mut()
        .add_packet_handler(
            Box::new(|packet, protocol| {
                if let Ok(event) = DataEvent::try_from_packet(&packet) {
                    match firmware_upgrader.borrow_mut().handle_data_event(&event) {
                        Err(FirmwareUpgraderError::UpgradeNotStarted) => {
                            match config_upgrader.borrow_mut().handle_data_event(&event, &mut delay) {
                                Err(ConfigUpgraderError::UpgradeNotStarted) => {
                                    log_warning!(
                                        logger,
                                        "Unexpected `data_event`.",
                                    );
                                },
                                Err(err) => {
                                    log_error!(
                                        logger,
                                        "Config upgrader failed to handle data packet with error ({:?}).",
                                        err
                                    );
                                },
                                Ok(finished_upgrading) => {
                                    transmit_ack_event(
                                        protocol,
                                        &device_info,
                                        event.transmitter_address,
                                        &logger,
                                    );

                                    if finished_upgrading {
                                        log_debug!(logger, "Config upgrade completed. Booting.");
                                        boot();
                                    }
                                }
                            }
                        }
                        Err(err) => {
                            log_error!(
                                logger,
                                "Firmware upgrader failed to handle data packet with error ({:?}).",
                                err
                            );
                        }
                        Ok(finished_upgrading) => {
                            transmit_ack_event(
                                protocol,
                                &device_info,
                                event.transmitter_address,
                                &logger,
                            );

                            if finished_upgrading {
                                log_debug!(logger, "Firmware upgrade completed. Booting.");
                                boot();
                            }
                        }
                    }
                }
            }),
            false,
        )
        .unwrap();

    loop {
        if let Err(err) = protocol.borrow_mut().tick() {
            log_warning!(logger, "Unexpected error occurred ({:?}).", err);
        }
    }
}

fn boot() -> ! {
    unsafe {
        bootload(PROGRAM_ADDRESS as *const u32);
    }
}

fn transmit_bootloader_hello_event(
    protocol: &mut CanProtocol,
    device_info: &DeviceInfo,
    programmer_hello_event: &ProgrammerHelloEvent,
    logger: &RefCell<Logger>,
) {
    let event = BootloaderHelloEvent {
        bootloader_address: device_info.device_address,
        programmer_address: programmer_hello_event.programmer_address,
    };

    if let Err(err) = protocol.send_packet(&event.to_packet()) {
        log_error!(
            logger,
            "Failed to send `bootloader_hello_event` ({:?}).",
            err
        );
    } else {
        log_debug!(logger, "Sent `bootloader_hello_event` ({:?}).", event);
    }
}

fn transmit_ack_event(
    protocol: &mut CanProtocol,
    device_info: &DeviceInfo,
    receiver_address: u16,
    logger: &RefCell<Logger>,
) {
    let event = AckEvent {
        receiver_address,
        transmitter_address: device_info.device_address,
    };

    if let Err(err) = protocol.send_packet(&event.to_packet()) {
        log_error!(logger, "Failed to send `ack_event` ({:?}).", err);
    } else {
        log_debug!(logger, "Sent `ack_event` ({:?}).", event);
    }
}

#[alloc_error_handler]
fn oom(_: Layout) -> ! {
    loop {}
}