use std::fs;
use std::io::{Read, Write};
use std::path::Path;
use std::fs::File;
use image::GenericImageView;
use csv::Reader;

fn main() {
    linker_be_nice();
    // make sure linkall.x is the last linker script (otherwise might cause problems with flip-link)
    println!("cargo:rustc-link-arg=-Tlinkall.x");

    #[cfg(feature="simulation")]
    write_sim_data();

    compile_assets();

}

fn compile_assets() {
    let asset_path = Path::new("assets");

    let mut image_output = File::create(Path::new("target/images.rs")).unwrap();
    for image in fs::read_dir(asset_path).unwrap() {
        let fname = image.unwrap().file_name();
        let fname_str = fname.to_str().unwrap();
        if fname_str.ends_with(".pbm") {
            let img = image::open(asset_path.join(fname_str)).unwrap();
            let img_name = fname_str.rsplit_once('.').unwrap().0.to_uppercase().replace("-", "_");

            let mut converted_row = Vec::new();
            let mut byte_buf = String::new();
            image_output.write_all(format!("pub const {img_name}: ImageRaw<BinaryColor> = ImageRaw::new(&[\n").as_bytes()).unwrap();
            for (x, _, pixel) in img.pixels() {
                if pixel.0 == [0, 0, 0, 255] {
                    byte_buf.push('1');
                } else {
                    byte_buf.push('0');
                }
                if byte_buf.len() == 8 {
                    converted_row.push(byte_buf);
                    byte_buf = String::new();
                }
                if x == img.width() - 1 {
                    if !byte_buf.is_empty() {
                        byte_buf.push('_');
                        for _ in 0..(9 - byte_buf.len()) {
                            byte_buf.push('0');
                        }
                        converted_row.push(byte_buf);
                        byte_buf = String::new();
                    }
                    image_output.write_all(b"    ").unwrap();
                    for pix in converted_row.iter() {
                        image_output.write_all(format!("0b{pix}, ").as_bytes()).unwrap();
                    }
                    image_output.write_all(b"\n").unwrap();
                    converted_row = Vec::new();
                }
            }

            image_output.write_all(format!("], {});\n", img.width()).as_bytes()).unwrap();
            println!("cargo::rerun-if-changed={fname_str}");
        }
    }
}

fn write_sim_data() {
    let test_data_path = Path::new("test-data");
    let output_path = Path::new("target");
    let gps_input = test_data_path.join("LocationGps.csv");
    let gyro_input = test_data_path.join("GyroscopeUncalibrated.csv");
    let accel_input = test_data_path.join("AccelerometerUncalibrated.csv");

    let gps_output = output_path.join("gps_test_data.msgpack");
    let motion_output = output_path.join("motion_test_data.msgpack");
    let unified_output = output_path.join("unified.msgpack");

    println!("cargo::rerun-if-changed={}", gps_input.to_str().unwrap());
    if !gps_output.exists() || gps_output.metadata().unwrap().modified().unwrap() < gps_input.metadata().unwrap().modified().unwrap() {
        let mut gps_data = Reader::from_reader(File::open(gps_input.clone()).unwrap());
        let record_count = gps_data.records().count() as u32;
        let mut gps_data = Reader::from_reader(File::open(gps_input).unwrap());

        let headers = gps_data.headers().unwrap();
        let (timestamp_idx, lat_idx, lon_idx) = (
            headers.iter().position(|x| { x == "seconds_elapsed" }).unwrap(),
            headers.iter().position(|x| { x == "longitude" }).unwrap(),
            headers.iter().position(|x| { x == "latitude" }).unwrap(),
        );
        let mut gps_output = File::create(gps_output.clone()).unwrap();
        rmp::encode::write_array_len(&mut gps_output, record_count).unwrap();
        let mut last_stamp = 0.0;
        for record in gps_data.records().flatten() {
            let (timestamp, lat, lon) = (
                record.get(timestamp_idx).unwrap().parse().unwrap(),
                record.get(lat_idx).unwrap().parse().unwrap(),
                record.get(lon_idx).unwrap().parse().unwrap()
            );
            let next_delay = timestamp - last_stamp;
            last_stamp = timestamp;
            rmp::encode::write_array_len(&mut gps_output, 3).unwrap();
            rmp::encode::write_f64(&mut gps_output, next_delay).unwrap();
            rmp::encode::write_f64(&mut gps_output, lat).unwrap();
            rmp::encode::write_f64(&mut gps_output, lon).unwrap();
        }
    }

    println!("cargo::rerun-if-changed={}", accel_input.to_str().unwrap());
    println!("cargo::rerun-if-changed={}", gyro_input.to_str().unwrap());
    let rebuild_motion = {
        if motion_output.exists() {
            let motion_stamp = motion_output.metadata().unwrap().modified().unwrap();
            motion_stamp < accel_input.metadata().unwrap().modified().unwrap() || motion_stamp < gyro_input.metadata().unwrap().modified().unwrap() 
        } else {
            true
        }
    };
    if rebuild_motion {
        let mut accel_data = Reader::from_reader(File::open(accel_input.clone()).unwrap());
        let mut gyro_data = Reader::from_reader(File::open(gyro_input).unwrap());

        let record_count = accel_data.records().count() as u32;
        let mut accel_data = Reader::from_reader(File::open(accel_input).unwrap());

        let headers = accel_data.headers().unwrap();
        let (timestamp_idx, accel_x_idx, accel_y_idx, accel_z_idx) = (
            headers.iter().position(|x| { x == "seconds_elapsed" }).unwrap(),
            headers.iter().position(|x| { x == "x" }).unwrap(),
            headers.iter().position(|x| { x == "y" }).unwrap(),
            headers.iter().position(|x| { x == "z" }).unwrap(),
        );

        let headers = gyro_data.headers().unwrap();
        let (gyro_x_idx, gyro_y_idx, gyro_z_idx) = (
            headers.iter().position(|x| { x == "x" }).unwrap(),
            headers.iter().position(|x| { x == "y" }).unwrap(),
            headers.iter().position(|x| { x == "z" }).unwrap(),
        );

        let mut motion_output = File::create(motion_output.clone()).unwrap();
        rmp::encode::write_array_len(&mut motion_output, record_count).unwrap();
        let mut last_stamp = 0.0;
        for (accel_record, gyro_record) in accel_data.records().flatten().zip(gyro_data.records().flatten()) {
            let (timestamp, accel_x, accel_y, accel_z) = (
                accel_record.get(timestamp_idx).unwrap().parse().unwrap(),
                accel_record.get(accel_x_idx).unwrap().parse().unwrap(),
                accel_record.get(accel_y_idx).unwrap().parse().unwrap(),
                accel_record.get(accel_z_idx).unwrap().parse().unwrap()
            );
            let (gyro_x, gyro_y, gyro_z) = (
                gyro_record.get(gyro_x_idx).unwrap().parse().unwrap(),
                gyro_record.get(gyro_y_idx).unwrap().parse().unwrap(),
                gyro_record.get(gyro_z_idx).unwrap().parse().unwrap()
            );
            let next_delay = timestamp - last_stamp;
            if next_delay >= 0.02 {
                last_stamp = timestamp;
                rmp::encode::write_array_len(&mut motion_output, 7).unwrap();
                rmp::encode::write_f64(&mut motion_output, next_delay).unwrap();
                rmp::encode::write_f64(&mut motion_output, accel_x).unwrap();
                rmp::encode::write_f64(&mut motion_output, accel_y).unwrap();
                rmp::encode::write_f64(&mut motion_output, accel_z).unwrap();
                rmp::encode::write_f64(&mut motion_output, gyro_x).unwrap();
                rmp::encode::write_f64(&mut motion_output, gyro_y).unwrap();
                rmp::encode::write_f64(&mut motion_output, gyro_z).unwrap();
            }
        }
    }

    // GPS data = 2, motion data = 1
    
    let mut unified_fd = File::create(unified_output.clone()).unwrap();

    // Write out the stream index, which will be 2 (motion + gps)
    rmp::encode::write_array_len(&mut unified_fd, 2).unwrap();

    let mut motion_output = File::open(motion_output).unwrap();
    let mut gps_output = File::open(gps_output).unwrap();
    rmp::encode::write_ext_meta(&mut unified_fd, motion_output.metadata().unwrap().len() as u32, 1).unwrap();
    let mut buf = Vec::new();
    motion_output.read_to_end(&mut buf).unwrap();
    unified_fd.write_all(buf.as_slice()).unwrap();

    rmp::encode::write_ext_meta(&mut unified_fd, gps_output.metadata().unwrap().len() as u32, 2).unwrap();
    let mut buf = Vec::new();
    gps_output.read_to_end(&mut buf).unwrap();
    unified_fd.write_all(buf.as_slice()).unwrap();

    let mut partitions = Reader::from_reader(File::open("partitions.csv").unwrap());
    let mut data_offset = 0x9000; // Assumes default bootloader size (0x7000) plus partition table (0x2000)
    let mut data_size = 0;
    let mut found_sim_partition = false;
    for p_desc in partitions.records().flatten() {
        let offset = parse_partition_number(p_desc.get(3).unwrap().trim());
        data_size = parse_partition_number(p_desc.get(4).unwrap().trim()).unwrap();
        data_offset = offset.unwrap_or(data_offset);
        if let Some("sim") = p_desc.get(0) {
            found_sim_partition = true;
            break;
        } else {
            data_offset += data_size;
        }
    }

    if !found_sim_partition {
        panic!("Could not find a 'sim' partition in partitions.csv!");
    }

    if buf.len() >= data_size {
        panic!("Simulation data is too big! Cannot fit {:#x} bytes into a partition with a size of {data_size:#x} bytes.", buf.len());
    }

    let mut data_flash_script = File::create(output_path.join("flash-sim-data.sh")).unwrap();
    data_flash_script.write_all(format!("espflash write-bin {data_offset:#x?} {}", unified_output.to_str().unwrap()).as_bytes()).unwrap();
}

fn parse_partition_number(n: &str) -> Option<usize> {
    if n.is_empty() {
        None
    } else if let Some(hex_offset) = n.strip_prefix("0x") {
        Some(usize::from_str_radix(hex_offset, 16).unwrap())
    } else if let Some(mb_offset) = n.strip_suffix("M") {
        Some(mb_offset.parse::<usize>().unwrap() * 1024 * 1024)
    } else {
        Some(n.parse().unwrap())
    }
}

fn linker_be_nice() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() > 1 {
        let kind = &args[1];
        let what = &args[2];

        match kind.as_str() {
            "undefined-symbol" => match what.as_str() {
                "_defmt_timestamp" => {
                    eprintln!();
                    eprintln!("💡 `defmt` not found - make sure `defmt.x` is added as a linker script and you have included `use defmt_rtt as _;`");
                    eprintln!();
                }
                "_stack_start" => {
                    eprintln!();
                    eprintln!("💡 Is the linker script `linkall.x` missing?");
                    eprintln!();
                }
                "esp_wifi_preempt_enable"
                | "esp_wifi_preempt_yield_task"
                | "esp_wifi_preempt_task_create" => {
                    eprintln!();
                    eprintln!("💡 `esp-wifi` has no scheduler enabled. Make sure you have the `builtin-scheduler` feature enabled, or that you provide an external scheduler.");
                    eprintln!();
                }
                "embedded_test_linker_file_not_added_to_rustflags" => {
                    eprintln!();
                    eprintln!("💡 `embedded-test` not found - make sure `embedded-test.x` is added as a linker script for tests");
                    eprintln!();
                }
                _ => (),
            },
            // we don't have anything helpful for "missing-lib" yet
            _ => {
                std::process::exit(1);
            }
        }

        std::process::exit(0);
    }

    println!(
        "cargo:rustc-link-arg=-Wl,--error-handling-script={}",
        std::env::current_exe().unwrap().display()
    );
}