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)]
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,
) -> Option<(u8, u8)> {
    get_well_from_pos(response.hover_pos(), start_x, start_y, radius)
}

fn get_well_from_pos(
    position: Option<egui::Pos2>,
    start_x: f32,
    start_y: f32,
    radius: f32,
) -> Option<(u8, u8)> {
    // Some((row, column))
    position
        .map(|p| Into::<(f32, f32)>::into(p))
        .and_then(|(x, y)| {
            // Check bounds
            if x < start_x || y < start_y {
                return None;
            }
            // CHECK Bottom Right BOUND

            let solved_column: u8 = (x - start_x).div_euclid(radius * 2.0) as u8;
            let solved_row: u8 = (y - start_y).div_euclid(radius * 2.0) as u8;

            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_source(transfer),
            plate_tool_lib::plate::PlateType::Destination => cache.get_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 as usize * columns as usize + well.col as usize]
                {
                    x.volume += 5.0;
                    x.color = PALETTE.get_ordered(transfer.id, ordered_ids);
                } else {
                    well_infos[well.row as usize * columns as usize + well.col 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_grid_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);
    let hovered_well = get_hover_well(&response, start_x, start_y, radius);

    if response.clicked() {
        if let Some(cts) = current_transfer_state {
            let end_well: Option<Well> = hovered_well.map(|(row, col)| Well {
                row: row + 1,
                col: col + 1,
            });
            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 {
            let start_well: Well = drag_start_well
                .map(|(row, col)| Well {
                    // Lib uses 1-indexing!
                    row: row + 1,
                    col: col + 1,
                })
                .unwrap();
            let end_well: Option<Well> = hovered_well.map(|(row, col)| Well {
                row: row + 1,
                col: col + 1,
            });
            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 = {
        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()),
            })
                // Drop back to 0-indexing here
            .map(|xs| xs.iter().map(|x| (x.row as usize - 1, x.col as usize - 1)).collect())
        };
        if let Some(wells) = current_transfer_wells {
            for w in wells {
                let well_info = &mut well_infos[w.0 * columns as usize + w.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 0..rows {
        for c_column in 0..columns {
            let center = egui::pos2(
                start_x + radius + 2.0 * radius * c_column as f32,
                start_y + radius + 2.0 * radius * c_row as f32,
            );

            //
            // Draw fill first
            //

            if let Some(well_info) =
                well_infos[c_row as usize * columns as usize + c_column 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 + 1,
                        col: c_column + 1,
                    },
                    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_grid(
    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,
) {
    match pf {
        PlateFormat::W96 => add_grid_sub(
            size,
            8,
            12,
            transfers,
            plate_type,
            ordered_ids,
            cache,
            current_transfer_state,
            ui,
            state,
        ),
        PlateFormat::W384 => add_grid_sub(
            size,
            16,
            24,
            transfers,
            plate_type,
            ordered_ids,
            cache,
            current_transfer_state,
            ui,
            state,
        ),
        PlateFormat::W1536 => add_grid_sub(
            size,
            32,
            48,
            transfers,
            plate_type,
            ordered_ids,
            cache,
            current_transfer_state,
            ui,
            state,
        ),
        _ => unimplemented!(),
    }
}