use eframe::egui::{self, pos2, Color32, Rounding};
use plate_tool_lib::plate::PlateFormat;
use plate_tool_lib::transfer_region::Region;
use plate_tool_lib::uuid::Uuid;
use plate_tool_lib::Well;

use crate::transfer_menu::CurrentTransferState;

use std::sync::LazyLock;

const PALETTE: plate_tool_lib::util::ColorPalette = plate_tool_lib::util::Palettes::RAINBOW;

// Stroke types
static STROKE_DEFAULT: LazyLock<egui::Stroke> =
    LazyLock::new(|| egui::Stroke::new(2.0, egui::Color32::from_gray(128)));
static STROKE_CURRENT: LazyLock<egui::Stroke> =
    LazyLock::new(|| egui::Stroke::new(2.0, egui::Color32::from_gray(200)));
static STROKE_SELECT: LazyLock<egui::Stroke> =
    LazyLock::new(|| egui::Stroke::new(3.0, egui::Color32::from_gray(200)));
static STROKE_HOVER: LazyLock<egui::Stroke> =
    LazyLock::new(|| egui::Stroke::new(3.0, egui::Color32::from_gray(255)));

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct PlateUiState {
    pub drag_start_position: Option<egui::Pos2>,
}

impl Default for PlateUiState {
    fn default() -> Self {
        Self {
            drag_start_position: None,
        }
    }
}

#[derive(Clone, Copy, Debug)]
struct WellInfo {
    volume: f32,
    color: [f64; 3],
}

impl WellInfo {
    fn new(volume: f32, color: [f64; 3]) -> Self {
        WellInfo { volume, color }
    }
}

fn calculate_grid_params(
    ul_origin: (f32, f32),
    tw: f32,
    th: f32,
    rows: u8,
    columns: u8,
) -> (f32, f32, f32) {
    // (Start X, Start Y, Radius)
    const PADDING: f32 = 20.0;

    let usable_width = tw - 2.0 * PADDING;
    let usable_height = th - 2.0 * PADDING;

    let shared_width = f32::ceil(usable_width / columns as f32);
    let shared_height = f32::ceil(usable_height / rows as f32);
    let radius = f32::min(shared_width, shared_height) / 2.0;
    (ul_origin.0 + PADDING, ul_origin.1 + PADDING, radius)
}

fn get_hover_well(
    response: &egui::Response,
    start_x: f32,
    start_y: f32,
    radius: f32,
    rows: u8,
    columns: u8,
) -> Option<(u8, u8)> {
    get_well_from_pos(
        response.hover_pos(),
        start_x,
        start_y,
        radius,
        rows,
        columns,
        true,
    )
}

fn get_well_from_pos(
    position: Option<egui::Pos2>,
    start_x: f32,
    start_y: f32,
    radius: f32,
    rows: u8,
    columns: u8,
    saturate: bool,
) -> Option<(u8, u8)> {
    // cool fact: the bounds of our frame aren't actually the bounds of what counts
    // as hovering. as a result, we have to make sure that we check these bounds here.
    // yippee
    let max_width = start_x + radius * 2.0 * columns as f32;
    let max_height = start_y + radius * 2.0 * rows as f32;

    // Some((row, column))
    position
        .map(|p| Into::<(f32, f32)>::into(p))
        .and_then(|(x, y)| {
            // Check bounds if no saturation
            if !saturate && (x < start_x || y < start_y || x > max_width || y > max_height) {
                return None;
            }

            let solved_column: u8 = match x {
                x if (x < start_x) => 1,
                x if (x > max_width) => columns,
                _ => (x - start_x).div_euclid(radius * 2.0) as u8 + 1u8,
            };
            let solved_row: u8 = match y {
                y if (y < start_y) => 1,
                y if (y > max_height) => rows,
                _ => (y - start_y).div_euclid(radius * 2.0) as u8 + 1u8,
            };

            Some((solved_row, solved_column))
        })
}

fn calculate_shading_for_wells(
    rows: u8,
    columns: u8,
    transfers: &Vec<plate_tool_lib::transfer::Transfer>,
    plate_type: plate_tool_lib::plate::PlateType,
    ordered_ids: &[Uuid],
    cache: &plate_tool_lib::transfer_region_cache::TransferRegionCache,
) -> Box<[Option<WellInfo>]> {
    let box_size: usize = rows as usize * columns as usize;
    let mut well_infos: Box<[Option<WellInfo>]> = vec![None; box_size].into_boxed_slice();

    for transfer in transfers {
        let cache_result = match plate_type {
            plate_tool_lib::plate::PlateType::Source => cache.get_or_calculate_source(transfer),
            plate_tool_lib::plate::PlateType::Destination => {
                cache.get_or_calculate_destination(transfer)
            }
        };
        if let Some(wells) = cache_result {
            for well in wells.iter().filter(|x| x.row <= rows && x.col <= columns) {
                if let Some(mut x) =
                    well_infos[(well.row - 1) as usize * columns as usize + (well.col - 1) as usize]
                {
                    x.volume += 5.0;
                    x.color = PALETTE.get_ordered(transfer.id, ordered_ids);
                } else {
                    well_infos
                        [(well.row - 1) as usize * columns as usize + (well.col - 1) as usize] =
                        Some(WellInfo::new(
                            5.0,
                            PALETTE.get_ordered(transfer.id, ordered_ids),
                        ));
                }
            }
        }
    }

    well_infos
}

fn f64_to_color32(x: [f64; 3]) -> Color32 {
    let r = x[0] as u8;
    let g = x[1] as u8;
    let b = x[2] as u8;
    Color32::from_rgb(r, g, b)
}

fn add_plate_sub(
    size: egui::Vec2,
    rows: u8,
    columns: u8,
    transfers: &Vec<plate_tool_lib::transfer::Transfer>,
    plate_type: plate_tool_lib::plate::PlateType,
    ordered_ids: &[Uuid],
    cache: &plate_tool_lib::transfer_region_cache::TransferRegionCache,
    current_transfer_state: Option<&CurrentTransferState>,
    ui: &mut egui::Ui,
    state: &mut PlateUiState,
) {
    let (response, painter) = ui.allocate_painter(size, egui::Sense::click_and_drag());

    let rect = response.rect;

    let ul_origin = rect.left_top();
    let total_width = rect.width();
    let total_height = rect.height();

    let (start_x, start_y, radius) =
        calculate_grid_params(ul_origin.into(), total_width, total_height, rows, columns);

    // Manage clicks and drags
    if response.drag_started() {
        state.drag_start_position = Some(response.hover_pos().unwrap());
    }

    let drag_start_well = get_well_from_pos(
        state.drag_start_position,
        start_x,
        start_y,
        radius,
        rows,
        columns,
        false,
    );
    let hovered_well = get_hover_well(&response, start_x, start_y, radius, rows, columns);

    if response.clicked() {
        if let Some(cts) = current_transfer_state {
            let end_well: Option<Well> = hovered_well.map(|(row, col)| Well { row, col });
            if let Some(end_well) = end_well {
                let new_region = Region::new_from_wells(end_well, None);
                let mut cts = cts.lock().unwrap();
                match plate_type {
                    plate_tool_lib::plate::PlateType::Source => cts.source_region = new_region,
                    plate_tool_lib::plate::PlateType::Destination => {
                        cts.destination_region = new_region
                    }
                }
            }
        }
    }
    if response.drag_stopped() {
        if let Some(cts) = current_transfer_state {
            if let Some(start_well) = drag_start_well.map(|(row, col)| Well { row, col }) {
                let end_well: Option<Well> = hovered_well.map(|(row, col)| Well { row, col });
                let new_region = Region::new_from_wells(start_well, end_well);
                let mut cts = cts.lock().unwrap();
                match plate_type {
                    plate_tool_lib::plate::PlateType::Source => cts.source_region = new_region,
                    plate_tool_lib::plate::PlateType::Destination => {
                        cts.destination_region = new_region
                    }
                }
            }
        }
        state.drag_start_position = None;
    }

    let current_selection: Option<Region> = {
        if let Some(cts) = current_transfer_state {
            let cts = cts.lock().unwrap();
            match plate_type {
                plate_tool_lib::plate::PlateType::Source => Some(cts.source_region.clone()),
                plate_tool_lib::plate::PlateType::Destination => {
                    Some(cts.destination_region.clone())
                }
            }
        } else {
            None
        }
    };
    let well_infos = {
        // Get non-active transfer info
        let mut well_infos =
            calculate_shading_for_wells(rows, columns, transfers, plate_type, ordered_ids, cache);

        // Get wells in the current transfer to tack on to well_infos separately
        let current_transfer_wells: Option<Box<[(usize, usize)]>> = {
            (match plate_type {
                plate_tool_lib::plate::PlateType::Source => current_transfer_state
                    .and_then(|x| x.lock().ok())
                    .map(|mut x| x.get_source_wells()),
                plate_tool_lib::plate::PlateType::Destination => current_transfer_state
                    .and_then(|x| x.lock().ok())
                    .map(|mut x| x.get_destination_wells()),
            })
            .map(|xs| {
                xs.iter()
                    .map(|x| (x.row as usize, x.col as usize))
                    .collect()
            })
        };
        if let Some(wells) = current_transfer_wells {
            for w in wells {
                let well_info = &mut well_infos[(w.0 - 1) * columns as usize + (w.1 - 1)];
                let volume = well_info.map(|x| x.volume).unwrap_or(0.0)
                    + current_transfer_state
                        .and_then(|x| x.lock().ok())
                        .map(|x| x.volume)
                        .unwrap_or(0.0);

                *well_info = Some(WellInfo {
                    color: [255.0, 255.0, 255.0],
                    volume: 1.0,
                })
            }
        }
        well_infos
    };

    // Plate Frame
    painter.rect_stroke(
        egui::Rect {
            min: pos2(start_x, start_y),
            max: pos2(
                start_x + 2.0 * radius * columns as f32,
                start_y + 2.0 * radius * rows as f32,
            ),
        },
        Rounding::default(),
        *STROKE_DEFAULT,
    );

    // Draw wells
    for c_row in 1..=rows {
        for c_column in 1..=columns {
            let center = egui::pos2(
                start_x + radius + 2.0 * radius * (c_column - 1) as f32,
                start_y + radius + 2.0 * radius * (c_row - 1) as f32,
            );

            //
            // Draw fill first
            //

            if let Some(well_info) =
                well_infos[(c_row - 1) as usize * columns as usize + (c_column - 1) as usize]
            {
                painter.circle_filled(center, radius * 0.80, f64_to_color32(well_info.color));
            }

            //
            // Draw stroke on top
            //

            // Base stroke
            painter.circle_stroke(center, radius * 0.80, *STROKE_DEFAULT);

            if current_selection.as_ref().is_some_and(|x| {
                x.contains_well(
                    &Well {
                        row: c_row,
                        col: c_column,
                    },
                    None,
                )
            }) {
                painter.circle_stroke(center, radius * 0.80, *STROKE_CURRENT);
            }
            if response.dragged() {
                if let (Some(hw), Some(dw)) = (hovered_well, drag_start_well) {
                    if c_column <= u8::max(hw.1, dw.1)
                        && c_column >= u8::min(hw.1, dw.1)
                        && c_row <= u8::max(hw.0, dw.0)
                        && c_row >= u8::min(hw.0, dw.0)
                    {
                        painter.circle_stroke(center, radius * 0.80, *STROKE_SELECT);
                    }
                }
            }
            if Some((c_row, c_column)) == hovered_well {
                painter.circle_stroke(center, radius * 0.80, *STROKE_HOVER);
            }
        }
    }

    // Draw row/column labels
    for c_row in 0..rows {
        painter.text(
            egui::pos2(
                start_x - 10.0,
                start_y + radius + 2.0 * radius * c_row as f32,
            ),
            egui::Align2::CENTER_CENTER,
            (c_row + 1).to_string(),
            egui::FontId::monospace(f32::min(radius, 14.0)),
            egui::Color32::from_gray(128),
        );
    }
    for c_column in 0..columns {
        painter.text(
            egui::pos2(
                start_x + radius + 2.0 * radius * c_column as f32,
                start_y - 12.0,
            ),
            egui::Align2::CENTER_CENTER,
            plate_tool_lib::util::num_to_letters(c_column + 1).unwrap(),
            egui::FontId::monospace(f32::min(radius, 14.0)),
            egui::Color32::from_gray(128),
        );
    }
}

pub fn add_plate(
    size: egui::Vec2,
    pf: PlateFormat,
    transfers: &Vec<plate_tool_lib::transfer::Transfer>,
    plate_type: plate_tool_lib::plate::PlateType,
    ordered_ids: &[Uuid],
    cache: &plate_tool_lib::transfer_region_cache::TransferRegionCache,
    current_transfer_state: Option<&CurrentTransferState>,
    ui: &mut egui::Ui,
    state: &mut PlateUiState,
) {
    add_plate_sub(
        size,
        pf.rows(),
        pf.columns(),
        transfers,
        plate_type,
        ordered_ids,
        cache,
        current_transfer_state,
        ui,
        state,
    );
}