use bone_render::{
    AtlasGrid, AtlasPage, ClearColor, OffscreenContext, Result, SnapshotFrame, SolidRenderer,
    Style, TILE_PAD, ViewportExtent, ViewportPx,
};
use bone_types::{DisplayMode, IconId};

use crate::icon::{IconModel, icon_camera, icon_style};
use crate::models::model;
use crate::offscreen_context;

pub const MASTER_TILE: u32 = 192;

const SHADOW_OFFSET: usize = 7;
const SHADOW_BLUR: usize = 9;
const SHADOW_ALPHA: u32 = 90;
const SHADOW_LEVEL: u32 = 10;

#[derive(Clone, Debug)]
pub struct BakedPage {
    pub page: AtlasPage,
    pub extent: ViewportExtent,
    pub rgba: Vec<u8>,
}

pub fn bake() -> Result<Vec<BakedPage>> {
    let ctx = offscreen_context(ViewportExtent::square(ViewportPx::new(MASTER_TILE)))?;
    let mut renderer = SolidRenderer::new(ctx.gpu(), ctx.color_format());
    let master = to_usize(MASTER_TILE);
    let shadowed: Vec<(IconId, Vec<u8>)> = IconId::ALL
        .iter()
        .map(|&id| {
            let premul = render_master(&mut renderer, &ctx, &model(id))?;
            Ok((id, with_shadow(master, &premul)))
        })
        .collect::<Result<Vec<_>>>()?;

    let grid = AtlasGrid::DEFAULT;
    Ok(AtlasPage::ALL
        .into_iter()
        .map(|page| {
            let page_side = to_usize(page.side_px());
            let tiles: Vec<(IconId, Vec<u8>)> = shadowed
                .iter()
                .map(|(id, premul)| (*id, downsample(master, premul, page_side)))
                .collect();
            let (width, height) = grid.pixel_size(page);
            BakedPage {
                page,
                extent: ViewportExtent::new(ViewportPx::new(width), ViewportPx::new(height)),
                rgba: pack(grid, page, &tiles),
            }
        })
        .collect())
}

fn render_master(
    renderer: &mut SolidRenderer,
    ctx: &OffscreenContext,
    m: &IconModel,
) -> Result<Vec<u8>> {
    let black = render_one(renderer, ctx, m, ClearColor::opaque(0.0, 0.0, 0.0))?;
    let white = render_one(renderer, ctx, m, ClearColor::opaque(1.0, 1.0, 1.0))?;
    Ok(black
        .rgba()
        .chunks_exact(4)
        .zip(white.rgba().chunks_exact(4))
        .flat_map(|(dark, light)| {
            let drop =
                |channel: usize| u32::from(light[channel]).saturating_sub(u32::from(dark[channel]));
            let coverage_drop = (drop(0) + drop(1) + drop(2)) / 3;
            let alpha = u8_of(255u32.saturating_sub(coverage_drop));
            [dark[0], dark[1], dark[2], alpha]
        })
        .collect())
}

fn render_one(
    renderer: &mut SolidRenderer,
    ctx: &OffscreenContext,
    m: &IconModel,
    background: ClearColor,
) -> Result<SnapshotFrame> {
    let (scene, edges) = m.tessellate();
    let style: Style = icon_style().with_background(background);
    renderer.render_display(
        ctx,
        &scene,
        &edges,
        icon_camera(m.view()),
        &style,
        DisplayMode::ShadedWithEdges,
    )
}

fn with_shadow(side: usize, icon: &[u8]) -> Vec<u8> {
    let alpha: Vec<u8> = icon.chunks_exact(4).map(|px| px[3]).collect();
    let blurred = box_blur(side, &alpha, SHADOW_BLUR);
    icon.chunks_exact(4)
        .enumerate()
        .flat_map(|(index, top)| {
            let x = index % side;
            let y = index / side;
            let source = if y >= SHADOW_OFFSET {
                u32::from(blurred[(y - SHADOW_OFFSET) * side + x])
            } else {
                0
            };
            let shadow_alpha = source * SHADOW_ALPHA / 255;
            let shadow_rgb = u8_of(SHADOW_LEVEL * shadow_alpha / 255);
            let shadow = [shadow_rgb, shadow_rgb, shadow_rgb, u8_of(shadow_alpha)];
            over([top[0], top[1], top[2], top[3]], shadow)
        })
        .collect()
}

fn box_blur(side: usize, src: &[u8], radius: usize) -> Vec<u8> {
    let horizontal = blur_axis(side, src, radius, true);
    blur_axis(side, &horizontal, radius, false)
}

fn blur_axis(side: usize, src: &[u8], radius: usize, horizontal: bool) -> Vec<u8> {
    (0..side)
        .flat_map(|y| (0..side).map(move |x| (x, y)))
        .map(|(x, y)| {
            let center = if horizontal { x } else { y };
            let lo = center.saturating_sub(radius);
            let hi = (center + radius).min(side - 1);
            let (sum, count) = (lo..=hi).fold((0u32, 0u32), |(sum, count), p| {
                let (sx, sy) = if horizontal { (p, y) } else { (x, p) };
                (sum + u32::from(src[sy * side + sx]), count + 1)
            });
            div_round(sum, count)
        })
        .collect()
}

fn downsample(master_side: usize, master: &[u8], page_side: usize) -> Vec<u8> {
    assert_eq!(
        master_side % page_side,
        0,
        "master tile must divide evenly into each atlas page side",
    );
    let factor = master_side / page_side;
    let area = to_u32(factor * factor);
    (0..page_side)
        .flat_map(|oy| (0..page_side).map(move |ox| (ox, oy)))
        .flat_map(|(ox, oy)| {
            let sums = (0..factor)
                .flat_map(|dy| (0..factor).map(move |dx| (dx, dy)))
                .fold([0u32; 4], |mut acc, (dx, dy)| {
                    let base = ((oy * factor + dy) * master_side + (ox * factor + dx)) * 4;
                    acc[0] += u32::from(master[base]);
                    acc[1] += u32::from(master[base + 1]);
                    acc[2] += u32::from(master[base + 2]);
                    acc[3] += u32::from(master[base + 3]);
                    acc
                });
            [
                div_round(sums[0], area),
                div_round(sums[1], area),
                div_round(sums[2], area),
                div_round(sums[3], area),
            ]
        })
        .collect()
}

fn pack(grid: AtlasGrid, page: AtlasPage, tiles: &[(IconId, Vec<u8>)]) -> Vec<u8> {
    let (width, height) = grid.pixel_size(page);
    let width = to_usize(width);
    let height = to_usize(height);
    let page_side = to_usize(page.side_px());
    let cell = to_usize(page.cell_px());
    let pad = to_usize(TILE_PAD);
    tiles
        .iter()
        .fold(vec![0u8; width * height * 4], |mut atlas, (id, tile)| {
            let grid_cell = grid.cell(id.tile());
            let cell_x = grid_cell.col * cell;
            let cell_y = grid_cell.row * cell;
            (0..cell).for_each(|ly| {
                (0..cell).for_each(|lx| {
                    let sx = lx.saturating_sub(pad).min(page_side - 1);
                    let sy = ly.saturating_sub(pad).min(page_side - 1);
                    let src = (sy * page_side + sx) * 4;
                    let dst = ((cell_y + ly) * width + (cell_x + lx)) * 4;
                    atlas[dst..dst + 4].copy_from_slice(&tile[src..src + 4]);
                });
            });
            atlas
        })
}

fn over(top: [u8; 4], bottom: [u8; 4]) -> [u8; 4] {
    let inverse = 255 - u32::from(top[3]);
    let blend = |channel: usize| {
        u8_of(u32::from(top[channel]) + (inverse * u32::from(bottom[channel]) + 127) / 255)
    };
    [blend(0), blend(1), blend(2), blend(3)]
}

fn div_round(sum: u32, count: u32) -> u8 {
    u8_of((sum + count / 2) / count)
}

fn u8_of(value: u32) -> u8 {
    let Ok(byte) = u8::try_from(value.min(255)) else {
        panic!("value clamped to 255 fits a u8");
    };
    byte
}

fn to_usize(value: u32) -> usize {
    let Ok(out) = usize::try_from(value) else {
        panic!("atlas dimension fits usize");
    };
    out
}

fn to_u32(value: usize) -> u32 {
    let Ok(out) = u32::try_from(value) else {
        panic!("atlas count fits u32");
    };
    out
}
