mod auth;
mod config;
mod effects;
mod input;
mod lock;
mod render;
mod screenshot;
mod system;
mod util;

use config::Config;
use lock::LockManager;
use screenshot::{CaptureData, Screenshot, ScreenshotManager};
use std::error::Error;
use std::fs::OpenOptions;
use std::io::Write;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use system::SystemManager;
use wayland_client::protocol::wl_output::WlOutput;
use wayland_client::protocol::wl_surface::WlSurface;
use wayland_client::{Connection, Dispatch, Proxy, QueueHandle};
use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::{
    Flags, ZwlrScreencopyFrameV1,
};
use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1;
use zeroize::Zeroizing;

use smithay_client_toolkit::reexports::calloop::{self, EventLoop, LoopHandle};
use smithay_client_toolkit::reexports::calloop_wayland_source::WaylandSource;

use smithay_client_toolkit::{
    compositor::{CompositorHandler, CompositorState},
    output::{OutputHandler, OutputState},
    registry::{ProvidesRegistryState, RegistryState},
    registry_handlers,
    seat::{
        keyboard::{KeyEvent, KeyboardHandler, Modifiers},
        pointer::{PointerEvent, PointerEventKind, PointerHandler},
        SeatHandler, SeatState,
    },
    session_lock::{
        SessionLock, SessionLockHandler, SessionLockState, SessionLockSurface,
        SessionLockSurfaceConfigure,
    },
    shm::{slot::SlotPool, Shm, ShmHandler},
};

static FILE_LOGGER: std::sync::LazyLock<std::sync::Mutex<Option<std::fs::File>>> =
    std::sync::LazyLock::new(|| std::sync::Mutex::new(None));

fn setup_file_logging(config: &Config) {
    if let Some(ref path) = config.log_path {
        match OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(path)
        {
            Ok(file) => {
                *FILE_LOGGER.lock().unwrap() = Some(file);
                eprintln!("Logging to: {}", path.display());
            }
            Err(e) => {
                eprintln!("Failed to open log file {}: {}", path.display(), e);
            }
        }
    } else if config.log_file {
        let default_path =
            std::path::PathBuf::from(std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string()))
                .join(".rustlock.log");
        match OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(&default_path)
        {
            Ok(file) => {
                *FILE_LOGGER.lock().unwrap() = Some(file);
                eprintln!("Logging to: {}", default_path.display());
            }
            Err(e) => {
                eprintln!("Failed to open log file {}: {}", default_path.display(), e);
            }
        }
    }
}

fn write_to_file(msg: &str) {
    if let Ok(mut guard) = FILE_LOGGER.lock() {
        if let Some(ref mut file) = *guard {
            let _ = writeln!(file, "{}", msg);
        }
    }
}

struct DualLogger;

impl log::Log for DualLogger {
    fn enabled(&self, metadata: &log::Metadata) -> bool {
        metadata.level() <= log::Level::Debug
    }

    fn log(&self, record: &log::Record) {
        if self.enabled(record.metadata()) {
            let msg = format!(
                "[{}] {}: {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"),
                record.level(),
                record.args()
            );
            eprintln!("{}", msg);
            write_to_file(&msg);
        }
    }

    fn flush(&self) {
        if let Ok(mut guard) = FILE_LOGGER.lock() {
            if let Some(ref mut file) = *guard {
                let _ = file.flush();
            }
        }
    }
}

struct WaylandLock {
    conn: Connection,
    loop_handle: LoopHandle<'static, Self>,
    lock_manager: Arc<Mutex<LockManager>>,
    config: Config,
    ctrlc_exit: Arc<std::sync::atomic::AtomicBool>,
    auth_tx: Option<calloop::channel::Sender<(Zeroizing<String>, u64)>>,
    auth_seq: u64,
    auth_pending_seq: Option<u64>,
    compositor_state: CompositorState,
    output_state: OutputState,
    registry_state: RegistryState,
    session_lock_state: SessionLockState,
    seat_state: SeatState,
    shm_state: Shm,
    pool: SlotPool,
    session_lock: Option<SessionLock>,
    lock_surfaces: Vec<SessionLockSurface>,
    outputs: Vec<WlOutput>,
    screenshot_frames: Vec<Option<ZwlrScreencopyFrameV1>>,
    captured_backgrounds: Vec<Option<cairo::ImageSurface>>,
    pending_screenshots: usize,
    exit: bool,
    screenshot_manager: Option<ScreenshotManager>,
    grace_until: Option<Instant>,
    auth_pending_at: Option<Instant>,
    system_manager: Arc<SystemManager>,
    modifiers: Modifiers,
    current_layout: u32,
}

impl WaylandLock {
    fn handle_auth_result(&mut self, success: bool) {
        // Clear grace period and auth pending on any auth result
        self.grace_until = None;
        self.auth_pending_at = None;
        self.auth_pending_seq = None;

        if success {
            log::info!("✅ Authentication successful - unlocking session");
            // ext-session-lock-v1 recommends destroying lock surfaces before
            // issuing unlock_and_destroy. With multiple outputs Hyprland
            // otherwise treats the unlock as a client crash and shows its
            // failsafe screen.
            self.lock_surfaces.clear();
            if let Some(session_lock) = &self.session_lock {
                session_lock.unlock();
                let _ = self.conn.flush();
                self.exit = true;
                log::info!("Unlock requested - exiting");
            } else {
                log::error!("No session_lock available to unlock!");
                self.exit = true;
            }
        } else {
            log::error!("❌ Authentication failed - wrong password");
            if let Ok(mut lock_manager) = self.lock_manager.lock() {
                for surface in &mut lock_manager.surfaces {
                    surface.show_wrong_password();
                }
            }
        }
    }

    fn get_output_dimensions(&self, output: &WlOutput) -> (i32, i32) {
        if let Some(info) = self.output_state.info(output) {
            if let Some(mode) = info.modes.first() {
                let (w, h) = mode.dimensions;
                return (w, h);
            }
        }
        (1920, 1080)
    }

    fn handle_key_event(&mut self, event: KeyEvent) {
        use crate::input::InputAction;
        use smithay_client_toolkit::seat::keyboard::Keysym;

        // Media and Session keys (Always handle these first and don't trigger grace unlock)
        match event.keysym {
            Keysym::XF86_AudioPlay | Keysym::XF86_AudioPause => {
                self.system_manager.media_play_pause();
                return;
            }
            Keysym::XF86_AudioNext => {
                self.system_manager.media_next();
                return;
            }
            Keysym::XF86_AudioPrev => {
                self.system_manager.media_prev();
                return;
            }
            Keysym::F1 => {
                self.system_manager
                    .send_command(system::BackendCommand::Suspend);
                return;
            }
            Keysym::F2 => {
                self.system_manager
                    .send_command(system::BackendCommand::Reboot);
                return;
            }
            Keysym::F3 => {
                self.system_manager
                    .send_command(system::BackendCommand::PowerOff);
                return;
            }
            _ => {}
        }

        // Check if we're in the grace period (any key unlocks without password)
        if let Some(grace_until) = self.grace_until {
            if Instant::now() < grace_until {
                self.handle_auth_result(true);
                return;
            } else {
                self.grace_until = None;
            }
        }

        if event.keysym == Keysym::Return {
            // Debounce: skip if auth is already pending (user pressed Enter twice)
            if self.auth_pending_at.is_some() {
                return;
            }

            log::info!("Enter pressed - submitting password");
            let password: Option<Zeroizing<String>> = self
                .lock_manager
                .lock()
                .ok()
                .and_then(|mut lm| lm.handle_key_event(event, self.modifiers))
                .and_then(|action| {
                    if let InputAction::SubmitPassword(p) = action {
                        (!p.is_empty()).then_some(p)
                    } else {
                        None
                    }
                });

            if let Some(password) = password {
                self.auth_seq += 1;
                self.auth_pending_at = Some(Instant::now());
                self.auth_pending_seq = Some(self.auth_seq);
                // Show verifying feedback on ALL surfaces BEFORE sending to PAM.
                if let Ok(mut lock_manager) = self.lock_manager.lock() {
                    for surface in &mut lock_manager.surfaces {
                        surface.show_verifying();
                    }
                }
                if let Some(tx) = &self.auth_tx {
                    let _ = tx.send((password, self.auth_seq));
                }
            }
        } else {
            self.lock_manager
                .lock()
                .ok()
                .and_then(|mut lm| lm.handle_key_event(event, self.modifiers));
        }
    }
}

impl ProvidesRegistryState for WaylandLock {
    fn registry(&mut self) -> &mut RegistryState {
        &mut self.registry_state
    }
    registry_handlers![OutputState, SeatState];
}

impl WaylandLock {}

impl CompositorHandler for WaylandLock {
    fn scale_factor_changed(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _surface: &WlSurface,
        _new_factor: i32,
    ) {
    }
    fn transform_changed(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _surface: &WlSurface,
        _new_transform: wayland_client::protocol::wl_output::Transform,
    ) {
    }
    fn frame(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _surface: &WlSurface,
        _time: u32,
    ) {
    }
    fn surface_enter(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _surface: &WlSurface,
        _output: &WlOutput,
    ) {
    }
    fn surface_leave(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _surface: &WlSurface,
        _output: &WlOutput,
    ) {
    }
}

impl OutputHandler for WaylandLock {
    fn output_state(&mut self) -> &mut OutputState {
        &mut self.output_state
    }
    fn new_output(&mut self, _conn: &Connection, qh: &QueueHandle<Self>, output: WlOutput) {
        // If we are already locked, we must create a lock surface for this newly
        // available output. This happens e.g. when a monitor powers back on after
        // "niri msg action power-off-monitors" — niri re-advertises the output and
        // the compositor requires a lock surface on every output or it shows a
        // compositor-defined fallback (typically a solid red/black screen).
        if let Some(session_lock) = &self.session_lock {
            let surface = self.compositor_state.create_surface(qh);
            let (width, height) = self.get_output_dimensions(&output);
            let lock_surface = session_lock.create_lock_surface(surface.clone(), &output, qh);
            self.lock_surfaces.push(lock_surface);
            if !self.outputs.contains(&output) {
                self.outputs.push(output.clone());
            }
            if let Ok(mut lm) = self.lock_manager.lock() {
                lm.add_surface(width, height, output);
                let count = lm.surface_count();
                if let Some(ls) = lm.get_surface_mut(count - 1) {
                    ls.set_wayland_surface(surface);
                }
            }
            log::info!("Created lock surface for newly available output");
        }
    }
    fn update_output(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _output: WlOutput) {
        // Dimension changes while locked are handled by the compositor sending a configure
        // event on the lock surface, which the SessionLockHandler::configure callback
        // already processes via locked_surface.resize().
    }
    fn output_destroyed(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, output: WlOutput) {
        // Clean up the lock surface and associated state for this output.
        // This happens e.g. when a monitor is powered off with its physical power switch.
        // When the monitor comes back on, new_output() will fire and recreate everything.
        //
        // outputs and captured_backgrounds are kept at the same indices, so we remove
        // from both using the same position.
        let output_id = Proxy::id(&output);
        if let Some(idx) = self.outputs.iter().position(|o| Proxy::id(o) == output_id) {
            self.outputs.remove(idx);
            if idx < self.captured_backgrounds.len() {
                self.captured_backgrounds.remove(idx);
            }
        }

        // lock_manager.surfaces and lock_surfaces are built in tandem and share indices,
        // so the index returned from the lock_manager removal applies to lock_surfaces too.
        if let Ok(mut lm) = self.lock_manager.lock() {
            if let Some(idx) = lm.remove_surface_by_output(&output) {
                if idx < self.lock_surfaces.len() {
                    drop(self.lock_surfaces.remove(idx));
                }
            }
        }

        log::info!("Removed lock surface for destroyed output");
    }
}

impl ShmHandler for WaylandLock {
    fn shm_state(&mut self) -> &mut Shm {
        &mut self.shm_state
    }
}

impl SessionLockHandler for WaylandLock {
    fn locked(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _session_lock: SessionLock) {
        log::info!("Session LOCKED confirmed by compositor");
        self.grace_until = Some(Instant::now() + Duration::from_secs_f32(self.config.grace));
    }

    fn finished(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _lock: SessionLock) {
        log::info!("Session lock finished");
        self.exit = true;
    }

    fn configure(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        session_lock_surface: SessionLockSurface,
        configure: SessionLockSurfaceConfigure,
        _serial: u32,
    ) {
        let (width, height) = configure.new_size;
        log::debug!("CONFIGURE callback: {}x{}", width, height);

        if let Ok(mut lock_manager) = self.lock_manager.lock() {
            if let Some(locked_surface) =
                lock_manager.find_surface_by_wayland_surface(session_lock_surface.wl_surface())
            {
                locked_surface.resize(width as i32, height as i32);
                locked_surface.set_configured();

                let output = locked_surface.output();
                let output_id = Proxy::id(output);
                let output_idx = self.outputs.iter().position(|o| Proxy::id(o) == output_id);

                if let Some(idx) = output_idx {
                    if let Some(bg) = self.captured_backgrounds.get(idx).and_then(|b| b.as_ref()) {
                        log::info!(
                            "Applying background for output {} (ID: {:?})",
                            idx,
                            output_id
                        );
                        locked_surface.set_background(bg.clone());
                    }
                }

                locked_surface.update();
                let _ = locked_surface.commit(&mut self.pool);
            }
        }
    }
}

impl KeyboardHandler for WaylandLock {
    fn enter(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
        _surface: &WlSurface,
        _serial: u32,
        _raw: &[u32],
        _keysyms: &[smithay_client_toolkit::seat::keyboard::Keysym],
    ) {
    }
    fn leave(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
        _surface: &WlSurface,
        _serial: u32,
    ) {
    }
    fn press_key(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
        _serial: u32,
        event: KeyEvent,
    ) {
        self.handle_key_event(event);
    }
    fn release_key(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
        _serial: u32,
        _event: KeyEvent,
    ) {
    }
    fn update_modifiers(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
        _serial: u32,
        modifiers: Modifiers,
        layout: u32,
    ) {
        let ctrl_changed = self.modifiers.ctrl != modifiers.ctrl;
        self.modifiers = modifiers;
        self.current_layout = layout;
        if ctrl_changed {
            log::debug!(
                "update_modifiers: ctrl {} -> {}",
                !modifiers.ctrl,
                modifiers.ctrl
            );
        }
        if let Ok(mut lock_manager) = self.lock_manager.lock() {
            lock_manager.set_ctrl_held(modifiers.ctrl);
        }
    }
}

impl SeatHandler for WaylandLock {
    fn seat_state(&mut self) -> &mut SeatState {
        &mut self.seat_state
    }
    fn new_seat(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _seat: wayland_client::protocol::wl_seat::WlSeat,
    ) {
    }
    fn new_capability(
        &mut self,
        _conn: &Connection,
        qh: &QueueHandle<Self>,
        seat: wayland_client::protocol::wl_seat::WlSeat,
        capability: smithay_client_toolkit::seat::Capability,
    ) {
        if capability == smithay_client_toolkit::seat::Capability::Keyboard {
            let _ = self.seat_state.get_keyboard_with_repeat(
                qh,
                &seat,
                None,
                self.loop_handle.clone(),
                Box::new(|state, _kbd, event| {
                    state.handle_key_event(event);
                }),
            );
        }
        if capability == smithay_client_toolkit::seat::Capability::Pointer {
            let _ = self.seat_state.get_pointer(qh, &seat);
        }
    }
    fn remove_capability(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _seat: wayland_client::protocol::wl_seat::WlSeat,
        _capability: smithay_client_toolkit::seat::Capability,
    ) {
    }
    fn remove_seat(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _seat: wayland_client::protocol::wl_seat::WlSeat,
    ) {
    }
}

impl Dispatch<ZwlrScreencopyManagerV1, ()> for WaylandLock {
    fn event(
        _state: &mut Self,
        _proxy: &ZwlrScreencopyManagerV1,
        _event: <ZwlrScreencopyManagerV1 as Proxy>::Event,
        _data: &(),
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
    ) {
    }
}

impl Dispatch<ZwlrScreencopyFrameV1, CaptureData> for WaylandLock {
    fn event(
        state: &mut Self,
        frame: &ZwlrScreencopyFrameV1,
        event: <ZwlrScreencopyFrameV1 as Proxy>::Event,
        data: &CaptureData,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
    ) {
        use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::Event;
        match event {
            Event::Buffer {
                format,
                width,
                height,
                stride,
            } => {
                let Ok(format) = format.into_result() else {
                    log::error!("Screencopy: invalid buffer format, skipping capture");
                    return;
                };

                let mut info = data.info.lock().unwrap();
                *info = Some(screenshot::BufferInfo {
                    width,
                    height,
                    stride,
                    format,
                });

                // Create a dedicated pool for this screenshot to ensure offset 0
                // This matches swaylock-effects and fixes "invalid buffer" on some compositors
                let size = (stride as usize) * (height as usize);
                match SlotPool::new(size, &state.shm_state) {
                    Ok(mut pool) => {
                        match pool.create_buffer(width as i32, height as i32, stride as i32, format)
                        {
                            Ok((buffer, _canvas)) => {
                                frame.copy(buffer.wl_buffer());
                                *data.buffer.lock().unwrap() = Some(buffer);
                                *data.pool.lock().unwrap() = Some(pool);
                            }
                            Err(e) => log::error!("Screencopy: Buffer creation failed: {:?}", e),
                        }
                    }
                    Err(e) => log::error!("Screencopy: Pool creation failed: {:?}", e),
                }
            }
            Event::Flags { flags } => {
                if let Ok(f) = flags.into_result() {
                    *data.flags.lock().unwrap() = Some(f);
                } else {
                    log::error!("Screencopy: invalid flags, skipping");
                }
            }
            Event::Ready { .. } => {
                log::info!("Screencopy: Ready for output {}", data.output_idx);
                if let Some(mgr) = &state.screenshot_manager {
                    let buffer = data.buffer.lock().unwrap().take();
                    let info = data.info.lock().unwrap().take();
                    let flags = data.flags.lock().unwrap().take();
                    let pool = data.pool.lock().unwrap().take();

                    if let (Some(buffer), Some(info), Some(flags), Some(mut pool)) =
                        (buffer, info, flags, pool)
                    {
                        let handle = screenshot::ScreencopyBufferHandle {
                            buffer,
                            info,
                            y_invert: flags.contains(Flags::YInvert),
                        };
                        if let Ok(surface) = mgr.buffer_to_surface(handle, &mut pool) {
                            let mut ss = Screenshot::new(surface);
                            if let Err(e) = ss.apply_effects(&state.config) {
                                log::error!(
                                    "Failed to apply effects to screenshot {}: {e}",
                                    data.output_idx
                                );
                            }
                            if data.output_idx < state.captured_backgrounds.len() {
                                state.captured_backgrounds[data.output_idx] = Some(ss.into_inner());
                            }
                        }
                    }
                }
                state.pending_screenshots -= 1;
                frame.destroy();
            }
            Event::Failed => {
                log::warn!("Screencopy: Failed for output {}", data.output_idx);
                state.pending_screenshots -= 1;
                frame.destroy();
            }
            _ => {}
        }
    }
}

impl wayland_client::Dispatch<wayland_client::protocol::wl_registry::WlRegistry, ()>
    for WaylandLock
{
    fn event(
        _state: &mut Self,
        _proxy: &wayland_client::protocol::wl_registry::WlRegistry,
        _event: wayland_client::protocol::wl_registry::Event,
        _data: &(),
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
    ) {
    }
}

impl PointerHandler for WaylandLock {
    fn pointer_frame(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _pointer: &wayland_client::protocol::wl_pointer::WlPointer,
        events: &[PointerEvent],
    ) {
        for event in events {
            match event.kind {
                PointerEventKind::Press { button: 0x110, .. } => {
                    let (x, y) = event.position;
                    if let Ok(lm) = self.lock_manager.lock() {
                        for surface in &lm.surfaces {
                            if surface.matches_surface(&event.surface) {
                                // Check media controls first
                                let mut handled = false;
                                for (action, rx, ry, rw, rh) in &surface.renderer.media_rects {
                                    if x >= *rx && x <= rx + rw && y >= *ry && y <= ry + rh {
                                        match *action {
                                            "play_pause" => self.system_manager.media_play_pause(),
                                            "stop" => self.system_manager.media_stop(),
                                            "next" => self.system_manager.media_next(),
                                            "prev" => self.system_manager.media_prev(),
                                            _ => {}
                                        }
                                        handled = true;
                                        break;
                                    }
                                }
                                if handled {
                                    return;
                                }
                                // Check indicator ring hit — hold to peek (like Ctrl)
                                let cx = surface.renderer.width as f64 / 2.0;
                                let cy = surface.renderer.height as f64 / 2.0;
                                let r = surface.renderer.config.indicator_radius as f64;
                                let shape = surface.renderer.config.ring_shape;
                                if crate::render::ring_shape::point_in_shape(cx, cy, r, shape, x, y)
                                {
                                    drop(lm);
                                    if let Ok(mut lm) = self.lock_manager.lock() {
                                        if let Some(s) = lm
                                            .surfaces
                                            .iter_mut()
                                            .find(|s| s.matches_surface(&event.surface))
                                        {
                                            s.set_peek_held(true);
                                            s.update();
                                            let _ = s.commit(&mut self.pool);
                                        }
                                    }
                                    return;
                                }
                            }
                        }
                    }
                }
                PointerEventKind::Release { button: 0x110, .. } => {
                    if let Ok(mut lm) = self.lock_manager.lock() {
                        // Clear peek on ALL surfaces. Wayland button release events
                        // can reference a different wl_surface proxy than the press
                        // event (compositor-dependent pointer grab semantics), so
                        // matching by surface might miss the LockedSurface that has
                        // peek_toggled=true. Iterating every surface guarantees peek
                        // always ends on button release — matching Ctrl-hold behavior.
                        for surface in &mut lm.surfaces {
                            surface.set_peek_held(false);
                            surface.update();
                            let _ = surface.commit(&mut self.pool);
                        }
                    }
                }
                _ => {}
            }
        }
    }
}

smithay_client_toolkit::delegate_pointer!(WaylandLock);
smithay_client_toolkit::delegate_compositor!(WaylandLock);
smithay_client_toolkit::delegate_output!(WaylandLock);
smithay_client_toolkit::delegate_shm!(WaylandLock);
smithay_client_toolkit::delegate_seat!(WaylandLock);
smithay_client_toolkit::delegate_keyboard!(WaylandLock);
smithay_client_toolkit::delegate_registry!(WaylandLock);
smithay_client_toolkit::delegate_session_lock!(WaylandLock);
wayland_client::delegate_noop!(WaylandLock: ignore wayland_client::protocol::wl_buffer::WlBuffer);

fn main() -> Result<(), Box<dyn Error>> {
    let config = Config::load();

    setup_file_logging(&config);
    static LOGGER: DualLogger = DualLogger;
    let max_level = if config.debug {
        log::LevelFilter::Debug
    } else {
        log::LevelFilter::Info
    };
    log::set_logger(&LOGGER).map(|()| log::set_max_level(max_level))?;

    log::info!("Starting rustlock v{}", env!("CARGO_PKG_VERSION"));
    #[allow(clippy::arc_with_non_send_sync)]
    let lock_manager = Arc::new(Mutex::new(LockManager::new(config.clone())));
    let ctrlc_exit = Arc::new(std::sync::atomic::AtomicBool::new(false));

    let conn = Connection::connect_to_env()?;
    let (globals, mut event_queue) =
        wayland_client::globals::registry_queue_init::<WaylandLock>(&conn)?;
    let qh: QueueHandle<WaylandLock> = event_queue.handle();

    let shm_state = Shm::bind(&globals, &qh).map_err(|_| "wl_shm not supported")?;

    let system_manager = Arc::new(SystemManager::new(&config));

    let (auth_tx_actual, auth_feedback_rx_actual) = match auth::create_and_run_auth_loop(
        config.pam_service.clone(),
    ) {
        Some(channels) => channels,
        None => {
            log::error!("Failed to initialize authentication. This usually means PAM is not configured correctly.");
            log::error!(
                "Please ensure you have a PAM service file at /etc/pam.d/{}",
                config.pam_service
            );
            std::process::exit(1);
        }
    };
    let mut event_loop: EventLoop<WaylandLock> = EventLoop::try_new()?;

    // Initialize state without pool first
    // We'll create the pool after we know the output dimensions

    // Initialize pool with a minimal size, it will be resized once outputs are detected
    let pool = SlotPool::new(1, &shm_state)?;

    let mut state = WaylandLock {
        conn: conn.clone(),
        loop_handle: event_loop.handle(),
        lock_manager: lock_manager.clone(),
        config: config.clone(),
        ctrlc_exit: ctrlc_exit.clone(),
        auth_tx: Some(auth_tx_actual),
        compositor_state: CompositorState::bind(&globals, &qh)?,
        output_state: OutputState::new(&globals, &qh),
        registry_state: RegistryState::new(&globals),
        session_lock_state: SessionLockState::new(&globals, &qh),
        seat_state: SeatState::new(&globals, &qh),
        shm_state,
        pool,
        session_lock: None,
        lock_surfaces: Vec::new(),
        outputs: Vec::new(),
        screenshot_frames: Vec::new(),
        captured_backgrounds: Vec::new(),
        pending_screenshots: 0,
        exit: false,
        screenshot_manager: ScreenshotManager::new(&globals, &qh).ok(),
        grace_until: None,
        auth_pending_at: None,
        auth_seq: 0,
        auth_pending_seq: None,
        system_manager: system_manager.clone(),
        modifiers: Modifiers::default(),
        current_layout: 0,
    };

    event_queue.blocking_dispatch(&mut state)?;

    // Handle custom background image if provided (prioritize over screenshots)
    if let Some(ref image_path) = state.config.image {
        log::info!("Loading custom background image from {:?}", image_path);
        if let Ok(img) = image::open(image_path) {
            let img = img.to_rgba8();
            let (w, h) = img.dimensions();
            if let Ok(mut surface) =
                cairo::ImageSurface::create(cairo::Format::ARgb32, w as i32, h as i32)
            {
                let image_ok = {
                    if let Ok(mut surface_data) = surface.data() {
                        for y in 0..h {
                            for x in 0..w {
                                let pixel = img.get_pixel(x, y);
                                let idx = ((y * w + x) * 4) as usize;
                                surface_data[idx] = pixel[2];
                                surface_data[idx + 1] = pixel[1];
                                surface_data[idx + 2] = pixel[0];
                                surface_data[idx + 3] = pixel[3];
                            }
                        }
                        true
                    } else {
                        log::error!("Failed to get background image surface data");
                        false
                    }
                };
                if image_ok {
                    let mut ss = Screenshot::new(surface);
                    if let Err(e) = ss.apply_effects(&state.config) {
                        log::error!("Failed to apply effects to custom background image: {e}");
                    }
                    let surface = ss.into_inner();

                    let num_outputs = state.output_state.outputs().count();
                    state.captured_backgrounds = vec![Some(surface); num_outputs];
                    state.config.screenshots = false;
                }
            } else {
                log::error!("Failed to create Cairo surface for background image");
            }
        } else {
            log::error!(
                "Failed to load custom background image from {:?}",
                image_path
            );
        }
    }

    // Now that we have output info, we can resize the pool if needed
    let mut total_size = 0;
    for output in state.output_state.outputs() {
        if let Some(info) = state.output_state.info(&output) {
            if let Some(mode) = info.modes.first() {
                let (w, h) = mode.dimensions;
                total_size += (w as usize) * (h as usize) * 4;
            }
        }
    }
    if total_size > 0 {
        // Resize pool to fit all outputs
        // SlotPool doesn't have a direct resize, but we can just create a new one if needed,
        // or rely on its internal growing if we use it that way.
        // Actually SCTK SlotPool handles resizing when creating buffers if needed,
        // but it's better to have a large enough base.
        state.pool = SlotPool::new(total_size, &state.shm_state)?;
    }

    let _wayland_source =
        WaylandSource::new(conn.clone(), event_queue).insert(event_loop.handle())?;

    event_loop
        .handle()
        .insert_source(auth_feedback_rx_actual, |event, _, state| {
            if let calloop::channel::Event::Msg((success, seq)) = event {
                // Ignore stale auth results from previous requests (e.g. after timeout or retry).
                if state.auth_pending_seq == Some(seq) {
                    state.handle_auth_result(success);
                }
            }
        })?;

    let timer = calloop::timer::Timer::from_duration(Duration::from_millis(16));
    event_loop.handle().insert_source(timer, |_, _, state| {
        // Clear expired grace period
        if let Some(grace_until) = state.grace_until {
            if Instant::now() >= grace_until {
                state.grace_until = None;
            }
        }

        // Auth timeout: if PAM thread doesn't respond within config.auth_timeout ms,
        // treat as auth failure so the user gets feedback instead of hanging forever.
        if state.auth_pending_seq.is_some() {
            if let Some(at) = state.auth_pending_at {
                if Instant::now().duration_since(at)
                    >= Duration::from_millis(state.config.auth_timeout)
                {
                    log::warn!(
                        "Authentication timed out after {} ms",
                        state.config.auth_timeout
                    );
                    // Clear pending seq so the eventual PAM result is ignored as stale
                    state.auth_pending_seq = None;
                    state.handle_auth_result(false);
                }
            }
        }

        let mut status = state.system_manager.get_status();
        status.keyboard_layout = Some(state.current_layout.to_string());

        if let Ok(mut lm) = state.lock_manager.lock() {
            lm.set_system_status(status);
            // Only commit surfaces that actually re-rendered this tick. An idle
            // lock screen renders nothing and commits nothing.
            for surface in &mut lm.surfaces {
                if surface.update() {
                    let _ = surface.commit(&mut state.pool);
                }
            }
        }
        if state.exit {
            calloop::timer::TimeoutAction::Drop
        } else {
            calloop::timer::TimeoutAction::ToDuration(Duration::from_millis(16))
        }
    })?;

    // Dispatch to ensure outputs are ready before capture
    if state.config.screenshots {
        event_loop.dispatch(Duration::from_millis(50), &mut state)?;
    }

    if state.config.screenshots {
        state.outputs = state.output_state.outputs().collect();
        log::info!(
            "Capturing screenshots for {} outputs...",
            state.outputs.len()
        );
        state.screenshot_frames = vec![None; state.outputs.len()];
        state.captured_backgrounds = vec![None; state.outputs.len()];
        state.pending_screenshots = state.outputs.len();

        for (i, output) in state.outputs.iter().enumerate() {
            if let Some(mgr) = &state.screenshot_manager {
                let data = CaptureData::new(i);
                match mgr.capture_output(output, &qh, data) {
                    Ok(frame) => state.screenshot_frames[i] = Some(frame),
                    Err(e) => {
                        log::error!("Screencopy failed for output {}: {:?}", i, e);
                        state.pending_screenshots -= 1;
                    }
                }
            } else {
                state.pending_screenshots -= 1;
            }
        }

        let start = std::time::Instant::now();
        while state.pending_screenshots > 0 && start.elapsed() < Duration::from_secs(2) {
            event_loop.dispatch(Duration::from_millis(50), &mut state)?;
        }
        log::info!("Screenshot capture phase complete");
    }

    log::info!("Attempting to lock Wayland session...");
    let session_lock = state.session_lock_state.lock(&qh).map_err(|e| {
        log::error!("Lock failed: {}", e);
        e
    })?;

    let outputs_to_lock: Vec<WlOutput> = state.output_state.outputs().collect();
    for output in outputs_to_lock {
        let surface = state.compositor_state.create_surface(&qh);
        let (width, height) = state.get_output_dimensions(&output);
        let lock_surface = session_lock.create_lock_surface(surface.clone(), &output, &qh);
        state.lock_surfaces.push(lock_surface);
        if !state.outputs.contains(&output) {
            state.outputs.push(output.clone());
        }
        if let Ok(mut lm) = state.lock_manager.lock() {
            lm.add_surface(width, height, output.clone());
            let count = lm.surface_count();
            if let Some(ls) = lm.get_surface_mut(count - 1) {
                ls.set_wayland_surface(surface);
            }
        }
    }
    state.session_lock = Some(session_lock);

    while !state.exit && !state.ctrlc_exit.load(std::sync::atomic::Ordering::SeqCst) {
        event_loop.dispatch(Duration::from_millis(16), &mut state)?;
    }

    // After unlock_and_destroy the compositor sends keyboard/pointer leave and
    // delete_id events that must be processed before we disconnect, otherwise
    // Hyprland treats the disconnect as a client crash and shows its failsafe
    // screen (only reproducible on multi-output setups).
    let _ = state.conn.roundtrip();

    log::info!("Exiting rustlock");
    Ok(())
}
