#![feature(if_let_guard)]

use quanta::{Clock, Instant};
use std::collections::HashMap;
use std::ops::DerefMut;
use std::sync::Arc;
use std::time::Duration;
use tiny_skia::{Color, *};
use tokio::sync::{OnceCell, mpsc};
use winit::{
    application::ApplicationHandler,
    event::{ElementState, MouseButton, TouchPhase, WindowEvent},
    event_loop::{ActiveEventLoop, ControlFlow, EventLoop},
    window::{Window, WindowAttributes, WindowId},
};

use crate::ws::LaserMessage;
use crate::{
    renderer::{Renderer, skia_rgba_to_bgra_u32},
    ws::WsMessage,
};

mod renderer;
mod utils;
mod ws;

type BoxedError = Box<dyn std::error::Error + 'static>;
type AppResult<T> = Result<T, BoxedError>;

const BASE_MAX_AGE: Duration = Duration::from_millis(200);
const MIN_AGE: Duration = Duration::from_millis(50);
const FAST_DECAY_THRESHOLD: usize = 30;
#[cfg(target_arch = "wasm32")]
const TARGET_FPS: u32 = 120;
#[cfg(not(target_arch = "wasm32"))]
const TARGET_FPS: u32 = 60;
const FRAME_TIME_MS: u64 = 1000 / TARGET_FPS as u64; // ~16.67ms

#[derive(Clone)]
struct LaserPoint {
    x: f32,
    y: f32,
    color: [u8; 3],
    created_at: Instant,
}

impl LaserPoint {
    fn new(msg: LaserMessage, now: Instant) -> Self {
        Self {
            x: msg.x as f32,
            y: msg.y as f32,
            color: utils::id_to_color(msg.id),
            created_at: now,
        }
    }
}

pub type WindowHandle = Arc<OnceCell<Arc<Window>>>;

pub struct Graphics {
    window: Arc<Window>,
    renderer: Renderer,
    pixmap: Pixmap,
}

impl Graphics {
    fn resize(&mut self, width: u32, height: u32) -> AppResult<()> {
        self.renderer.resize(width, height)?;
        self.pixmap = Pixmap::new(width, height).unwrap();
        Ok(())
    }
}

pub struct LaserOverlay {
    gfx: Option<Graphics>,
    window: WindowHandle,
    server_mouse_pos: (f32, f32),
    laser_points: HashMap<(u64, u8), Vec<LaserPoint>, ahash::RandomState>,
    in_chan: (mpsc::Sender<WsMessage>, mpsc::Receiver<WsMessage>),
    out_tx: mpsc::Sender<WsMessage>,
    last_render: Instant,
    last_cleanup: Instant,
    clock: Clock,
    client_id: u64,
    mouse_pressed: bool,
    mouse_pos: (f32, f32),
    current_line_id: u8,
    next_line_id: u8,
    needs_redraw: bool,
    has_any_points: bool,
}

impl LaserOverlay {
    pub fn new() -> (mpsc::Sender<WsMessage>, mpsc::Receiver<WsMessage>, Self) {
        let in_chan = mpsc::channel(1024);
        let (out_tx, out_rx) = mpsc::channel(512);

        let clock = Clock::new();
        let now = clock.now();

        let this = Self {
            gfx: None,
            window: WindowHandle::default(),
            server_mouse_pos: (0.0, 0.0),
            laser_points: Default::default(),
            in_chan,
            out_tx,
            last_render: now,
            last_cleanup: now,
            clock,
            client_id: fastrand::u64(..),
            mouse_pressed: false,
            mouse_pos: (0.0, 0.0),
            current_line_id: 0,
            next_line_id: 1,
            needs_redraw: false,
            has_any_points: false,
        };

        (this.in_chan.0.clone(), out_rx, this)
    }

    #[cfg(feature = "server")]
    pub fn start_mouse_listener(send_tx: tokio::sync::broadcast::Sender<(u64, WsMessage)>) {
        std::thread::spawn({
            use enigo::{Enigo, Mouse, Settings};

            move || {
                let enigo = Enigo::new(&Settings::default()).unwrap();
                loop {
                    let res = enigo.location();
                    let Ok(pos) = res else {
                        eprintln!("failed to get mouse position: {res:?}");
                        continue;
                    };
                    let msg = WsMessage::Mouse(ws::MouseMessage {
                        x: pos.0 as u32,
                        y: pos.1 as u32,
                    });
                    let _ = send_tx.send((0, msg));
                    std::thread::sleep(Duration::from_millis(1000 / 30));
                }
            }
        });
    }

    pub fn window_handle(&self) -> WindowHandle {
        self.window.clone()
    }

    pub fn init_graphics(&mut self, window: Arc<Window>) -> AppResult<()> {
        #[cfg(target_arch = "wasm32")]
        let size = {
            use winit::platform::web::WindowExtWebSys;
            let canvas = window.canvas().unwrap();
            let (w, h) = (canvas.client_width(), canvas.client_height());
            canvas.set_width(w.try_into().unwrap());
            canvas.set_height(h.try_into().unwrap());
            winit::dpi::PhysicalSize::new(w as u32, h as u32)
        };
        #[cfg(not(target_arch = "wasm32"))]
        let size = window.inner_size();

        let _ = self.window.set(window.clone());
        self.gfx = Some(Graphics {
            renderer: Renderer::new(window.clone(), size.width, size.height)?,
            pixmap: Pixmap::new(size.width, size.height).unwrap(),
            window,
        });

        self.needs_redraw = true;

        Ok(())
    }

    fn handle_mouse_press(&mut self, position: (f32, f32)) {
        self.mouse_pressed = true;
        self.mouse_pos = position;
        self.current_line_id = self.next_line_id;
        self.next_line_id = self.next_line_id.wrapping_add(1);
        self.needs_redraw = true;

        let msg = WsMessage::Laser(LaserMessage {
            x: position.0 as u32,
            y: position.1 as u32,
            id: self.client_id,
            line_id: self.current_line_id,
        });
        let _ = self.in_chan.0.try_send(msg);
        let _ = self.out_tx.try_send(msg);
    }

    fn handle_mouse_move(&mut self, position: (f32, f32)) {
        if !self.mouse_pressed {
            return;
        }

        let dx = position.0 - self.mouse_pos.0;
        let dy = position.1 - self.mouse_pos.1;
        let distance = (dx * dx + dy * dy).sqrt();

        if distance < 1.0 {
            return;
        }

        self.mouse_pos = position;
        self.needs_redraw = true;
        let msg = WsMessage::Laser(LaserMessage {
            x: position.0 as u32,
            y: position.1 as u32,
            id: self.client_id,
            line_id: self.current_line_id,
        });
        let _ = self.in_chan.0.try_send(msg);
        let _ = self.out_tx.try_send(msg);
    }

    fn handle_mouse_release(&mut self) {
        self.mouse_pressed = false;
    }

    fn calculate_point_max_age(point_index: usize, total_points: usize) -> Duration {
        if total_points <= FAST_DECAY_THRESHOLD {
            return BASE_MAX_AGE;
        }

        if point_index < FAST_DECAY_THRESHOLD {
            let progress = point_index as f32 / FAST_DECAY_THRESHOLD as f32;
            let age_range = BASE_MAX_AGE.as_millis() - MIN_AGE.as_millis();
            let calculated_age = MIN_AGE.as_millis() + (age_range as f32 * progress) as u128;
            Duration::from_millis(calculated_age as u64)
        } else {
            BASE_MAX_AGE
        }
    }

    #[inline(always)]
    fn ingest_points(&mut self) {
        let mut has_any_points = false;

        let should_cleanup = self.last_cleanup.elapsed().as_millis() >= FRAME_TIME_MS as u128;

        if should_cleanup {
            self.last_cleanup = self.clock.now();

            let mut ids_to_remove = Vec::new();
            for (id, points) in self.laser_points.iter_mut() {
                let points_len = points.len();
                if points_len == 0 {
                    continue;
                }

                let mut new_points = Vec::with_capacity(points_len);
                for (index, point) in points.iter().enumerate() {
                    let max_age = Self::calculate_point_max_age(index, points_len);
                    if point.created_at.elapsed() < max_age {
                        new_points.push(point.clone());
                    }
                }

                if new_points.len() != points_len {
                    self.needs_redraw = true;
                }
                *points = new_points;

                if points.is_empty() {
                    ids_to_remove.push(*id);
                } else {
                    has_any_points = true;
                }
            }

            for id in ids_to_remove {
                self.laser_points.remove(&id);
            }
        } else {
            has_any_points = self.laser_points.values().any(|points| !points.is_empty());
        }

        while let Ok(msg) = self.in_chan.1.try_recv() {
            match msg {
                WsMessage::Laser(msg) => {
                    self.laser_points
                        .entry((msg.id, msg.line_id))
                        .or_default()
                        .push(LaserPoint::new(msg, self.clock.now()));
                    has_any_points = true;
                }
                #[cfg(feature = "client")]
                WsMessage::Mouse(msg) => {
                    self.server_mouse_pos = (msg.x as f32, msg.y as f32);
                    self.needs_redraw = true;
                }
                #[cfg(not(feature = "client"))]
                WsMessage::Mouse(_) => {}
            }
        }

        self.has_any_points = has_any_points;
    }

    #[cfg(feature = "client")]
    #[inline(always)]
    fn draw_server_mouse(mut pixmap: PixmapMut, mouse_pos: (f32, f32)) {
        let (x, y) = mouse_pos;
        let radius = 10.0;
        let color = Color::WHITE;

        let mut pb = PathBuilder::new();
        pb.push_circle(x, y, radius);

        if let Some(path) = pb.finish() {
            let mut paint = Paint::default();
            paint.set_color(color);
            paint.anti_alias = true;
            paint.blend_mode = BlendMode::Source;

            let mut stroke = Stroke::default();
            stroke.width = radius * 2.0;
            stroke.line_cap = LineCap::Round;
            stroke.line_join = LineJoin::Round;

            pixmap.stroke_path(&path, &paint, &stroke, Transform::identity(), None);
        }
    }

    #[inline(always)]
    fn draw_tapering_laser_line(mut pixmap: PixmapMut, points: &[LaserPoint], now: Instant) {
        if points.len() < 2 {
            return;
        }

        let max_width = 8.0;
        let min_width = 0.5;
        let points_len = points.len();

        let data = points
            .windows(2)
            .enumerate()
            .map(|(i, chunk)| {
                let current = &chunk[0];
                let next = &chunk[1];

                let progress = i as f32 / (points_len - 2) as f32;
                let width = min_width + (max_width - min_width) * progress;

                let point_index = i;
                let max_age = Self::calculate_point_max_age(point_index, points_len);
                let age = now.duration_since(current.created_at);
                let age_progress = age.as_millis() as f32 / max_age.as_millis() as f32;
                let alpha = (255.0 * (1.0 - age_progress.clamp(0.0, 1.0))) as u8;

                let mut pb = PathBuilder::new();
                pb.move_to(current.x, current.y);
                pb.line_to(next.x, next.y);

                (pb.finish(), width, next.color, alpha)
            })
            .collect::<Vec<_>>();

        // draw glow first so we can draw the actual line on top
        // otherwise the glow would cover the line
        for (path, width, color, alpha) in data.iter().filter(|p| p.1 > 2.0) {
            let Some(path) = path else {
                continue;
            };

            let mut glow_paint = Paint::default();
            glow_paint.set_color_rgba8(color[0], color[1], color[2], alpha / 5);
            glow_paint.anti_alias = true;
            // replace the existing alpha
            glow_paint.blend_mode = BlendMode::Source;

            let mut glow_stroke = Stroke::default();
            glow_stroke.width = width * 1.8;
            glow_stroke.line_cap = LineCap::Round;
            glow_stroke.line_join = LineJoin::Round;

            pixmap.stroke_path(
                &path,
                &glow_paint,
                &glow_stroke,
                Transform::identity(),
                None,
            );
        }

        for (path, width, color, alpha) in data {
            let Some(path) = path else {
                continue;
            };

            let mut paint = Paint::default();
            paint.set_color_rgba8(color[0], color[1], color[2], alpha.max(10));
            paint.anti_alias = true;
            // replace the existing alpha (so it doesn't blend with existing pixels, looks bad)
            paint.blend_mode = BlendMode::Source;

            let mut stroke = Stroke::default();
            stroke.width = width;
            stroke.line_cap = LineCap::Round;
            stroke.line_join = LineJoin::Round;

            pixmap.stroke_path(&path, &paint, &stroke, Transform::identity(), None);
        }
    }

    #[inline(always)]
    pub fn render(&mut self) -> AppResult<()> {
        let Some(gfx) = self.gfx.as_mut() else {
            return Ok(());
        };

        let mut frame = gfx.renderer.frame_mut()?;

        gfx.pixmap.fill(Color::TRANSPARENT);
        for points in self.laser_points.values() {
            Self::draw_tapering_laser_line(gfx.pixmap.as_mut(), points, self.clock.now());
        }
        #[cfg(feature = "client")]
        Self::draw_server_mouse(gfx.pixmap.as_mut(), self.server_mouse_pos);
        skia_rgba_to_bgra_u32(gfx.pixmap.data(), frame.deref_mut());

        gfx.window.pre_present_notify();
        gfx.renderer.present()?;

        self.last_render = self.clock.now();
        self.needs_redraw = false;

        Ok(())
    }

    #[inline(always)]
    pub fn should_render(&self) -> bool {
        self.has_any_points || self.needs_redraw
    }
}

impl ApplicationHandler for LaserOverlay {
    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
        let attrs = WindowAttributes::default()
            .with_title("laser overlay")
            .with_transparent(true);

        #[cfg(feature = "server")]
        let attrs = attrs
            .with_fullscreen(Some(winit::window::Fullscreen::Borderless(None)))
            .with_window_level(winit::window::WindowLevel::AlwaysOnTop)
            .with_decorations(false);

        #[cfg(target_arch = "wasm32")]
        let attrs = winit::platform::web::WindowAttributesExtWebSys::with_append(attrs, true);

        let window = event_loop.create_window(attrs).unwrap();
        #[cfg(feature = "server")]
        let _ = window.set_cursor_hittest(false);

        self.init_graphics(window.into()).unwrap();

        event_loop.set_control_flow(ControlFlow::wait_duration(Duration::from_millis(
            FRAME_TIME_MS,
        )));
    }

    fn window_event(
        &mut self,
        event_loop: &ActiveEventLoop,
        _window_id: WindowId,
        event: WindowEvent,
    ) {
        match event {
            WindowEvent::CloseRequested => {
                event_loop.exit();
                return;
            }
            WindowEvent::Touch(touch) => {
                let pos = (touch.location.x as f32, touch.location.y as f32);
                match touch.phase {
                    TouchPhase::Started => {
                        self.handle_mouse_press(pos);
                    }
                    TouchPhase::Moved => {
                        self.handle_mouse_move(pos);
                    }
                    TouchPhase::Ended | TouchPhase::Cancelled => {
                        self.handle_mouse_release();
                    }
                }
            }
            WindowEvent::MouseInput { state, button, .. } => {
                if button == MouseButton::Left {
                    match state {
                        ElementState::Pressed => {
                            self.handle_mouse_press(self.mouse_pos);
                        }
                        ElementState::Released => {
                            self.handle_mouse_release();
                        }
                    }
                }
            }
            WindowEvent::CursorMoved { position, .. } => {
                let pos = (position.x as f32, position.y as f32);
                self.handle_mouse_move(pos);
                self.mouse_pos = pos;
            }
            WindowEvent::Resized(size) if let Some(gfx) = self.gfx.as_mut() => {
                gfx.resize(size.width, size.height).unwrap();
            }
            WindowEvent::RedrawRequested => {
                self.ingest_points();
                self.render().unwrap();
            }
            _ => {}
        }
    }

    fn about_to_wait(&mut self, event_loop: &ActiveEventLoop) {
        if self.should_render() {
            if let Some(gfx) = self.gfx.as_ref() {
                if self.last_render.elapsed().as_millis() >= FRAME_TIME_MS as u128 {
                    gfx.window.request_redraw();
                }
            }
            event_loop.set_control_flow(ControlFlow::wait_duration(Duration::from_millis(
                FRAME_TIME_MS,
            )));
        } else {
            event_loop.set_control_flow(ControlFlow::Wait);
        }
    }
}

#[allow(unused_mut)]
fn run_app(event_loop: EventLoop<()>, mut app: LaserOverlay) -> AppResult<()> {
    #[cfg(not(any(target_arch = "wasm32", target_arch = "wasm64")))]
    event_loop.run_app(&mut app)?;
    #[cfg(any(target_arch = "wasm32", target_arch = "wasm64"))]
    {
        console_error_panic_hook::set_once();
        winit::platform::web::EventLoopExtWebSys::spawn_app(event_loop, app);
    }
    Ok(())
}

fn run() -> AppResult<()> {
    let event_loop = EventLoop::new()?;
    let (_tx, _rx, app) = LaserOverlay::new();

    #[cfg(any(feature = "server", feature = "client"))]
    let window = app.window_handle();

    #[cfg(feature = "server")]
    tokio::spawn({
        let window = window.clone();
        let tx = _tx.clone();
        async move {
            let (server, send_tx) = ws::server::listen(3111, window, tx).await.unwrap();
            LaserOverlay::start_mouse_listener(send_tx);
            server.await;
        }
    });
    #[cfg(feature = "client")]
    {
        #[cfg(not(any(target_arch = "wasm32", target_arch = "wasm64")))]
        let _ = tokio_rustls::rustls::crypto::ring::default_provider().install_default();
        let client_id = app.client_id;
        let fut = async move {
            ws::client::connect(env!("SERVER_URL"), window, _rx, _tx, client_id)
                .await
                .unwrap()
        };
        #[cfg(not(target_arch = "wasm32"))]
        tokio::spawn(fut);
        #[cfg(target_arch = "wasm32")]
        wasm_bindgen_futures::spawn_local(fut);
    }

    run_app(event_loop, app)?;

    Ok(())
}

#[cfg(not(target_arch = "wasm32"))]
#[tokio::main]
async fn main() -> AppResult<()> {
    run()
}

#[cfg(target_arch = "wasm32")]
fn main() -> AppResult<()> {
    run()
}