// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (C) 2026, Nathan Gill

use std::{
    collections::HashMap,
    fs::File,
    io::Seek,
    os::fd::AsRawFd,
    sync::{Arc, atomic::AtomicBool},
};

use anyhow::{Result, anyhow, bail};
use cairo::ImageSurface;
use tracing::{debug, trace, warn};
use wayland_client::protocol::{wl_region, wl_subcompositor, wl_subsurface};
use wayland_client::{
    Connection, Dispatch, QueueHandle, delegate_noop,
    protocol::{
        wl_callback, wl_compositor, wl_display, wl_output, wl_registry, wl_seat, wl_shm,
        wl_shm_pool, wl_surface,
    },
};
use wayland_protocols::ext::session_lock::v1::client::{
    ext_session_lock_manager_v1, ext_session_lock_v1,
};
use zeroize::Zeroizing;

#[cfg(feature = "gdk-pixbuf")]
use {crate::cairo_ext::ImageSurfaceExt, gdk_pixbuf::Pixbuf};

use crate::{
    auth::AuthChannel,
    cairo_ext::SubpixelOrderExt,
    comm::PipeCommChannel,
    event_loop::{EventSource, NLockEventLoop},
    signal::install_debug_handler,
};
use crate::{
    auth::{AtomicAuthState, AuthState},
    util::detect_png,
};
use crate::{config::NLockConfig, event::EventType, util::BackgroundType};
use crate::{
    seat::{NLockSeat, NLockXkb},
    surface::NLockSurface,
};

pub struct NLockStateArgs {
    pub config: NLockConfig,
    pub display: wl_display::WlDisplay,
    pub auth_comm: Arc<AuthChannel>,
    pub debug: bool,
}

pub struct NLockState {
    pub config: NLockConfig,
    pub running: Arc<AtomicBool>,
    pub locked: bool,
    pub unlocked: bool,
    pub state_changed: Arc<AtomicBool>,
    pub display: wl_display::WlDisplay,
    pub registry: Option<wl_registry::WlRegistry>,
    pub compositor: Option<wl_compositor::WlCompositor>,
    pub subcompositor: Option<wl_subcompositor::WlSubcompositor>,
    pub shm: Option<wl_shm::WlShm>,
    pub r_seat: Option<wl_seat::WlSeat>,
    pub session_lock_manager: Option<ext_session_lock_manager_v1::ExtSessionLockManagerV1>,
    pub session_lock: Option<ext_session_lock_v1::ExtSessionLockV1>,
    pub surfaces: HashMap<u32, NLockSurface>,
    pub seat: NLockSeat,
    pub xkb: NLockXkb,
    pub password: Zeroizing<String>,
    pub auth_comm: Arc<AuthChannel>,
    pub auth_state: Arc<AtomicAuthState>,
    pub background_image: Option<cairo::ImageSurface>,
    pub event_loop: NLockEventLoop,
    pub debug_comm: Option<Arc<PipeCommChannel<()>>>,
    pub interrupt: PipeCommChannel<()>,
}

impl NLockState {
    pub fn new(args: NLockStateArgs) -> Result<Self> {
        let debug_comm = if args.debug {
            Some(Arc::new(PipeCommChannel::new()?))
        } else {
            None
        };

        let interrupt = PipeCommChannel::new()?;

        let mut s = Self {
            config: args.config,
            running: Arc::new(AtomicBool::new(true)),
            locked: false,
            unlocked: false,
            state_changed: Arc::new(AtomicBool::new(false)),
            display: args.display,
            registry: None,
            compositor: None,
            subcompositor: None,
            shm: None,
            r_seat: None,
            session_lock_manager: None,
            session_lock: None,
            surfaces: HashMap::new(),
            seat: NLockSeat::default(),
            xkb: NLockXkb::default(),
            password: Zeroizing::new("".to_string()),
            auth_comm: args.auth_comm,
            auth_state: Arc::new(AtomicAuthState::new(AuthState::Idle)),
            background_image: None,
            event_loop: NLockEventLoop::default(),
            debug_comm,
            interrupt,
        };

        if let Err(e) = s.try_load_background_image() {
            bail!(
                "Failed to load background image: {}: {}",
                s.config.image.path.display(),
                e
            );
        }

        s.event_loop.add(
            EventSource::Fd(s.auth_comm.response.rx().as_raw_fd()),
            EventType::AuthStateChanged.into(),
        )?;

        s.event_loop.add(
            EventSource::Fd(s.interrupt.rx().as_raw_fd()),
            EventType::Interrupt.into(),
        )?;

        if args.debug
            && let Some(debug_comm) = &s.debug_comm
        {
            s.event_loop.add(
                EventSource::Fd(debug_comm.rx().as_raw_fd()),
                EventType::Debug.into(),
            )?;

            // install a handler on SIGUSR1
            install_debug_handler(debug_comm.clone())?;
        }

        Ok(s)
    }

    pub fn get_registry(&mut self, qh: &QueueHandle<Self>) {
        let registry = self.display.get_registry(qh, ());
        self.registry = Some(registry);
    }

    pub fn lock(&mut self, qh: &QueueHandle<Self>) {
        if let Some(session_lock_manager) = &self.session_lock_manager {
            let session_lock = session_lock_manager.lock(qh, ());
            self.session_lock = Some(session_lock);
        }
    }

    pub fn unlock(&mut self) {
        if let Some(session_lock) = &self.session_lock {
            if self.locked {
                session_lock.unlock_and_destroy();
            } else {
                session_lock.destroy();
            }

            // free any held surfaces
            self.surfaces.clear();

            self.session_lock = None;
            self.locked = false;
            self.unlocked = true;

            self.clear_password();

            debug!("Session is unlocked");
        }
    }

    pub fn clear_password(&mut self) {
        self.password.clear();
    }

    /// Write the current password into the auth channel and clear it
    pub fn submit_password(&mut self) {
        if let Err(e) = self.auth_comm.request.write(self.password.to_string()) {
            warn!("Failed to write auth request: {e}");
        }

        self.clear_password();
    }

    fn try_load_background_image(&mut self) -> Result<()> {
        if self.config.general.bg_type == BackgroundType::Color {
            self.config.general.bg_type = BackgroundType::Color;
            return Ok(());
        }

        let mut image_file = File::open(&self.config.image.path)?;
        let is_png = detect_png(&mut image_file)?;
        image_file.rewind()?;

        // if it's a PNG, Cairo can handle it directly, pixbuf conversion is expensive
        if is_png {
            let image_surface = ImageSurface::create_from_png(&mut image_file)?;
            self.background_image = Some(image_surface);
        } else {
            #[cfg(feature = "gdk-pixbuf")]
            {
                let pixbuf = Pixbuf::from_read(image_file)?;
                let pixbuf = pixbuf
                    .apply_embedded_orientation()
                    .ok_or(anyhow!("Failed to apply embedded image orientation"))?;

                let image_surface = ImageSurface::create_from_pixbuf(&pixbuf)?;
                self.background_image = Some(image_surface);
            }

            #[cfg(not(feature = "gdk-pixbuf"))]
            {
                return Err(anyhow!(
                    "Non-PNG formats require GDK-Pixbuf support, but nlock was compiled without it"
                ));
            }
        }

        self.config.general.bg_type = BackgroundType::Image;

        Ok(())
    }
}

impl Dispatch<wl_registry::WlRegistry, ()> for NLockState {
    fn event(
        state: &mut Self,
        registry: &wl_registry::WlRegistry,
        event: wl_registry::Event,
        _: &(),
        _: &Connection,
        qh: &QueueHandle<NLockState>,
    ) {
        match event {
            wl_registry::Event::Global {
                name,
                interface,
                version,
            } => match &interface[..] {
                "wl_compositor" => {
                    let compositor =
                        registry.bind::<wl_compositor::WlCompositor, _, _>(name, version, qh, ());
                    state.compositor = Some(compositor);
                }
                "wl_subcompositor" => {
                    let subcompositor = registry.bind::<wl_subcompositor::WlSubcompositor, _, _>(
                        name,
                        version,
                        qh,
                        (),
                    );
                    state.subcompositor = Some(subcompositor);
                }
                "wl_shm" => {
                    let shm = registry.bind::<wl_shm::WlShm, _, _>(name, version, qh, ());
                    state.shm = Some(shm);
                }
                "wl_seat" => {
                    let seat = registry.bind::<wl_seat::WlSeat, _, _>(name, version, qh, ());
                    state.r_seat = Some(seat);
                }
                "wl_output" => {
                    trace!("Output ( name: {:?}, ... )", name);

                    let output =
                        registry.bind::<wl_output::WlOutput, _, _>(name, version, qh, name);

                    state
                        .surfaces
                        .entry(name)
                        .or_insert_with(|| NLockSurface::new(output, name));
                }
                "ext_session_lock_manager_v1" => {
                    let session_lock_manager = registry
                        .bind::<ext_session_lock_manager_v1::ExtSessionLockManagerV1, _, _>(
                        name,
                        version,
                        qh,
                        (),
                    );
                    state.session_lock_manager = Some(session_lock_manager);
                }
                _ => {}
            },
            wl_registry::Event::GlobalRemove { name } => {
                trace!("GlobalRemove ( name: {:?}, ... )", name);

                // old surfaces dropped here
                let _ = state.surfaces.remove(&name);
            }
            _ => {}
        }
    }
}

delegate_noop!(NLockState: ignore wl_compositor::WlCompositor);
delegate_noop!(NLockState: ignore wl_subcompositor::WlSubcompositor);
delegate_noop!(NLockState: ignore wl_shm::WlShm);
delegate_noop!(NLockState: ignore wl_surface::WlSurface);
delegate_noop!(NLockState: ignore wl_subsurface::WlSubsurface);
delegate_noop!(NLockState: ignore ext_session_lock_manager_v1::ExtSessionLockManagerV1);
delegate_noop!(NLockState: ignore wl_shm_pool::WlShmPool);
delegate_noop!(NLockState: ignore wl_region::WlRegion);

impl Dispatch<wl_callback::WlCallback, u32> for NLockState {
    fn event(
        state: &mut Self,
        _: &wl_callback::WlCallback,
        event: <wl_callback::WlCallback as wayland_client::Proxy>::Event,
        data: &u32,
        _: &Connection,
        _: &QueueHandle<Self>,
    ) {
        if let wl_callback::Event::Done { callback_data: _ } = event
            && let Some(surface) = state.surfaces.get_mut(data)
        {
            surface.tracking.ready = true;
            if let Err(e) = state.interrupt.write(()) {
                warn!("Failed to write interrupt in frame callback: {:?}", e);
            }
        }
    }
}

impl Dispatch<ext_session_lock_v1::ExtSessionLockV1, ()> for NLockState {
    fn event(
        state: &mut Self,
        _: &ext_session_lock_v1::ExtSessionLockV1,
        event: <ext_session_lock_v1::ExtSessionLockV1 as wayland_client::Proxy>::Event,
        _: &(),
        _: &Connection,
        _: &QueueHandle<Self>,
    ) {
        match event {
            ext_session_lock_v1::Event::Locked => {
                state.locked = true;

                debug!("Session is locked");
            }
            ext_session_lock_v1::Event::Finished => {
                state.unlock();
            }
            _ => {}
        }
    }
}

impl Dispatch<wl_output::WlOutput, u32> for NLockState {
    fn event(
        state: &mut Self,
        _: &wl_output::WlOutput,
        event: <wl_output::WlOutput as wayland_client::Proxy>::Event,
        data: &u32,
        _: &Connection,
        qh: &QueueHandle<Self>,
    ) {
        let Some(surface) = state.surfaces.get_mut(data) else {
            warn!("could not find surface for {}", data);
            return;
        };

        match event {
            wl_output::Event::Geometry {
                x: _,
                y: _,
                physical_width,
                physical_height,
                subpixel,
                make: _,
                model: _,
                transform: _,
            } => {
                surface.set_subpixel_order(cairo::SubpixelOrder::from_wl_subpixel(subpixel));

                if let Err(e) = surface.set_physical_dimensions(physical_width, physical_height) {
                    warn!("Failed to set output physical dimensions: {e}");
                }
            }
            wl_output::Event::Name { name } => {
                debug!("Found output '{name}'");
                surface.output_name = Some(name);
            }
            wl_output::Event::Scale { factor } => {
                if let Err(e) = surface.set_scale(factor) {
                    warn!("Failed to set output scale: {e}");
                } else {
                    debug!(
                        "Set output scale for '{}' to {factor}",
                        surface.output_name.as_ref().unwrap_or(&"".to_string())
                    );
                }
            }
            wl_output::Event::Done => {
                if let (Some(compositor), Some(subcompositor), Some(session_lock)) =
                    (&state.compositor, &state.subcompositor, &state.session_lock)
                {
                    surface.create_surface(compositor, subcompositor, session_lock, qh);
                }
            }
            _ => {}
        }
    }
}