// src/bin/main.rs

#![no_std]
#![no_main]
#![deny(
    clippy::mem_forget,
    reason = "mem::forget is generally not safe to do with esp_hal types"
)]
// TODO WARNING core 1 should be logic, core 0 wifi, its flipped now

use embassy_executor::Spawner;
use embassy_futures::select::{select, Either, select3, Either3};
use embassy_net::{Runner, StackResources};
use embassy_sync::signal::Signal;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_time::{Duration, Timer};
use projekt_final::bus::I2cInner;
use projekt_final::mqtt::client::mqtt_set_imu;

use esp_alloc as _;
use esp_backtrace as _;

use esp_hal::{
    gpio::InputConfig,
    clock::CpuClock,
    gpio::{Input, Pull},
    i2c::master::{Config as I2cConfig, I2c},
    rng::Rng,
    system::{CpuControl, Stack},
    timer::timg::TimerGroup,
};
use esp_wifi::{
    wifi::{ClientConfiguration, Configuration, WifiController, WifiDevice, WifiEvent, WifiState},
    EspWifiController,
};

use pages_tui::input::Key;
use log::info;
use rust_mqtt::packet::v5::publish_packet::QualityOfService;
use static_cell::StaticCell;
use core::cell::RefCell;

use projekt_final::{
    bus,
    display,
    mpu,
    mqtt::client::{mqtt_events, mqtt_try_publish, mqtt_publish, mqtt_subscribe, mqtt_task, IncomingMsg},
};

extern crate alloc;
use alloc::format;

static I2C_BUS: StaticCell<RefCell<I2cInner>> = StaticCell::new();
static APP_CORE_STACK: StaticCell<Stack<8192>> = StaticCell::new();
static EXECUTOR_CORE1: StaticCell<esp_hal_embassy::Executor> = StaticCell::new();
static NETWORK_READY: Signal<CriticalSectionRawMutex, ()> = Signal::new();

macro_rules! mk_static {
    ($t:ty,$val:expr) => {{
        static STATIC_CELL: static_cell::StaticCell<$t> = static_cell::StaticCell::new();
        #[deny(unused_attributes)]
        let x = STATIC_CELL.uninit().write(($val));
        x
    }};
}

const SSID: &str = env!("SSID");
const PASSWORD: &str = env!("PASSWORD");
const MQTT_PUBLISH_DIVIDER: u32 = 10;

esp_bootloader_esp_idf::esp_app_desc!();

#[esp_hal_embassy::main]
async fn main(spawner: Spawner) -> ! {
    esp_println::logger::init_logger_from_env();
    info!("===============================");
    info!("  ESP32 IoT Firmware Starting");
    info!("===============================");

    let config = esp_hal::Config::default().with_cpu_clock(CpuClock::max());
    let peripherals = esp_hal::init(config);
    esp_alloc::heap_allocator!(size: 72 * 1024);

    info!("Initializing I2C bus...");
    let i2c = I2c::new(peripherals.I2C0, I2cConfig::default())
        .expect("Failed to create I2C instance")
        .with_sda(peripherals.GPIO21)
        .with_scl(peripherals.GPIO22)
        .into_async();

    let i2c_bus = I2C_BUS.init(RefCell::new(i2c));
    let display_i2c = bus::new_device(i2c_bus);
    let mpu_i2c = bus::new_device(i2c_bus);

    info!("Initializing WiFi...");
    let timg0 = TimerGroup::new(peripherals.TIMG0);
    let mut rng = Rng::new(peripherals.RNG);

    let esp_wifi_ctrl = mk_static!(
        EspWifiController<'static>,
        esp_wifi::init(timg0.timer0, rng.clone()).unwrap()
    );

    let (controller, interfaces) =
        esp_wifi::wifi::new(esp_wifi_ctrl, peripherals.WIFI).unwrap();

    let wifi_interface = interfaces.sta;
    let timg1 = TimerGroup::new(peripherals.TIMG1);
    esp_hal_embassy::init(timg1.timer0);

    let seed = (rng.random() as u64) << 32 | rng.random() as u64;

    // Start core 1 for WiFi and MQTT (network stack created there)
    let mut cpu_control = CpuControl::new(peripherals.CPU_CTRL);
    let _guard = cpu_control.start_app_core(
        APP_CORE_STACK.init(Stack::new()),
        move || {
            let executor = EXECUTOR_CORE1.init(esp_hal_embassy::Executor::new());
            executor.run(|spawner| {
                spawner.spawn(core1_network_task(spawner, controller, wifi_interface, seed)).ok();
            });
        }
    ).unwrap();

    // Wait for network to be ready (signaled from core 1)
    NETWORK_READY.wait().await;
    info!("Network ready, starting core 0 tasks");

    let config = InputConfig::default().with_pull(Pull::Down);
    let button_select = Input::new(peripherals.GPIO32, config);
    let button_next = Input::new(peripherals.GPIO35, config);
    spawner.spawn(button_detection_task(button_select, button_next)).unwrap();

    // Core 0: display and MPU tasks
    spawner.spawn(display::task::display_task(display_i2c)).expect("spawn display_task");
    spawner.spawn(mpu::task::mpu_task(mpu_i2c)).expect("spawn mpu_task");

    display::api::set_status("Booting...").await;
    mqtt_subscribe("esp32/imu").await;
    mqtt_publish("esp32/imu", b"online", QualityOfService::QoS1, false).await;

    display::api::set_status("Running").await;
    display::api::set_mqtt_status(true, 0).await;

    let mqtt_rx = mqtt_events();
    let imu_rx = mpu::api::events();
    let mut imu_reading_count: u32 = 0;
    let mut mqtt_msg_count: u32 = 0;
    let mut mqtt_publish_drops: u32 = 0;

    loop {
        match select3(
            mqtt_rx.receive(),
            imu_rx.receive(),
            Timer::after(Duration::from_secs(5)),
        ).await {
            Either3::First(msg) => {
                mqtt_msg_count += 1;
                handle_mqtt_message(msg).await;
                display::api::set_mqtt_status(true, mqtt_msg_count).await;
            }
            Either3::Second(mut reading) => {
                // Drain any queued IMU messages and keep only the latest
                let mut drained = 0;
                while let Ok(next) = imu_rx.try_receive() {
                    reading = next;
                    drained += 1;
                }
                if drained > 0 {
                    log::info!("IMU drained {} stale readings before display", drained);
                }
                imu_reading_count += 1;
                
                // Show_imu is now non-blocking (uses try_send internally)
                display::api::show_imu(reading);
                
                if imu_reading_count % MQTT_PUBLISH_DIVIDER == 0 {
                    let payload = format!(
                        "{{\"ax\":{:.2},\"ay\":{:.2},\"az\":{:.2},\"t\":{:.1}}}",
                        reading.accel_g[0], reading.accel_g[1], reading.accel_g[2], reading.temp_c
                    );
                    mqtt_set_imu(payload.as_bytes());
                }
            }
            Either3::Third(_) => {
                crate::mpu::api::IMU_CHANNEL.clear();
                info!("IMU heartbeat: force-cleared queue, {} readings total, {} mqtt drops", 
                      imu_reading_count, mqtt_publish_drops);
            }
        }
    }
}

// Runs on core 1 - creates and owns the network stack
#[embassy_executor::task]
async fn core1_network_task(
    spawner: Spawner,
    controller: WifiController<'static>,
    wifi_interface: WifiDevice<'static>,
    seed: u64,
) {
    spawner.spawn(connection_task(controller)).ok();

    let net_config = embassy_net::Config::dhcpv4(Default::default());
    let (stack, runner) = embassy_net::new(
        wifi_interface,
        net_config,
        mk_static!(StackResources<3>, StackResources::<3>::new()),
        seed,
    );

    spawner.spawn(net_task(runner)).ok();

    // Wait for network
    loop {
        if stack.is_link_up() { break; }
        Timer::after(Duration::from_millis(500)).await;
    }
    loop {
        if let Some(config) = stack.config_v4() {
            info!("Got IP: {}", config.address);
            break;
        }
        Timer::after(Duration::from_millis(500)).await;
    }

    // Signal core 0 that network is ready
    NETWORK_READY.signal(());

    spawner.spawn(mqtt_task(stack)).ok();
}

async fn handle_mqtt_message(msg: IncomingMsg) {
    if let Ok(txt) = core::str::from_utf8(&msg.payload) {
        match txt {
            "clear" => { display::api::clear().await; }
            "status" => { mqtt_publish("esp32/status", b"running", QualityOfService::QoS1, false).await; }
            _ => { display::api::add_chat_message(txt).await; }
        }
    }
}

#[embassy_executor::task]
async fn connection_task(mut controller: WifiController<'static>) {
    loop {
        if esp_wifi::wifi::wifi_state() == WifiState::StaConnected {
            controller.wait_for_event(WifiEvent::StaDisconnected).await;
            Timer::after(Duration::from_millis(5000)).await;
        }
        if !matches!(controller.is_started(), Ok(true)) {
            let client_config = Configuration::Client(ClientConfiguration {
                ssid: SSID.into(),
                password: PASSWORD.into(),
                ..Default::default()
            });
            controller.set_configuration(&client_config).unwrap();
            info!("Wi-Fi starting...");
            controller.start_async().await.unwrap();
        }
        match controller.connect_async().await {
            Ok(_) => info!("Wifi connected!"),
            Err(e) => {
                info!("Failed to connect to wifi: {e:#?}");
                Timer::after(Duration::from_millis(5000)).await
            }
        }
    }
}

#[embassy_executor::task]
async fn net_task(mut runner: Runner<'static, WifiDevice<'static>>) {
    runner.run().await
}

#[embassy_executor::task]
async fn button_detection_task(mut select_btn: Input<'static>, mut next_btn: Input<'static>) {
    loop {
        match select(
            select_btn.wait_for_rising_edge(),
            next_btn.wait_for_rising_edge(),
        ).await {
            Either::First(_) => {
                info!("Detection: GPIO 32 (Select) triggered!");
                display::api::push_key(Key::enter()).await;
            }
            Either::Second(_) => {
                info!("Detection: GPIO 35 (Next) triggered!");
                display::api::push_key(Key::tab()).await
            }
        }
        // Debounce: prevent mechanical bouncing from double-triggering
        embassy_time::Timer::after(embassy_time::Duration::from_millis(200)).await;
    }
}