mod auth;
mod config;
mod input;
mod lock;
mod render;
mod screenshot;
mod timer;
mod util;

use config::Config;
use lock::LockManager;
use screenshot::{CaptureData, ScreenshotManager, Screenshot};
use std::error::Error;
use std::fs::OpenOptions;
use std::io::Write;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use wayland_client::protocol::wl_output::WlOutput;
use wayland_client::protocol::wl_surface::WlSurface;
use wayland_client::{Connection, Dispatch, QueueHandle, Proxy};
use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::{ZwlrScreencopyFrameV1, Flags};
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},
        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) {
    let log_path = std::env::var("HOME")
        .unwrap_or_else(|_| "/tmp".to_string())
        + "/.wayrustlock.log";

    match OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(&log_path)
    {
        Ok(file) => {
            *FILE_LOGGER.lock().unwrap() = Some(file);
            eprintln!("Logging to: {}", log_path);
        }
        Err(e) => {
            eprintln!("Failed to open log file {}: {}", log_path, 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 {
    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>>>,
    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,
    unlocking: bool,
    screenshot_manager: Option<ScreenshotManager>,
}

impl WaylandLock {
    fn handle_auth_result(&mut self, success: bool) {
        if success {
            log::info!("✅ Authentication successful - unlocking session");
            if let Some(session_lock) = &self.session_lock {
                session_lock.unlock();
                self.unlocking = true;
                log::debug!("Unlock requested - waiting for compositor finished event");
            } else {
                log::error!("No session_lock available to unlock!");
                self.exit = true;
            }
        } else {
            log::warn!("❌ 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 as i32, h as i32);
            }
        }
        (1920, 1080)
    }

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

        if event.keysym == Keysym::Return {
            log::info!("Enter pressed - submitting password");
            if let Ok(mut lock_manager) = self.lock_manager.lock() {
                let mut password = Zeroizing::new(String::new());
                for surface in &mut lock_manager.surfaces {
                    if let Some(InputAction::SubmitPassword(p)) = surface.handle_key_event(event.clone()) {
                        password = p;
                    }
                }
                if !password.is_empty() {
                    if let Some(tx) = &self.auth_tx {
                        let _ = tx.send(password);
                    }
                }
            }
        } else {
            let action = self
                .lock_manager
                .lock()
                .map(|mut lm| lm.handle_key_event(event))
                .unwrap_or(None);

            if let Some(InputAction::TempScreenshot) = action {
                if let Ok(mut lm) = self.lock_manager.lock() {
                    lm.toggle_peek();
                }
            }
        }
    }
}

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

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) {}
    fn update_output(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _output: WlOutput) {}
    fn output_destroyed(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _output: WlOutput) {}
}

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");
    }

    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) {}
}

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| {}));
        }
    }
    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 format = format.into_result().unwrap();
                let mut info = data.info.lock().unwrap();
                *info = Some(screenshot::BufferInfo { width, height, stride, format });
                match state.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);
                    }
                    Err(e) => log::error!("Screencopy: Buffer creation failed: {:?}", e),
                }
            }
            Event::Flags { flags } => { *data.flags.lock().unwrap() = Some(flags.into_result().unwrap()); }
            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();
                    if let (Some(buffer), Some(info), Some(flags)) = (buffer, info, flags) {
                        let handle = screenshot::ScreencopyBufferHandle { buffer, info, y_invert: flags.contains(Flags::YInvert) };
                        if let Ok(surface) = mgr.buffer_to_surface(handle, &mut state.pool) {
                            let mut ss = Screenshot::new(surface);
                            let _ = ss.apply_effects(&state.config);
                            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>) {}
}

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;
    log::set_logger(&LOGGER).map(|()| log::set_max_level(log::LevelFilter::Debug))?;

    log::info!("Starting wayrustlock v{}", env!("CARGO_PKG_VERSION"));
    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 pool = SlotPool::new(1920 * 1080 * 4, &shm_state)?;

    let (auth_tx_actual, auth_feedback_rx_actual) = auth::create_and_run_auth_loop();
    let mut event_loop: EventLoop<WaylandLock> = EventLoop::try_new()?;

    let mut state = WaylandLock {
        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,
        unlocking: false,
        screenshot_manager: ScreenshotManager::new(&globals, &qh).ok(),
    };

    event_queue.blocking_dispatch(&mut 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) = event {
            state.handle_auth_result(success);
            if success { state.lock_surfaces.clear(); }
        }
    })?;

    let timer = calloop::timer::Timer::from_duration(Duration::from_millis(16));
    event_loop.handle().insert_source(timer, |_, _, state| {
        if state.unlocking {
            return calloop::timer::TimeoutAction::ToDuration(Duration::from_millis(100));
        }

        if let Ok(mut lm) = state.lock_manager.lock() {
            lm.update();
            for surface in &mut lm.surfaces { let _ = surface.commit(&mut state.pool); }
        }
        if state.exit { calloop::timer::TimeoutAction::Drop } else { calloop::timer::TimeoutAction::ToDuration(Duration::from_millis(16)) }
    })?;

    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)?;
    }

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