hinoirisetr 1.6.2

use rustix::event::EventfdFlags;
use rustix::event::eventfd;
use rustix::event::{PollFd, PollFlags, poll};
use rustix::fd::AsFd;
use std::collections::HashMap;
use std::fs::File;
use std::num::NonZero;
use std::sync::mpsc::{self, Sender};
use std::thread;
use std::thread::JoinHandle;
use wayland_client::protocol::wl_output;
use wayland_client::{Connection, Dispatch, QueueHandle, protocol::wl_registry};
use wayland_protocols_wlr::gamma_control::v1::client::{
    zwlr_gamma_control_manager_v1, zwlr_gamma_control_v1,
};

use crate::debug;

#[derive(Debug)]
pub struct WlrGammaBackend {
    tx: Sender<Cmd>,
    wake_fd: rustix::fd::OwnedFd,
    handle: Option<JoinHandle<()>>,
}

impl WlrGammaBackend {
    /// spawn the backend thread and return a handle
    pub fn new() -> Self {
        let wake_fd = eventfd(0, EventfdFlags::NONBLOCK).expect("eventfd");
        let wake_fd_thread = wake_fd.try_clone().expect("eventfd clone");
        let (tx, rx) = mpsc::channel::<Cmd>();

        let handle = thread::spawn(move || run(rx, wake_fd_thread));

        Self {
            tx,
            wake_fd,
            handle: Some(handle),
        }
    }

    pub fn set_gamma(&self, gamma: u8) {
        self.send(Cmd::SetGamma(gamma));
    }

    pub fn set_temperature(&self, temp: u32) {
        self.send(Cmd::SetTemperature(temp));
    }

    pub fn stop(&mut self) {
        self.send(Cmd::Quit);
        if let Some(handle) = self.handle.take() {
            handle.join().ok();
        }
    }

    fn send(&self, cmd: Cmd) {
        self.tx.send(cmd).expect("backend thread died");
        // wake the poll() in the backend thread
        rustix::io::write(&self.wake_fd, &1u64.to_ne_bytes()).ok();
    }
}

impl Default for WlrGammaBackend {
    fn default() -> Self {
        Self::new()
    }
}

enum Cmd {
    SetGamma(u8),
    SetTemperature(u32),
    Quit,
}

struct State {
    outputs: HashMap<u32, OutputData>,
    gamma_control_manager: Option<zwlr_gamma_control_manager_v1::ZwlrGammaControlManagerV1>,
    current_gamma: u8,
    current_temp: u32,
}

struct OutputData {
    output: wl_output::WlOutput,
    gamma_control: Option<zwlr_gamma_control_v1::ZwlrGammaControlV1>,
    ramp_size: Option<NonZero<usize>>,
}

impl State {
    fn set_gamma(&mut self, gamma: u8) {
        self.current_gamma = gamma.clamp(1, 100);
    }

    fn set_temperature(&mut self, temp: u32) {
        self.current_temp = temp.clamp(1000, 40000);
    }

    fn apply_changes(&mut self) {
        for data in self.outputs.values() {
            if let (Some(control), Some(ramp_size)) = (&data.gamma_control, data.ramp_size) {
                let file = create_gamma_fd(ramp_size.into(), self.current_gamma, self.current_temp);
                control.set_gamma(file.as_fd());
            }
        }
    }
}

// backend thread
fn run(rx: mpsc::Receiver<Cmd>, wake_fd: rustix::fd::OwnedFd) {
    let conn = Connection::connect_to_env().expect("wayland connection");
    let display = conn.display();
    let mut event_queue = conn.new_event_queue();
    let qh = event_queue.handle();

    let _registry = display.get_registry(&qh, ());

    let mut state = State {
        outputs: HashMap::new(),
        gamma_control_manager: None,
        current_gamma: 100,
        current_temp: 6500,
    };

    // find outputs and manager
    event_queue.roundtrip(&mut state).expect("roundtrip");

    // main loop: wait in poll() until wayland or wake_fd fires
    loop {
        // reset the wake_fd
        let mut buf = [0u8; 8];
        let _ = rustix::io::read(&wake_fd, &mut buf);
        // get all pending commands
        while let Ok(cmd) = rx.try_recv() {
            match cmd {
                Cmd::SetGamma(g) => state.set_gamma(g),
                Cmd::SetTemperature(t) => state.set_temperature(t),
                Cmd::Quit => return,
            }
            state.apply_changes();
        }

        // flush any pending outgoing messages
        event_queue.flush().expect("flush");

        // block until wayland socket or wake_fd has data
        let guard = if let Some(g) = event_queue.prepare_read() {
            g
        } else {
            event_queue.dispatch_pending(&mut state).expect("dispatch");
            continue;
        };

        let wayland_pollfd = PollFd::new(&conn, PollFlags::IN);
        let wake_pollfd = PollFd::new(&wake_fd, PollFlags::IN);
        poll(&mut [wayland_pollfd, wake_pollfd], None).expect("poll");

        guard.read().ok();
        event_queue.dispatch_pending(&mut state).expect("dispatch");
    }
}

// wayland dispatch impls

impl Dispatch<wl_registry::WlRegistry, ()> for State {
    fn event(
        state: &mut Self,
        registry: &wl_registry::WlRegistry,
        event: wl_registry::Event,
        (): &(),
        _: &Connection,
        qh: &QueueHandle<Self>,
    ) {
        match event {
            wl_registry::Event::Global {
                name,
                interface,
                version,
            } => match interface.as_str() {
                "wl_output" => {
                    let output = registry.bind::<wl_output::WlOutput, _, _>(name, version, qh, ());
                    let mut control = None;
                    if let Some(manager) = &state.gamma_control_manager {
                        control = Some(manager.get_gamma_control(&output, qh, name));
                    }

                    let outputs_data = OutputData {
                        output,
                        gamma_control: control,
                        ramp_size: None,
                    };

                    debug!("Monitor attached: global name {name}");

                    state.outputs.insert(name, outputs_data);
                    state.apply_changes();
                },
                "zwlr_gamma_control_manager_v1" => {
                    let manager = registry
                        .bind::<zwlr_gamma_control_manager_v1::ZwlrGammaControlManagerV1, _, _>(
                            name,
                            version,
                            qh,
                            (),
                        );
                    for (output_name, output) in &mut state.outputs {
                        let control = manager.get_gamma_control(&output.output, qh, *output_name);

                        output.gamma_control = Some(control);
                    }

                    state.gamma_control_manager = Some(manager);
                    state.apply_changes();
                },
                _ => {},
            },
            wl_registry::Event::GlobalRemove { name } => {
                if let Some(removed_output) = state.outputs.remove(&name) {
                    if let Some(control) = removed_output.gamma_control {
                        control.destroy();
                    }
                    removed_output.output.release();
                    debug!("Monitor removed: global name {name}");
                }
            },
            _ => {},
        }
    }
}

impl Dispatch<zwlr_gamma_control_v1::ZwlrGammaControlV1, u32> for State {
    fn event(
        state: &mut Self,
        _proxy: &zwlr_gamma_control_v1::ZwlrGammaControlV1,
        event: zwlr_gamma_control_v1::Event,
        name: &u32,
        _: &Connection,
        _: &QueueHandle<Self>,
    ) {
        match event {
            zwlr_gamma_control_v1::Event::GammaSize { size } => {
                if let Some(output_data) = state.outputs.get_mut(name) {
                    output_data.ramp_size = Some(NonZero::new(size as usize).unwrap());
                }
                state.apply_changes();
            },
            zwlr_gamma_control_v1::Event::Failed => {
                // warn!("gamma control failed for {proxy:?}");
            },
            _ => {},
        }
    }
}

impl Dispatch<zwlr_gamma_control_manager_v1::ZwlrGammaControlManagerV1, ()> for State {
    fn event(
        _: &mut Self,
        _: &zwlr_gamma_control_manager_v1::ZwlrGammaControlManagerV1,
        _: zwlr_gamma_control_manager_v1::Event,
        (): &(),
        _: &Connection,
        _: &QueueHandle<Self>,
    ) {
    }
}

impl Dispatch<wl_output::WlOutput, ()> for State {
    fn event(
        _: &mut Self,
        _: &wl_output::WlOutput,
        _: wl_output::Event,
        (): &(),
        _: &Connection,
        _: &QueueHandle<Self>,
    ) {
    }
}

// gamma math

/// Tanner Helland empirical curve fit: Kelvin -> linear RGB multipliers (0.0–1.0)
/// https://tannerhelland.com/2012/09/18/convert-temperature-rgb-algorithm-code.html
fn kelvin_to_rgb(temp: u32) -> (f64, f64, f64) {
    let t = f64::from(temp.clamp(1000, 40000) / 100);

    let r = if t <= 66.0 {
        1.0
    } else {
        (329.698727446 * (t - 60.0).powf(-0.1332047592) / 255.0).clamp(0.0, 1.0)
    };

    let g = if t <= 66.0 {
        99.4708025861f64.mul_add(t.ln(), -161.1195681661) / 255.0
    } else {
        288.1221695283 * (t - 60.0).powf(-0.0755148492) / 255.0
    }
    .clamp(0.0, 1.0);

    let b = if t >= 66.0 {
        1.0
    } else if t <= 19.0 {
        0.0
    } else {
        138.5177312231f64.mul_add((t - 10.0).ln(), -305.0447927307) / 255.0
    }
    .clamp(0.0, 1.0);

    (r, g, b)
}

fn generate_ramp_channel(ramp_size: usize, gamma: u8, color_mul: f64) -> Vec<u16> {
    let brightness = f64::from(gamma.clamp(1, 100)) / 100.0;
    let max_index = (ramp_size - 1).max(1) as f64;

    (0..ramp_size)
        .map(|i| {
            let t = i as f64 / max_index;
            let val = (t * brightness * color_mul).mul_add(65535.0, 0.5);
            val.clamp(0.0, 65535.0) as u16
        })
        .collect()
}

fn create_gamma_fd(ramp_size: usize, gamma: u8, temp: u32) -> File {
    let total_bytes = 3 * ramp_size * 2;

    let mfd = memfd::MemfdOptions::default()
        .create("gamma")
        .expect("memfd");
    let file = mfd.into_file();
    file.set_len(total_bytes as u64).unwrap();

    let mut buffer = unsafe {
        memmap2::MmapOptions::new()
            .len(total_bytes)
            .map_mut(&file)
            .expect("mmap")
    };

    let buffer_u16: &mut [u16] =
        unsafe { std::slice::from_raw_parts_mut(buffer.as_mut_ptr() as *mut u16, total_bytes / 2) };

    let (r_mul, g_mul, b_mul) = kelvin_to_rgb(temp);
    let red = generate_ramp_channel(ramp_size, gamma, r_mul);
    let green = generate_ramp_channel(ramp_size, gamma, g_mul);
    let blue = generate_ramp_channel(ramp_size, gamma, b_mul);

    buffer_u16[..ramp_size].copy_from_slice(&red);
    buffer_u16[ramp_size..2 * ramp_size].copy_from_slice(&green);
    buffer_u16[2 * ramp_size..].copy_from_slice(&blue);

    buffer.flush().unwrap();
    file
}