mod utils;

use std::cell::RefCell;
use std::rc::Rc;
use wasm_bindgen::prelude::*;
use wasm_bindgen::Clamped;
use wasm_bindgen::JsCast;

use js_sys::Math;
use web_sys::{CanvasRenderingContext2d, HtmlCanvasElement, ImageData, Window};

const CELL_SIZE: u32 = 10;
const TICK_MS: f64 = 1000.0 / 60.0;
const HUE_PERIOD_MS: f64 = 3000.0;

/// Cells: 0 = dead, 1-255 = alive with that hue value.
const DEAD: u8 = 0;

#[inline(always)]
fn safe_hue(h: u8) -> u8 {
    if h == DEAD { 1 } else { h }
}

// ── Lookup tables (computed once) ──────────────────────────────────

struct Luts {
    /// RGBA for each hue. Index 0 = black (dead cell).
    rgb: [[u8; 4]; 256],
    sin: [f32; 256],
    cos: [f32; 256],
}

impl Luts {
    fn new() -> Self {
        let mut t = Self {
            rgb: [[0, 0, 0, 255]; 256],
            sin: [0.0; 256],
            cos: [0.0; 256],
        };
        for i in 1..256u16 {
            let h = i as f32 / 255.0 * 6.0;
            let x = 1.0 - (h % 2.0 - 1.0).abs();
            let (r, g, b) = match h as u32 {
                0 => (1.0f32, x, 0.0),
                1 => (x, 1.0, 0.0),
                2 => (0.0, 1.0, x),
                3 => (0.0, x, 1.0),
                4 => (x, 0.0, 1.0),
                _ => (1.0, 0.0, x),
            };
            t.rgb[i as usize] = [(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8, 255];
        }
        for i in 0..256 {
            let rad = i as f32 * std::f32::consts::TAU / 256.0;
            t.sin[i] = rad.sin();
            t.cos[i] = rad.cos();
        }
        t
    }

    #[inline]
    fn mix(&self, hues: &[u8]) -> u8 {
        let mut s = 0.0f32;
        let mut c = 0.0f32;
        for &h in hues {
            s += self.sin[h as usize];
            c += self.cos[h as usize];
        }
        safe_hue((s.atan2(c).rem_euclid(std::f32::consts::TAU) * 256.0 / std::f32::consts::TAU) as u8)
    }
}

// ── Grid ───────────────────────────────────────────────────────────

struct Grid {
    w: u32,
    h: u32,
    cells: Vec<u8>,
    buf: Vec<u8>,
    hues: Vec<u8>,
}

impl Grid {
    fn new(w: u32, h: u32) -> Self {
        let n = (w * h) as usize;
        let mut g = Self {
            w,
            h,
            cells: vec![DEAD; n],
            buf: vec![DEAD; n],
            hues: Vec::with_capacity(8),
        };
        g.randomize();
        g
    }

    fn randomize(&mut self) {
        for c in &mut self.cells {
            *c = if Math::random() < 0.5 {
                safe_hue((Math::random() * 255.0) as u8)
            } else {
                DEAD
            };
        }
    }

    fn tick(&mut self, luts: &Luts) {
        let (w, h) = (self.w, self.h);

        for row in 0..h {
            let n_off = (if row == 0 { h - 1 } else { row - 1 }) * w;
            let r_off = row * w;
            let s_off = (if row == h - 1 { 0 } else { row + 1 }) * w;

            for col in 0..w {
                let we = if col == 0 { w - 1 } else { col - 1 };
                let ea = if col == w - 1 { 0 } else { col + 1 };

                // Read all 8 neighbors
                let ns = [
                    self.cells[(n_off + we) as usize],
                    self.cells[(n_off + col) as usize],
                    self.cells[(n_off + ea) as usize],
                    self.cells[(r_off + we) as usize],
                    self.cells[(r_off + ea) as usize],
                    self.cells[(s_off + we) as usize],
                    self.cells[(s_off + col) as usize],
                    self.cells[(s_off + ea) as usize],
                ];

                // Count alive neighbors (branchless)
                let count = ns.iter().fold(0u8, |a, &v| a + (v != DEAD) as u8);

                let i = (r_off + col) as usize;
                let cell = self.cells[i];
                let alive = cell != DEAD;

                self.buf[i] = if alive {
                    if count == 2 || count == 3 { cell } else { DEAD }
                } else if count == 3 {
                    // Only collect hues for births
                    self.hues.clear();
                    for &v in &ns {
                        if v != DEAD {
                            self.hues.push(v);
                        }
                    }
                    luts.mix(&self.hues)
                } else {
                    DEAD
                };
            }
        }
        std::mem::swap(&mut self.cells, &mut self.buf);
    }

    fn stamp(&mut self, cr: i32, cc: i32, half: i32, hue: u8) {
        let (h, w) = (self.h as i32, self.w as i32);
        let hue = safe_hue(hue);
        for dr in -half..=half {
            for dc in -half..=half {
                let r = (cr + dr).rem_euclid(h) as u32;
                let c = (cc + dc).rem_euclid(w) as u32;
                self.cells[(r * self.w + c) as usize] = hue;
            }
        }
    }
}

// ── App state ──────────────────────────────────────────────────────

struct App {
    win: Window,
    canvas: HtmlCanvasElement,
    ctx: CanvasRenderingContext2d,
    luts: Luts,
    grid: Grid,
    pixels: Vec<u8>,
    cw: u32,
    ch: u32,
    crow: i32,
    ccol: i32,
    cursor_on: bool,
    last_tick: f64,
    dirty: bool,
}

impl App {
    fn new() -> Result<Self, JsValue> {
        utils::set_panic_hook();
        let win = web_sys::window().ok_or("no window")?;
        let doc = win.document().ok_or("no document")?;
        let canvas: HtmlCanvasElement = doc
            .get_element_by_id("canvas")
            .ok_or("no canvas")?
            .dyn_into()?;
        let ctx: CanvasRenderingContext2d = canvas.get_context("2d")?.ok_or("no 2d ctx")?.dyn_into()?;
        ctx.set_image_smoothing_enabled(false);

        let mut app = Self {
            win,
            canvas,
            ctx,
            luts: Luts::new(),
            grid: Grid::new(1, 1),
            pixels: Vec::new(),
            cw: 0,
            ch: 0,
            crow: 0,
            ccol: 0,
            cursor_on: false,
            last_tick: 0.0,
            dirty: true,
        };
        app.resize();
        Ok(app)
    }

    fn now(&self) -> f64 {
        self.win.performance().unwrap().now()
    }

    fn resize(&mut self) {
        let vw = self.win.inner_width().unwrap().as_f64().unwrap();
        let vh = self.win.inner_height().unwrap().as_f64().unwrap();
        let cols = (vw as u32 / CELL_SIZE).max(1);
        let rows = (vh as u32 / CELL_SIZE).max(1);

        self.cw = cols;
        self.ch = rows;
        self.canvas.set_width(cols);
        self.canvas.set_height(rows);
        self.canvas
            .dyn_ref::<web_sys::Element>()
            .unwrap()
            .set_attribute(
                "style",
                &format!(
                    "position:fixed;inset:0;width:{vw}px;height:{vh}px;image-rendering:pixelated"
                ),
            )
            .ok();
        self.ctx.set_image_smoothing_enabled(false);

        self.grid = Grid::new(cols, rows);
        self.pixels = vec![0u8; (cols * rows * 4) as usize];
        self.dirty = true;
    }

    fn draw(&mut self) {
        if !self.dirty {
            return;
        }
        self.dirty = false;

        let cells = &self.grid.cells;
        let rgb = &self.luts.rgb;
        let px = &mut self.pixels;

        for (i, &cell) in cells.iter().enumerate() {
            px[i * 4..i * 4 + 4].copy_from_slice(&rgb[cell as usize]);
        }

        if let Ok(img) =
            ImageData::new_with_u8_clamped_array_and_sh(Clamped(px), self.cw, self.ch)
        {
            self.ctx.put_image_data(&img, 0.0, 0.0).ok();
        }
    }
}

// ── Entry point ────────────────────────────────────────────────────

#[wasm_bindgen]
pub fn start() -> Result<(), JsValue> {
    let app = App::new()?;
    let rc = Rc::new(RefCell::new(app));

    let win = rc.borrow().win.clone();
    let doc = win.document().unwrap();

    // Mouse
    let s = rc.clone();
    let cb = Closure::wrap(Box::new(move |e: web_sys::MouseEvent| {
        let mut a = s.borrow_mut();
        a.ccol = e.client_x() / CELL_SIZE as i32;
        a.crow = e.client_y() / CELL_SIZE as i32;
        a.cursor_on = true;
    }) as Box<dyn FnMut(_)>);
    doc.add_event_listener_with_callback("mousemove", cb.as_ref().unchecked_ref())?;
    cb.forget();

    // Touch move
    let s = rc.clone();
    let cb = Closure::wrap(Box::new(move |e: web_sys::TouchEvent| {
        e.prevent_default();
        if let Some(t) = e.touches().get(0) {
            let mut a = s.borrow_mut();
            a.ccol = t.client_x() / CELL_SIZE as i32;
            a.crow = t.client_y() / CELL_SIZE as i32;
            a.cursor_on = true;
        }
    }) as Box<dyn FnMut(_)>);
    doc.add_event_listener_with_callback("touchmove", cb.as_ref().unchecked_ref())?;
    cb.forget();

    // Touch start
    let s = rc.clone();
    let cb = Closure::wrap(Box::new(move |e: web_sys::TouchEvent| {
        e.prevent_default();
        if let Some(t) = e.touches().get(0) {
            let mut a = s.borrow_mut();
            a.ccol = t.client_x() / CELL_SIZE as i32;
            a.crow = t.client_y() / CELL_SIZE as i32;
            a.cursor_on = true;
        }
    }) as Box<dyn FnMut(_)>);
    doc.add_event_listener_with_callback("touchstart", cb.as_ref().unchecked_ref())?;
    cb.forget();

    // Touch end
    let s = rc.clone();
    let cb = Closure::wrap(Box::new(move |_: web_sys::TouchEvent| {
        s.borrow_mut().cursor_on = false;
    }) as Box<dyn FnMut(_)>);
    doc.add_event_listener_with_callback("touchend", cb.as_ref().unchecked_ref())?;
    cb.forget();

    // Resize
    let s = rc.clone();
    let cb = Closure::wrap(Box::new(move || {
        s.borrow_mut().resize();
    }) as Box<dyn FnMut()>);
    win.add_event_listener_with_callback("resize", cb.as_ref().unchecked_ref())?;
    cb.forget();

    // Animation loop
    let f: Rc<RefCell<Option<Closure<dyn FnMut()>>>> = Rc::new(RefCell::new(None));
    let g = f.clone();
    let s = rc.clone();
    let w = win.clone();

    *g.borrow_mut() = Some(Closure::wrap(Box::new(move || {
        let now = s.borrow().now();

        {
            let mut a = s.borrow_mut();

            if now - a.last_tick >= TICK_MS {
                a.last_tick = now;
                let luts = &a.luts as *const Luts;
                // SAFETY: luts is not mutated during tick()
                a.grid.tick(unsafe { &*luts });
                a.dirty = true;
            }

            if a.cursor_on {
                let (cr, cc) = (a.crow, a.ccol);
                let hue = ((now % HUE_PERIOD_MS) / HUE_PERIOD_MS * 256.0) as u8;
                a.grid.stamp(cr, cc, 2, hue);
                a.dirty = true;
            }

            a.draw();
        }

        w.request_animation_frame(f.borrow().as_ref().unwrap().as_ref().unchecked_ref())
            .ok();
    }) as Box<dyn FnMut()>));

    win.request_animation_frame(g.borrow().as_ref().unwrap().as_ref().unchecked_ref())?;
    Ok(())
}
