use avm_stats::{PixelType, PlaneType};
use convert_case::{Case, Casing};
use egui::{pos2, vec2, Align2, Color32, Mesh, Rect, Rounding, Shape, Stroke, TextStyle, Ui};

use crate::app_state::AppState;
use crate::image_manager::ImageType;
use crate::stream::{ChangeFrame, CurrentFrame, FrameStatus};

use crate::views::render_view::RenderView;

use super::frame_viewer::REF_FRAME_COLORS;
use super::{SelectedObjectKind, ViewMode};

// TODO(comc): There must be a built-in way to do horizontal scrolling with the mouse wheel in egui.
const SCROLL_FACTOR: f32 = 0.2;
pub struct FrameSelectViewer;
impl RenderView for FrameSelectViewer {
    fn title(&self) -> String {
        "Frame Select".into()
    }
    fn render(&self, ui: &mut Ui, state: &mut AppState) -> anyhow::Result<()> {
        let mut set_current_frame_index = None;
        let Some(stream) = state.stream.as_ref() else {
            return Ok(());
        };
        let scroll_to = state.settings.scroll_to_index.take();
        let mut motion_vector_arrows: Vec<(usize, i32)> = Vec::new();
        if state.settings.sharable.view_mode == ViewMode::Motion {
            if let Some(selected_object) = state.settings.selected_object.as_ref() {
                if let SelectedObjectKind::CodingUnit(cu) = selected_object.kind {
                    if let Some(frame) = stream.current_frame() {
                        if let Some(cu) = cu.try_resolve(frame) {
                            for mv in cu.coding_unit.prediction_mode.as_ref().unwrap().motion_vectors.iter() {
                                if mv.ref_frame_is_inter && !mv.ref_frame_is_tip {
                                    let order_hint = mv.ref_frame_order_hint;
                                    if let Some(frame_index) = stream.lookup_order_hint(order_hint) {
                                        if !motion_vector_arrows
                                            .iter()
                                            .any(|(other_frame_index, _)| *other_frame_index == frame_index)
                                        {
                                            motion_vector_arrows.push((frame_index, mv.ref_frame))
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        let sorted_frame_indices = stream.get_sorted_frames(state.settings.persistent.frame_sort_order);
        let mut frame_rects = vec![Rect::NOTHING; sorted_frame_indices.len()];
        egui::ScrollArea::horizontal()
            .max_height(200.0)
            .drag_to_scroll(true)
            .vscroll(false)
            .show(ui, |ui| {
                egui::Grid::new("img_grid").show(ui, |ui| {
                    for &frame_id in sorted_frame_indices.iter() {
                        let frame = stream.get_frame(frame_id);
                        let (id, rect) = ui.allocate_space(egui::vec2(100.0, 100.0));
                        frame_rects[frame_id] = rect;
                        if let Some(scroll_to) = scroll_to {
                            if scroll_to == frame_id {
                                ui.scroll_to_rect(rect, None);
                            }
                        }
                        let response = ui.interact(rect, id, egui::Sense::click());
                        let scroll_delta = ui.input(|i| i.scroll_delta);
                        if scroll_delta.length() > 0.0 {
                            let scroll_delta = SCROLL_FACTOR * vec2(scroll_delta.y, scroll_delta.x);
                            ui.scroll_with_delta(scroll_delta);
                        }
                        ui.painter()
                            .add(Shape::rect_filled(rect, Rounding::ZERO, Color32::GRAY));

                        let frame_id_text = if let FrameStatus::Loaded(frame) = frame {
                            let display_index = frame.frame_params.as_ref().unwrap().display_index;
                            format!("Frame {frame_id} ({display_index})")
                        } else {
                            format!("Frame {frame_id}")
                        };
                        let text = ui.fonts(|f| {
                            // TODO(comc): Use display or decode index instead of "frame_id":
                            Shape::text(
                                f,
                                rect.center_top(),
                                Align2::CENTER_TOP,
                                frame_id_text,
                                TextStyle::Body.resolve(ui.style()),
                                Color32::BLACK,
                            )
                        });
                        ui.painter().add(text);

                        match frame {
                            FrameStatus::OutOfBounds => {}
                            FrameStatus::Decoding => {
                                let text = ui.fonts(|f| {
                                    Shape::text(
                                        f,
                                        rect.center(),
                                        Align2::CENTER_CENTER,
                                        "Decoding...",
                                        TextStyle::Body.resolve(ui.style()),
                                        Color32::BLACK,
                                    )
                                });
                                ui.painter().add(text);
                            }
                            FrameStatus::Invalid => {
                                let text = ui.fonts(|f| {
                                    Shape::text(
                                        f,
                                        rect.center(),
                                        Align2::CENTER_CENTER,
                                        "ERROR",
                                        TextStyle::Body.resolve(ui.style()),
                                        Color32::BLACK,
                                    )
                                });
                                ui.painter().add(text);
                            }
                            FrameStatus::Pending => {
                                let text = ui.fonts(|f| {
                                    Shape::text(
                                        f,
                                        rect.center(),
                                        Align2::CENTER_CENTER,
                                        "Loading...",
                                        TextStyle::Body.resolve(ui.style()),
                                        Color32::BLACK,
                                    )
                                });
                                ui.painter().add(text);
                            }
                            FrameStatus::Loaded(frame) => {
                                // TODO(comc): Don't unwrap!
                                let image_type =
                                    ImageType::new(PlaneType::Rgb, PixelType::Reconstruction, false, false);
                                let Ok(texture_handle) = stream.images.get_or_create_image(
                                    ui.ctx(),
                                    &stream.pixel_data,
                                    frame,
                                    image_type,
                                    &state.settings.sharable.heatmap_settings,
                                ) else {
                                    return;
                                };
                                let mut image_mesh = Mesh::with_texture(texture_handle.id());
                                // TODO(comc): preserve aspect ratio.
                                let image_uv = Rect::from_min_size(pos2(0.0, 0.0), vec2(1.0, 1.0));
                                let mut image_rect = rect;
                                image_rect.set_top(rect.top() + 20.0);
                                image_rect.set_bottom(rect.bottom() - 20.0);
                                image_mesh.add_rect_with_uv(image_rect, image_uv, Color32::WHITE);
                                ui.painter().add(image_mesh);

                                if response.clicked() {
                                    set_current_frame_index = Some(frame_id);
                                }
                                let mut frame_type_name =
                                    frame.frame_type_name().trim_end_matches("_FRAME").to_case(Case::Pascal);
                                let tip_mode = frame.tip_mode_name();
                                if tip_mode == "TIP_FRAME_AS_OUTPUT" {
                                    frame_type_name = format!("{frame_type_name} (TIP)");
                                }

                                let text = ui.fonts(|f| {
                                    Shape::text(
                                        f,
                                        rect.center_bottom(),
                                        Align2::CENTER_BOTTOM,
                                        frame_type_name,
                                        TextStyle::Body.resolve(ui.style()),
                                        Color32::BLACK,
                                    )
                                });
                                ui.painter().add(text);
                            }
                            FrameStatus::Unloaded => {
                                let text = ui.fonts(|f| {
                                    Shape::text(
                                        f,
                                        rect.center(),
                                        Align2::CENTER_CENTER,
                                        "Click to load",
                                        TextStyle::Body.resolve(ui.style()),
                                        Color32::BLACK,
                                    )
                                });
                                ui.painter().add(text);

                                if response.clicked() {
                                    set_current_frame_index = Some(frame_id);
                                }
                            }
                        };

                        if let Some(current_frame) = state.stream.current_frame() {
                            if current_frame.decode_index() == frame_id {
                                ui.painter().add(Shape::rect_stroke(
                                    rect,
                                    Rounding::ZERO,
                                    Stroke::new(3.0, Color32::YELLOW),
                                ));
                            }
                        }
                    }
                });
                let painter = ui.painter();
                for (mv_index, &(frame_index, ref_frame)) in motion_vector_arrows.iter().enumerate() {
                    let start_rect = frame_rects[stream.current_frame_index];
                    let end_rect = frame_rects[frame_index];
                    let color_index = ref_frame as usize % 8;
                    let color = REF_FRAME_COLORS[color_index];
                    let pos_offset = vec2(0.0, mv_index as f32 * 5.0);
                    painter.add(Shape::line_segment(
                        [start_rect.center() + pos_offset, end_rect.center() + pos_offset],
                        Stroke::new(2.0, color),
                    ));
                    painter.add(Shape::circle_filled(
                        start_rect.center() + pos_offset,
                        2.0,
                        REF_FRAME_COLORS[color_index],
                    ));

                    let dir_sign = (start_rect.left() - end_rect.left()).signum();
                    let triangle_tip = end_rect.center() + pos_offset;
                    let triangle_vertices = vec![
                        triangle_tip,
                        triangle_tip + vec2(dir_sign * 3.0, 3.0),
                        triangle_tip + vec2(dir_sign * 3.0, -3.0),
                    ];
                    painter.add(Shape::convex_polygon(
                        triangle_vertices,
                        REF_FRAME_COLORS[color_index],
                        Stroke::NONE,
                    ));

                    painter.add(Shape::rect_stroke(end_rect, Rounding::ZERO, Stroke::new(3.0, color)));
                }
            });
        if let Some(current_frame_index) = set_current_frame_index {
            // Unwrapping is okay, since we already know stream exists at this point.
            let stream = state.stream.as_mut().unwrap();
            stream.change_frame(ChangeFrame::index(current_frame_index));
            state.settings.selected_object = None;
        }
        Ok(())
    }
}