#![no_std]
#![no_main]

use cyw43_pio::PioSpi;
use defmt::*;
use embassy_executor::Spawner;
use embassy_net::{Stack, StackResources};
use embassy_rp::bind_interrupts;
use embassy_rp::gpio::{Level, Output};
use embassy_rp::peripherals::{DMA_CH0, PIO0};
use embassy_rp::pio::{InterruptHandler, Pio};
use embassy_time::{Duration, Timer};
use static_cell::StaticCell;
use {defmt_rtt as _, panic_probe as _};

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => InterruptHandler<PIO0>;
});

// Task to run the CYW43 network stack
#[embassy_executor::task]
async fn wifi_task(
    runner: cyw43::Runner<'static, Output<'static>, PioSpi<'static, PIO0, 0, DMA_CH0>>
) -> ! {
    runner.run().await
}

// Task to run the network stack
#[embassy_executor::task]
async fn net_task(stack: &'static Stack<cyw43::NetDriver<'static>>) -> ! {
    stack.run().await
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    let p = embassy_rp::init(Default::default());
    
    // Initialize PIO for CYW43 communication
    let Pio { mut common, sm0, .. } = Pio::new(p.PIO0, Irqs);
    
    // Set up CYW43 pins
    let pwr = Output::new(p.PIN_23, Level::Low);
    let cs = Output::new(p.PIN_25, Level::High);
    let mut pio_spi = PioSpi::new(
        &mut common,
        sm0,
        p.DMA_CH0,
        p.PIN_24,
        p.PIN_29,
        p.PIN_25,
    );
    
    // Load CYW43 firmware
    static STATE: StaticCell<cyw43::State> = StaticCell::new();
    let state = STATE.init(cyw43::State::new());
    let fw = include_bytes!("../../../firmware/43439A0.bin");
    let clm = include_bytes!("../../../firmware/43439A0_clm.bin");
    
    // Initialize CYW43 driver
    let (net_device, mut control, runner) = cyw43::new(state, pwr, pio_spi, fw).await;
    unwrap!(spawner.spawn(wifi_task(runner)));
    
    // Initialize network stack
    static STACK: StaticCell<Stack<cyw43::NetDriver<'static>>> = StaticCell::new();
    static RESOURCES: StaticCell<StackResources<3>> = StaticCell::new();
    
    let config = embassy_net::Config::dhcpv4(Default::default());
    let seed = 0x1234_5678; // Random seed for network stack
    
    let stack = &*STACK.init(Stack::new(
        net_device,
        config,
        RESOURCES.init(StackResources::new()),
        seed,
    ));
    
    unwrap!(spawner.spawn(net_task(stack)));
    
    // Connect to WiFi
    control.init(clm).await;
    control.set_power_management(cyw43::PowerManagementMode::PowerSave).await;
    
    let wifi_network = "YOUR_WIFI_SSID";
    let wifi_password = "YOUR_WIFI_PASSWORD";
    
    info!("Connecting to WiFi...");
    loop {
        match control.join_wpa2(wifi_network, wifi_password).await {
            Ok(_) => {
                info!("WiFi connected!");
                break;
            }
            Err(err) => {
                warn!("Failed to connect to WiFi: {:?}", err);
                Timer::after(Duration::from_secs(1)).await;
            }
        }
    }
    
    // Wait for DHCP
    info!("Waiting for DHCP...");
    while !stack.is_config_up() {
        Timer::after_millis(100).await;
    }
    info!("DHCP configured!");
    
    // Print IP address
    if let Some(config) = stack.config_v4() {
        info!("IP address: {:?}", config.address);
    }
    
    // Blink LED to show we're connected
    loop {
        control.gpio_set(0, true).await;
        Timer::after_millis(100).await;
        control.gpio_set(0, false).await;
        Timer::after_millis(900).await;
    }
}

use embassy_net::tcp::TcpSocket;
use embassy_net::{Stack, StackResources};

#[embassy_executor::task]
async fn http_server_task(stack: &'static Stack<cyw43::NetDriver<'static>>) {
    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];
    
    loop {
        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
        socket.set_timeout(Some(Duration::from_secs(10)));
        
        info!("Listening on TCP:80...");
        if let Err(e) = socket.accept(80).await {
            warn!("Accept error: {:?}", e);
            continue;
        }
        
        info!("Received connection from {:?}", socket.remote_endpoint());
        
        let mut buf = [0; 1024];
        loop {
            let n = match socket.read(&mut buf).await {
                Ok(0) => {
                    warn!("Connection closed");
                    break;
                }
                Ok(n) => n,
                Err(e) => {
                    warn!("Read error: {:?}", e);
                    break;
                }
            };
            
            // Simple HTTP response
            let response = b"HTTP/1.1 200 OK\r\n\
                Content-Type: text/html\r\n\
                Connection: close\r\n\
                \r\n\
                <html><body><h1>Hello from Pico 2 W!</h1></body></html>";
            
            if let Err(e) = socket.write_all(response).await {
                warn!("Write error: {:?}", e);
            }
            break;
        }
    }
}