use super::color::{
    GENERATED_ASSET_SHEET_GREEN_SCREEN_MIN_SCORE, GENERATED_ASSET_SHEET_GREEN_SCREEN_SOFT_SCORE,
    GENERATED_ASSET_SHEET_HARD_GREEN_SCREEN_SCORE, clamp_generated_asset_sheet_unit,
    compute_generated_asset_sheet_green_screen_score,
    compute_generated_asset_sheet_key_color_score,
    compute_generated_asset_sheet_white_screen_score,
    is_generated_asset_sheet_soft_green_matte_pixel, lerp_generated_asset_sheet_channel,
    touches_generated_asset_sheet_background_mask,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct GeneratedAssetSheetKeyColor {
    pub red: u8,
    pub green: u8,
    pub blue: u8,
}

impl GeneratedAssetSheetKeyColor {
    pub const GREEN_SCREEN: Self = Self {
        red: 0,
        green: 255,
        blue: 0,
    };

    pub const MAGENTA_SCREEN: Self = Self {
        red: 255,
        green: 0,
        blue: 255,
    };

    pub fn is_green_screen(self) -> bool {
        self == Self::GREEN_SCREEN
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct GeneratedAssetSheetAlphaOptions {
    pub key_color: GeneratedAssetSheetKeyColor,
    pub remove_near_white_background: bool,
    pub remove_disconnected_hard_key_background: bool,
}

impl GeneratedAssetSheetAlphaOptions {
    pub const fn green_screen() -> Self {
        Self {
            key_color: GeneratedAssetSheetKeyColor::GREEN_SCREEN,
            remove_near_white_background: true,
            remove_disconnected_hard_key_background: true,
        }
    }

    pub const fn jump_hop_magenta_screen() -> Self {
        Self {
            key_color: GeneratedAssetSheetKeyColor::MAGENTA_SCREEN,
            remove_near_white_background: false,
            remove_disconnected_hard_key_background: false,
        }
    }
}

impl Default for GeneratedAssetSheetAlphaOptions {
    fn default() -> Self {
        Self::green_screen()
    }
}

pub fn apply_generated_asset_sheet_green_screen_alpha(
    source: image::DynamicImage,
) -> image::DynamicImage {
    apply_generated_asset_sheet_alpha_with_options(
        source,
        GeneratedAssetSheetAlphaOptions::default(),
    )
}

pub fn apply_generated_asset_sheet_alpha_with_options(
    source: image::DynamicImage,
    options: GeneratedAssetSheetAlphaOptions,
) -> image::DynamicImage {
    let mut image = source.to_rgba8();
    let (width, height) = image.dimensions();
    remove_generated_asset_sheet_green_screen_background(
        image.as_mut(),
        width as usize,
        height as usize,
        options,
    );
    image::DynamicImage::ImageRgba8(image)
}

fn remove_generated_asset_sheet_green_screen_background(
    pixels: &mut [u8],
    width: usize,
    height: usize,
    options: GeneratedAssetSheetAlphaOptions,
) -> bool {
    let pixel_count = width.saturating_mul(height);
    if pixel_count == 0 || pixels.len() < pixel_count.saturating_mul(4) {
        return false;
    }

    let mut key_scores = vec![0.0f32; pixel_count];
    let mut white_scores = vec![0.0f32; pixel_count];
    let mut background_hints = vec![0.0f32; pixel_count];
    let mut background_mask = vec![0u8; pixel_count];
    let mut queue = Vec::<usize>::new();
    let mut queue_index = 0usize;

    for pixel_index in 0..pixel_count {
        let offset = pixel_index * 4;
        let red = pixels[offset];
        let green = pixels[offset + 1];
        let blue = pixels[offset + 2];
        let alpha = pixels[offset + 3];
        let key_score =
            compute_generated_asset_sheet_key_score([red, green, blue, alpha], options.key_color);
        let white_score = if options.remove_near_white_background {
            compute_generated_asset_sheet_white_screen_score([red, green, blue, alpha])
        } else {
            0.0
        };
        let transparency_hint =
            clamp_generated_asset_sheet_unit((56.0 - alpha as f32) / 56.0) * 0.75;

        key_scores[pixel_index] = key_score;
        white_scores[pixel_index] = white_score;
        background_hints[pixel_index] = key_score.max(white_score).max(transparency_hint);
    }

    let seed_background_pixel =
        |pixel_index: usize, background_mask: &mut [u8], queue: &mut Vec<usize>| {
            if background_mask[pixel_index] != 0 {
                return;
            }
            let alpha = pixels[pixel_index * 4 + 3];
            let strong_candidate = alpha < 40
                || key_scores[pixel_index] >= GENERATED_ASSET_SHEET_HARD_GREEN_SCREEN_SCORE
                || (alpha < 224
                    && key_scores[pixel_index] > GENERATED_ASSET_SHEET_GREEN_SCREEN_MIN_SCORE)
                || (options.remove_near_white_background && white_scores[pixel_index] > 0.32);
            if !strong_candidate {
                return;
            }
            background_mask[pixel_index] = 1;
            queue.push(pixel_index);
        };

    for x in 0..width {
        seed_background_pixel(x, &mut background_mask, &mut queue);
        seed_background_pixel((height - 1) * width + x, &mut background_mask, &mut queue);
    }
    for y in 1..height.saturating_sub(1) {
        seed_background_pixel(y * width, &mut background_mask, &mut queue);
        seed_background_pixel(y * width + width - 1, &mut background_mask, &mut queue);
    }

    while queue_index < queue.len() {
        let pixel_index = queue[queue_index];
        queue_index += 1;

        let x = pixel_index % width;
        let y = pixel_index / width;
        let neighbor_indexes = [
            if x > 0 { Some(pixel_index - 1) } else { None },
            if x + 1 < width {
                Some(pixel_index + 1)
            } else {
                None
            },
            if y > 0 {
                Some(pixel_index - width)
            } else {
                None
            },
            if y + 1 < height {
                Some(pixel_index + width)
            } else {
                None
            },
        ];

        for next_pixel_index in neighbor_indexes.into_iter().flatten() {
            if background_mask[next_pixel_index] != 0 {
                continue;
            }
            let next_offset = next_pixel_index * 4;
            let alpha = pixels[next_offset + 3];
            let key_score = key_scores[next_pixel_index];
            let white_score = white_scores[next_pixel_index];
            let hint = background_hints[next_pixel_index];
            let reachable_soft_edge = hint > 0.08
                && alpha < 224
                && (key_score > 0.04
                    || (options.remove_near_white_background && white_score > 0.08)
                    || alpha < 180);
            let key_background = key_score >= GENERATED_ASSET_SHEET_HARD_GREEN_SCREEN_SCORE
                || (alpha < 224 && key_score > GENERATED_ASSET_SHEET_GREEN_SCREEN_MIN_SCORE);
            if alpha < 40
                || key_background
                || (options.remove_near_white_background && white_score > 0.32)
                || reachable_soft_edge
            {
                background_mask[next_pixel_index] = 1;
                queue.push(next_pixel_index);
            }
        }
    }

    if options.remove_disconnected_hard_key_background {
        for pixel_index in 0..pixel_count {
            if background_mask[pixel_index] == 0
                && key_scores[pixel_index] >= GENERATED_ASSET_SHEET_HARD_GREEN_SCREEN_SCORE
            {
                background_mask[pixel_index] = 1;
            }
        }
    }

    let soft_green_cleanup_rounds = (width.min(height) / 40).clamp(4, 14);
    for _ in 0..soft_green_cleanup_rounds {
        let mut expanded_mask = background_mask.clone();
        let mut changed_this_round = false;
        for y in 0..height {
            for x in 0..width {
                let pixel_index = y * width + x;
                if background_mask[pixel_index] != 0 {
                    continue;
                }
                let offset = pixel_index * 4;
                let pixel = [
                    pixels[offset],
                    pixels[offset + 1],
                    pixels[offset + 2],
                    pixels[offset + 3],
                ];
                let key_score = key_scores[pixel_index];
                let white_score = white_scores[pixel_index];
                if !is_generated_asset_sheet_soft_key_matte_pixel(
                    pixel,
                    key_score,
                    white_score,
                    options,
                ) {
                    continue;
                }
                if !touches_generated_asset_sheet_background_mask(
                    x,
                    y,
                    width,
                    height,
                    &background_mask,
                ) {
                    continue;
                }

                expanded_mask[pixel_index] = 1;
                changed_this_round = true;
            }
        }
        background_mask = expanded_mask;
        if !changed_this_round {
            break;
        }
    }

    for _ in 0..2 {
        let mut expanded_mask = background_mask.clone();
        for y in 0..height {
            for x in 0..width {
                let pixel_index = y * width + x;
                if background_mask[pixel_index] != 0 {
                    continue;
                }
                let alpha = pixels[pixel_index * 4 + 3];
                let key_score = key_scores[pixel_index];
                let white_score = white_scores[pixel_index];
                let hint = background_hints[pixel_index];
                let soft_matte_candidate = alpha < 224
                    || (options.remove_near_white_background && white_score > 0.10)
                    || key_score >= GENERATED_ASSET_SHEET_HARD_GREEN_SCREEN_SCORE;
                if hint < GENERATED_ASSET_SHEET_GREEN_SCREEN_SOFT_SCORE || !soft_matte_candidate {
                    continue;
                }

                let mut adjacent_background_count = 0usize;
                for offset_y in -1i32..=1 {
                    for offset_x in -1i32..=1 {
                        if offset_x == 0 && offset_y == 0 {
                            continue;
                        }
                        let next_x = x as i32 + offset_x;
                        let next_y = y as i32 + offset_y;
                        if next_x < 0
                            || next_x >= width as i32
                            || next_y < 0
                            || next_y >= height as i32
                        {
                            adjacent_background_count += 1;
                            continue;
                        }
                        if background_mask[next_y as usize * width + next_x as usize] != 0 {
                            adjacent_background_count += 1;
                        }
                    }
                }

                if adjacent_background_count >= 2
                    || (adjacent_background_count >= 1
                        && hint >= GENERATED_ASSET_SHEET_GREEN_SCREEN_MIN_SCORE)
                {
                    expanded_mask[pixel_index] = 1;
                }
            }
        }
        background_mask = expanded_mask;
    }

    let mut changed = false;
    for pixel_index in 0..pixel_count {
        if background_mask[pixel_index] == 0 {
            continue;
        }
        let alpha_offset = pixel_index * 4 + 3;
        if pixels[alpha_offset] != 0 {
            pixels[alpha_offset] = 0;
            changed = true;
        }
    }

    for y in 0..height {
        for x in 0..width {
            let pixel_index = y * width + x;
            let offset = pixel_index * 4;
            let alpha = pixels[offset + 3];
            if alpha == 0 {
                continue;
            }

            let mut touches_transparent_edge = false;
            for offset_y in -1i32..=1 {
                for offset_x in -1i32..=1 {
                    if offset_x == 0 && offset_y == 0 {
                        continue;
                    }
                    let next_x = x as i32 + offset_x;
                    let next_y = y as i32 + offset_y;
                    if next_x < 0 || next_x >= width as i32 || next_y < 0 || next_y >= height as i32
                    {
                        touches_transparent_edge = true;
                        continue;
                    }
                    let next_pixel_index = next_y as usize * width + next_x as usize;
                    if background_mask[next_pixel_index] != 0
                        || pixels[next_pixel_index * 4 + 3] < 16
                    {
                        touches_transparent_edge = true;
                    }
                }
            }

            if !touches_transparent_edge {
                continue;
            }

            let key_score = key_scores[pixel_index];
            let white_score = white_scores[pixel_index];
            let contamination = key_score.max(white_score).max(if alpha < 220 {
                ((220 - alpha) as f32 / 220.0) * 0.25
            } else {
                0.0
            });
            if contamination < 0.06 {
                continue;
            }

            let sample = collect_generated_asset_sheet_foreground_neighbor_color(
                pixels,
                width,
                height,
                x,
                y,
                &background_mask,
                &background_hints,
            );
            let mut red = pixels[offset] as f32;
            let mut green = pixels[offset + 1] as f32;
            let mut blue = pixels[offset + 2] as f32;
            let blend = if options.key_color.is_green_screen() {
                clamp_generated_asset_sheet_unit(contamination.max(0.22))
            } else {
                // 中文注释：洋红 / 青色等非绿幕 key 的残留更容易表现成彩边，
                // 需要比绿幕更强地向主体邻近色收敛，避免 PNG 边缘继续带 key 色。
                clamp_generated_asset_sheet_unit((key_score * 1.35).max(contamination).max(0.28))
            };

            if let Some((sample_red, sample_green, sample_blue)) = sample {
                red = lerp_generated_asset_sheet_channel(red, sample_red as f32, blend);
                green = lerp_generated_asset_sheet_channel(green, sample_green as f32, blend);
                blue = lerp_generated_asset_sheet_channel(blue, sample_blue as f32, blend);

                if options.key_color.is_green_screen() && key_score > 0.04 {
                    green = green.min(sample_green as f32 + 18.0);
                }
                if options.remove_near_white_background && white_score > 0.1 {
                    red = red.min(sample_red as f32 + 26.0);
                    green = green.min(sample_green as f32 + 26.0);
                    blue = blue.min(sample_blue as f32 + 26.0);
                }
                if !options.key_color.is_green_screen() && key_score > 0.04 {
                    let defringed = suppress_generated_asset_sheet_key_color_fringe(
                        [red, green, blue],
                        [sample_red as f32, sample_green as f32, sample_blue as f32],
                        key_score,
                        options.key_color,
                    );
                    red = defringed[0];
                    green = defringed[1];
                    blue = defringed[2];
                }
            } else {
                if options.key_color.is_green_screen() && key_score > 0.04 {
                    let toned_green = (green - (green - red.max(blue)) * 0.78)
                        .round()
                        .max(red.max(blue));
                    green = green.min(toned_green).min(red.max(blue) + 18.0);
                }

                if options.remove_near_white_background && white_score > 0.12 {
                    let spread = red.max(green).max(blue) - red.min(green).min(blue);
                    if spread < 20.0 {
                        let toned_value = ((red + green + blue) / 3.0 * 0.88).round();
                        red = red.min(toned_value);
                        green = green.min(toned_value);
                        blue = blue.min(toned_value);
                    }
                }
                if !options.key_color.is_green_screen() && key_score > 0.04 {
                    let neutral = (red + green + blue) / 3.0;
                    let defringed = suppress_generated_asset_sheet_key_color_fringe(
                        [red, green, blue],
                        [neutral, neutral, neutral],
                        key_score,
                        options.key_color,
                    );
                    red = defringed[0];
                    green = defringed[1];
                    blue = defringed[2];
                }
            }

            let mut next_alpha = alpha;
            let edge_fade = if options.key_color.is_green_screen() {
                (key_score * 0.35).max(white_score * 0.28)
            } else {
                (key_score * 0.48).max(white_score * 0.28)
            };
            if edge_fade > 0.08 {
                next_alpha = ((alpha as f32) * (1.0 - edge_fade)).round() as u8;
                if next_alpha < 10 {
                    next_alpha = 0;
                }
            }

            let next_red = red.round().clamp(0.0, 255.0) as u8;
            let next_green = green.round().clamp(0.0, 255.0) as u8;
            let next_blue = blue.round().clamp(0.0, 255.0) as u8;
            if next_red != pixels[offset]
                || next_green != pixels[offset + 1]
                || next_blue != pixels[offset + 2]
                || next_alpha != alpha
            {
                pixels[offset] = next_red;
                pixels[offset + 1] = next_green;
                pixels[offset + 2] = next_blue;
                pixels[offset + 3] = next_alpha;
                changed = true;
            }
        }
    }

    changed
}

pub(super) fn suppress_generated_asset_sheet_key_color_fringe(
    color: [f32; 3],
    target: [f32; 3],
    key_score: f32,
    key_color: GeneratedAssetSheetKeyColor,
) -> [f32; 3] {
    let strength = clamp_generated_asset_sheet_unit(key_score * 1.18);
    let key_channels = [
        key_color.red as f32 / 255.0,
        key_color.green as f32 / 255.0,
        key_color.blue as f32 / 255.0,
    ];
    let mut next = color;

    for index in 0..3 {
        if key_channels[index] >= 0.66 {
            let cap = target[index] + 18.0 + (1.0 - strength) * 28.0;
            next[index] = next[index].min(lerp_generated_asset_sheet_channel(
                next[index],
                cap,
                strength,
            ));
        } else if key_channels[index] <= 0.34 {
            next[index] =
                lerp_generated_asset_sheet_channel(next[index], target[index], strength * 0.72);
        }
    }

    next
}

fn compute_generated_asset_sheet_key_score(
    pixel: [u8; 4],
    key_color: GeneratedAssetSheetKeyColor,
) -> f32 {
    if key_color.is_green_screen() {
        return compute_generated_asset_sheet_green_screen_score(pixel);
    }

    compute_generated_asset_sheet_key_color_score(
        pixel,
        [key_color.red, key_color.green, key_color.blue],
    )
}

fn is_generated_asset_sheet_soft_key_matte_pixel(
    pixel: [u8; 4],
    key_score: f32,
    white_score: f32,
    options: GeneratedAssetSheetAlphaOptions,
) -> bool {
    if options.key_color.is_green_screen() {
        return is_generated_asset_sheet_soft_green_matte_pixel(pixel, key_score, white_score);
    }

    pixel[3] != 0
        && key_score >= GENERATED_ASSET_SHEET_GREEN_SCREEN_MIN_SCORE
        && (!options.remove_near_white_background || white_score < 0.34)
}

fn collect_generated_asset_sheet_foreground_neighbor_color(
    pixels: &[u8],
    width: usize,
    height: usize,
    x: usize,
    y: usize,
    background_mask: &[u8],
    background_hints: &[f32],
) -> Option<(u8, u8, u8)> {
    let mut total_weight = 0.0f32;
    let mut total_red = 0.0f32;
    let mut total_green = 0.0f32;
    let mut total_blue = 0.0f32;

    for offset_y in -2i32..=2 {
        for offset_x in -2i32..=2 {
            if offset_x == 0 && offset_y == 0 {
                continue;
            }
            let next_x = x as i32 + offset_x;
            let next_y = y as i32 + offset_y;
            if next_x < 0 || next_x >= width as i32 || next_y < 0 || next_y >= height as i32 {
                continue;
            }

            let next_pixel_index = next_y as usize * width + next_x as usize;
            if background_mask[next_pixel_index] != 0 || background_hints[next_pixel_index] >= 0.18
            {
                continue;
            }

            let next_offset = next_pixel_index * 4;
            let next_alpha = pixels[next_offset + 3];
            if next_alpha < 96 {
                continue;
            }
            let distance = offset_x.unsigned_abs() + offset_y.unsigned_abs();
            let weight = (next_alpha as f32 / 255.0)
                * if distance <= 1 {
                    1.8
                } else if distance == 2 {
                    1.2
                } else {
                    0.7
                };

            total_weight += weight;
            total_red += pixels[next_offset] as f32 * weight;
            total_green += pixels[next_offset + 1] as f32 * weight;
            total_blue += pixels[next_offset + 2] as f32 * weight;
        }
    }

    if total_weight <= 0.0 {
        return None;
    }

    Some((
        (total_red / total_weight).round() as u8,
        (total_green / total_weight).round() as u8,
        (total_blue / total_weight).round() as u8,
    ))
}