#![no_std]
#![no_main]

mod max6675;

use core::fmt::Write;
use core::write;

use cortex_m::delay::Delay;

use display_interface::DisplayError;

use embedded_graphics::{
    mono_font::{iso_8859_1::FONT_10X20, MonoTextStyle},
    pixelcolor::BinaryColor,
    prelude::*,
    text::{Alignment, Baseline, Text, TextStyleBuilder},
};

use embedded_hal::{
    digital::v2::{InputPin, OutputPin, ToggleableOutputPin},
    timer::CountDown as _,
};

use fugit::{ExtU64, MicrosDurationU64, RateExtU32};

use heapless::String;

use max6675::{Error, FixedU16, Max6675};

use nb::Error::WouldBlock;

use panic_halt as _;

use rp_pico::{
    entry,
    hal::{
        clocks::{self, Clock},
        gpio::PinState,
        i2c::I2C,
        timer::CountDown,
        usb::UsbBus,
        Sio, Timer, Watchdog,
    },
    pac::{CorePeripherals, Peripherals},
    Pins,
};

use ssd1306::{mode::BufferedGraphicsMode, prelude::*, I2CDisplayInterface, Ssd1306};

use usb_device::{bus::UsbBus as UsbBusTrait, bus::UsbBusAllocator, prelude::*};

use usbd_serial::{SerialPort, USB_CLASS_CDC};

const HYSTERESIS: FixedU16 = FixedU16::const_from_int(5);

const PEAK_TEMPERATURE: FixedU16 = FixedU16::const_from_int(240);

const PREHEAT_TEMPERATURE: FixedU16 = FixedU16::const_from_int(150);

const SOAK_TIME: MicrosDurationU64 = MicrosDurationU64::secs(90);

enum State<'a> {
    Off,
    Preheat,
    Soak(CountDown<'a>),
    Reflow,
    Cooling,
}

fn print_serial<'a, B: UsbBusTrait>(
    serial: &mut SerialPort<'a, B>,
    msg: &str,
) -> Result<(), UsbError> {
    let buf = msg.as_bytes();
    let mut pos = 0;

    let mut retry_count = 0;
    while pos < buf.len() {
        match serial.write(&buf[pos..]) {
            Ok(count) => {
                pos += count;
                retry_count = 0;
                continue;
            }
            Err(UsbError::WouldBlock) if retry_count < 100 => {
                retry_count += 1;
                continue;
            }
            Err(err) => return Err(err),
        }
    }

    Ok(())
}

fn println_serial<'a, B: UsbBusTrait>(
    serial: &mut SerialPort<'a, B>,
    msg: &str,
) -> Result<(), UsbError> {
    print_serial(serial, msg)?;
    print_serial(serial, "\r\n")
}

fn stringify_display_error(err: DisplayError) -> &'static str {
    match err {
        DisplayError::InvalidFormatError => "invalid format",
        DisplayError::BusWriteError => "bus write error",
        DisplayError::DCError => "dc error",
        DisplayError::CSError => "cs error",
        DisplayError::DataFormatNotImplemented => "data format not implemented",
        DisplayError::RSError => "rs error",
        DisplayError::OutOfBoundsError => "out of bounds",
        _ => "unknown error",
    }
}

fn display_text<DI, SIZE>(
    display: &mut Ssd1306<DI, SIZE, BufferedGraphicsMode<SIZE>>,
    text: &str,
    baseline: Baseline,
) -> Result<(), DisplayError>
where
    DI: WriteOnlyDataCommand,
    SIZE: DisplaySize,
{
    let (w, h) = display.dimensions();

    Text::with_text_style(
        text,
        Point::new((w / 2).into(), (h / 2).into()),
        MonoTextStyle::new(&FONT_10X20, BinaryColor::On),
        TextStyleBuilder::new()
            .alignment(Alignment::Center)
            .baseline(baseline)
            .build(),
    )
    .draw(display)?;

    Ok(())
}

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

    let mut watchdog = Watchdog::new(p.WATCHDOG);
    let clocks = clocks::init_clocks_and_plls(
        rp_pico::XOSC_CRYSTAL_FREQ,
        p.XOSC,
        p.CLOCKS,
        p.PLL_SYS,
        p.PLL_USB,
        &mut p.RESETS,
        &mut watchdog,
    )
    .ok()
    .unwrap();

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

    let sio = Sio::new(p.SIO);
    let pins = Pins::new(p.IO_BANK0, p.PADS_BANK0, sio.gpio_bank0, &mut p.RESETS);

    let usb_bus = UsbBusAllocator::new(UsbBus::new(
        p.USBCTRL_REGS,
        p.USBCTRL_DPRAM,
        clocks.usb_clock,
        true,
        &mut p.RESETS,
    ));

    let mut serial = SerialPort::new(&usb_bus);

    let mut usb_dev = UsbDeviceBuilder::new(&usb_bus, UsbVidPid(0x2342, 0x1337))
        .manufacturer("lujoga")
        .product("Soldering Iron")
        .device_class(USB_CLASS_CDC)
        .build();

    for _ in 0..1000 {
        delay.delay_ms(5);

        if !usb_dev.poll(&mut [&mut serial]) {
            continue;
        }

        let mut buf = [0u8];
        if let Ok(_) = serial.read(&mut buf[..]) {
            break;
        }
    }

    let _ = println_serial(&mut serial, "Hello, world!");

    let i2c = I2C::i2c0(
        p.I2C0,
        pins.gpio0.into_mode(),
        pins.gpio1.into_mode(),
        400.kHz(),
        &mut p.RESETS,
        clocks.peripheral_clock.freq(),
    );

    let mut display = Ssd1306::new(
        I2CDisplayInterface::new_custom_address(i2c, 0x3c),
        DisplaySize128x64,
        DisplayRotation::Rotate0,
    )
    .into_buffered_graphics_mode();

    if let Err(err) = display.init() {
        let _ = print_serial(&mut serial, stringify_display_error(err));
        panic!();
    }

    let timer = Timer::new(p.TIMER, &mut p.RESETS);

    let so = pins.gpio6.into_floating_input();
    let cs = pins.gpio7.into_push_pull_output_in_state(PinState::High);
    let sck = pins.gpio8.into_push_pull_output();
    let mut thermocouple = Max6675::new(cs, sck, so);

    let button = pins.gpio15.into_pull_up_input();

    let mut relay = pins.gpio17.into_push_pull_output_in_state(PinState::Low);

    let mut led = pins.led.into_push_pull_output_in_state(PinState::Low);

    let mut button_held = false;
    let mut button_cooldown = timer.count_down();
    button_cooldown.start(0.secs());

    let mut counter = 0;
    let mut state = State::Off;
    let mut heating = false;
    let mut error_counter = 0;
    loop {
        delay.delay_ms(5);
        counter = (counter + 1) % 2000;

        if let Ok(_) = button_cooldown.wait() {
            let button_pressed = button.is_low().unwrap();
            if button_pressed && !button_held {
                state = if let State::Off = state {
                    State::Preheat
                } else {
                    State::Off
                };

                button_cooldown.start(500.millis());
            } else {
                button_cooldown.start(0.secs());
            }

            button_held = button_pressed;
        }

        if counter % 50 == 0 {
            let thermocouple_result = thermocouple.read_temperature(&mut timer.count_down());

            if let Some(temperature) = thermocouple_result.as_ref().copied().ok() {
                error_counter = 0;

                match state {
                    State::Preheat => {
                        if temperature < PREHEAT_TEMPERATURE {
                            heating = true;
                        } else {
                            let mut time = timer.count_down();
                            time.start(SOAK_TIME);
                            state = State::Soak(time);
                        }
                    }
                    State::Soak(ref mut time) => match time.wait() {
                        Ok(_) => state = State::Reflow,
                        Err(WouldBlock) => {
                            if temperature < PREHEAT_TEMPERATURE {
                                heating = true;
                            }
                            if temperature > PREHEAT_TEMPERATURE + HYSTERESIS {
                                heating = false;
                            }
                        }
                        Err(_) => unreachable!(),
                    },
                    State::Reflow => {
                        if temperature < PEAK_TEMPERATURE {
                            heating = true;
                        } else {
                            heating = false;

                            state = State::Cooling;
                        }
                    }
                    State::Cooling if temperature < PREHEAT_TEMPERATURE => {
                        state = State::Preheat;
                    }
                    _ => (),
                }
            } else {
                error_counter += 1;
            }

            if error_counter >= 4 {
                state = State::Off;
            }

            let mut text: String<12> = String::new();
            let text_str;

            match thermocouple_result {
                Ok(temperature) => {
                    if let Err(_) = write!(&mut text, "{} °C", temperature) {
                        text_str = "format error";
                    } else {
                        text_str = text.as_str();
                    }
                }
                Err(Error::IOError) => {
                    text_str = "i/o error";
                }
                Err(Error::OpenThermocouple) => {
                    text_str = "n/c";
                }
            };

            let _ = println_serial(&mut serial, text_str);

            display.clear_buffer();

            if let Err(err) = display_text(&mut display, text_str, Baseline::Bottom) {
                let _ = println_serial(&mut serial, stringify_display_error(err));
            }

            let state_str = match state {
                State::Off => "OFF",
                State::Preheat => "PREHEAT",
                State::Soak(_) => "SOAK",
                State::Reflow => "REFLOW",
                State::Cooling => "COOLING",
            };

            if let Err(err) = display_text(&mut display, state_str, Baseline::Top) {
                let _ = println_serial(&mut serial, stringify_display_error(err));
            }

            if let Err(err) = display.flush() {
                let _ = println_serial(&mut serial, stringify_display_error(err));
            }
        }

        if let State::Off = state {
            heating = false;
        }

        if heating {
            relay.set_high().unwrap();
        } else {
            relay.set_low().unwrap();
        }

        if counter % 100 == 0 {
            led.toggle().unwrap();
        }

        if !usb_dev.poll(&mut [&mut serial]) {
            continue;
        }

        let mut buf = [0u8; 64];
        if let Ok(count) = serial.read(&mut buf[..]) {
            let _ = serial.write(&buf[..count]);
        }
    }
}