// SPDX-FileCopyrightText: 2023 CELESTIFYX Team
// SPDX-License-Identifier: GPL-3.0-or-later

use color_quant::NeuQuant;

use std::{
    time::Instant,
    cmp::min,
    collections::HashMap,

    os::fd::{
        AsFd,
        BorrowedFd
    },

    io::{
        Write,
        stdout,
        Stdin
    }
};

use image::{
    DynamicImage,
    RgbaImage,
    FlatSamples,
    Rgba,
    SubImage,
    ImageBuffer,
    Rgb,
    GenericImageView,
    Pixel,

    imageops::{
        FilterType,
        colorops
    }
};

use rustix::{
    io::read,

    event::{
        PollFd,
        PollFlags,
        Timespec,
        poll
    },

    termios::{
        OptionalActions,
        Termios,
        Winsize,
        tcgetwinsize,
        tcsetattr
    }
};

use anyhow::{
    Context as _,
    Result
};

use super::ImageBackend;

use crate::image::{
    terminal::enable_raw_mode,
    config::DisplayConfig
};

pub(super) struct SixelBackend;

impl SixelBackend {
    pub(super) fn supported() -> Result<bool> {
        let stdin: Stdin = std::io::stdin();

        // save terminal attributes and disable canonical input processing mode
        let old_attributes: Termios = enable_raw_mode(&stdin)?;

        // ask for the primary device attribute string
        print!("\x1B[c");
        stdout().flush()?;

        let start_time: Instant = Instant::now();
        let stdin_fd: BorrowedFd = stdin.as_fd();

        let mut stdin_pollfd: [PollFd; 1] = [PollFd::new(&stdin_fd, PollFlags::IN)];
        let mut buf: Vec<u8> = Vec::<u8>::new();

        loop {
            // check for timeout while polling to avoid blocking the main thread
            while poll(&mut stdin_pollfd, Some(&Timespec::default()))? < 1 {
                if start_time.elapsed().as_millis() > 50 {
                    tcsetattr(stdin, OptionalActions::Now, &old_attributes).context("Failed to update terminal attributes")?;
                    return Ok(false);
                }
            }

            let mut byte: [u8; 1] = [0];
            read(&stdin, &mut byte)?;

            (&mut buf).push(byte[0]);

            if buf.starts_with(&[0x1B, b'[', b'?']) && buf.ends_with(b"c") {
                for attribute in buf[3..(buf.len() - 1)].split(|x: &u8| *x == b';') {
                    if attribute == [b'4'] {
                        tcsetattr(stdin, OptionalActions::Now, &old_attributes).context("Failed to update terminal attributes")?;
                        return Ok(true);
                    }
                }
            }
        }
    }
}

impl ImageBackend for SixelBackend {
    #[allow(clippy::map_entry)]
    fn add_image(&self, config: &DisplayConfig, image: &DynamicImage) -> Result<String> {
        let tty_size: Winsize = tcgetwinsize(std::io::stdin())?;

        let cw: u16 = tty_size.ws_xpixel / tty_size.ws_col;
        let lh: u16 = tty_size.ws_ypixel / tty_size.ws_row;

        let width_ratio:  f64 = 1.0 / cw as f64;
        let height_ratio: f64 = 1.0 / lh as f64;

        // resize image to fit the text height with the Lanczos3 algorithm
        let image: DynamicImage = image.resize(u32::MAX, (config.lines.len() as f64 / height_ratio) as u32, FilterType::Lanczos3);
        let rgba_image: RgbaImage = image.to_rgba8(); // convert the image to rgba samples

        let image_columns: f64 = width_ratio * image.width() as f64;
        let image_rows:    f64 = height_ratio * image.height() as f64;

        let flat_samples: FlatSamples<&[u8]> = rgba_image.as_flat_samples();
        let mut rgba_image: RgbaImage = rgba_image.clone();

        // reduce the amount of colors using dithering
        let pixels: &[u8] = flat_samples.image_slice().context("Error while slicing the image")?;
        colorops::dither(&mut rgba_image, &NeuQuant::new(10, config.color_resolution, pixels));

        let rgb_image: ImageBuffer<Rgb<u8>, Vec<u8>> = ImageBuffer::from_fn(rgba_image.width(), rgba_image.height(), |x: u32, y: u32| {
            let rgba_pixel: &Rgba<u8> = rgba_image.get_pixel(x, y);
            let mut rgb_pixel: Rgb<u8> = rgba_pixel.to_rgb();

            for subpixel in &mut rgb_pixel.0 {
                *subpixel = (*subpixel as f32 / 255.0 * rgba_pixel[3] as f32) as u8;
            }

            rgb_pixel
        });

        let mut image_data: Vec<u8> = Vec::<u8>::new();
        (&mut image_data).extend(b"\x1BPq"); // start sixel data
        (&mut image_data).extend(format!("\"1;1;{};{}", image.width(), image.height()).as_bytes());

        let mut colors: HashMap<Rgb<u8>, u8> = HashMap::<Rgb<u8>, u8>::new();

        // subtract 1 -> divide -> add 1 to round up the integer division
        for i in 0..((rgb_image.height() - 1) / 6 + 1) {
            let sixel_row: SubImage<&ImageBuffer<Rgb<u8>, Vec<u8>>> = rgb_image.view(0, i * 6, rgb_image.width(), min(6, rgb_image.height() - i * 6));

            for (_, _, pixel) in sixel_row.pixels() {
                if !colors.contains_key(&pixel) {
                    // sixel uses percentages for rgb values
                    let color_multiplier: f32 = 100.0 / 255.0;
                    (&mut image_data).extend(format!("#{};2;{};{};{}", colors.len(), (pixel[0] as f32 * color_multiplier) as u32, (pixel[1] as f32 * color_multiplier) as u32, (pixel[2] as f32 * color_multiplier) as u32).as_bytes());

                    let index: u8 = colors.len() as u8;
                    (&mut colors).insert(pixel, index);
                }
            }

            for (color, color_index) in &colors {
                let mut sixel_samples: Vec<u8> = vec![0; sixel_row.width() as usize];
                (&mut sixel_samples).resize(sixel_row.width() as usize, 0);

                for (x, y, pixel) in sixel_row.pixels() {
                    if color == &pixel {
                        sixel_samples[x as usize] |= 1 << y;
                    }
                }

                (&mut image_data).extend(format!("#{}", color_index).bytes());
                (&mut image_data).extend(sixel_samples.iter().map(|x: &u8| x + 0x3F));
                (&mut image_data).push(b'$');
            }

            (&mut image_data).push(b'-');
        }

        (&mut image_data).extend(b"\x1B\\");
        (&mut image_data).extend(format!("\x1B[{}A", image_rows as u32 - 1).as_bytes()); // move cursor to top-left corner
        (&mut image_data).extend(format!("\x1B[{}C", image_columns as u32 + 1).as_bytes()); // move cursor to top-right corner of image

        let mut i: usize = 0;

        for line in &config.lines {
            (&mut image_data).extend(format!("\x1B[s{}{}\x1B[u\x1B[1B", " ".repeat(config.margin), line).as_bytes());
            i += 1;
        }

        (&mut image_data).extend(format!("\n\x1B[{}B", config.lines.len().max(image_rows as usize) - i).as_bytes()); // move cursor to end of image
        Ok(String::from_utf8(image_data)?)
    }
}
