//! Wayland state: output enumeration, layer surface creation, event dispatch.

use std::time::{Duration, Instant};

use anyhow::{bail, Context, Result};
use smithay_client_toolkit::{
    compositor::{CompositorHandler, CompositorState},
    delegate_compositor, delegate_layer, delegate_output, delegate_registry,
    output::{OutputHandler, OutputState},
    registry::{ProvidesRegistryState, RegistryState},
    registry_handlers,
    shell::{
        wlr_layer::{
            Anchor, KeyboardInteractivity, Layer, LayerShell, LayerShellHandler, LayerSurface,
            LayerSurfaceConfigure,
        },
        WaylandSurface,
    },
};
use wayland_client::{
    globals::registry_queue_init,
    protocol::{wl_output, wl_surface},
    Connection, Proxy, QueueHandle,
};

use crate::renderer::Renderer;
use crate::shader::ShaderSource;

/// Application state held across the Wayland event loop.
pub struct AppState {
    // SCTK required fields
    registry_state: RegistryState,
    output_state: OutputState,

    // Layer shell
    layer: Option<LayerSurface>,

    // Rendering
    renderer: Option<Renderer>,

    // Configuration that renderer needs at configure time
    display_ptr: *mut std::ffi::c_void,
    surface_ptr: *mut std::ffi::c_void,
    shader: ShaderSource,
    fps: u32,

    // Bookkeeping
    configured: bool,
    width: u32,
    height: u32,
    pub exit: bool,

    // Frame callback: set true when compositor says it's ready for next frame
    draw_ready: bool,
    last_frame: Instant,
}

// Safety: we only use the raw pointers in single-threaded context
unsafe impl Send for AppState {}
unsafe impl Sync for AppState {}

/// Parse the layer CLI string into an SCTK Layer.
pub fn parse_layer(s: &str) -> Result<Layer> {
    match s.to_lowercase().as_str() {
        "background" => Ok(Layer::Background),
        "bottom" => Ok(Layer::Bottom),
        "top" => Ok(Layer::Top),
        "overlay" => Ok(Layer::Overlay),
        other => bail!("unknown layer: {}", other),
    }
}

/// Set up a Wayland connection and run the render loop for one output.
pub fn run(
    target_output: String,
    shader: ShaderSource,
    fps: u32,
    layer_kind: Layer,
) -> Result<()> {
    let conn = Connection::connect_to_env().context("failed to connect to Wayland display")?;
    let display_ptr = conn.backend().display_ptr() as *mut std::ffi::c_void;

    let (globals, mut event_queue) =
        registry_queue_init::<AppState>(&conn).context("registry_queue_init")?;
    let qh = event_queue.handle();

    let compositor =
        CompositorState::bind(&globals, &qh).context("wl_compositor not available")?;
    let layer_shell =
        LayerShell::bind(&globals, &qh).context("zwlr_layer_shell_v1 not available")?;

    let mut state = AppState {
        registry_state: RegistryState::new(&globals),
        output_state: OutputState::new(&globals, &qh),
        layer: None,
        renderer: None,
        display_ptr,
        surface_ptr: std::ptr::null_mut(),
        shader,
        fps,
        configured: false,
        width: 0,
        height: 0,
        exit: false,
        draw_ready: false,
        last_frame: Instant::now(),
    };

    // One roundtrip to discover outputs.
    event_queue.roundtrip(&mut state).context("initial roundtrip")?;

    // Find the matching wl_output.
    let target_wl_output = find_output(&state.output_state, &target_output);

    // Create the wl_surface and layer surface.
    let surface = compositor.create_surface(&qh);
    state.surface_ptr = surface.id().as_ptr() as *mut std::ffi::c_void;

    let layer = layer_shell.create_layer_surface(
        &qh,
        surface,
        layer_kind,
        Some("gpupaper"),
        target_wl_output.as_ref(),
    );

    // Stretch to fill the entire output.
    layer.set_anchor(Anchor::TOP | Anchor::BOTTOM | Anchor::LEFT | Anchor::RIGHT);
    layer.set_exclusive_zone(-1);
    layer.set_size(0, 0);
    layer.set_keyboard_interactivity(KeyboardInteractivity::None);
    layer.commit();
    state.layer = Some(layer);

    // Wait for the configure event which gives us the actual dimensions.
    log::info!("waiting for configure…");
    while !state.configured {
        event_queue.blocking_dispatch(&mut state).context("event dispatch failed")?;
    }

    log::info!("configured: {}x{}", state.width, state.height);

    // Create renderer now that we have confirmed dimensions.
    state.renderer = Some(unsafe {
        Renderer::new(
            state.display_ptr,
            state.surface_ptr,
            state.width,
            state.height,
            &state.shader,
            state.fps,
        )
        .context("failed to create renderer")?
    });

    // Render loop driven by wl_surface.frame callbacks.
    // The compositor stops sending callbacks when the surface is fully occluded,
    // so the loop naturally idles to zero CPU when another app is fullscreen.
    let frame_duration = if fps > 0 {
        Some(Duration::from_nanos(1_000_000_000 / fps as u64))
    } else {
        None
    };

    // Register the first frame callback then render immediately.  On Wayland
    // the compositor only fires frame callbacks for commits that contain a
    // buffer; an empty commit would stall the loop forever on many compositors.
    // Rendering here attaches the first buffer and commits it together with the
    // frame callback so the compositor has something to display right away.
    {
        let wl_surface = state.layer.as_ref().unwrap().wl_surface();
        wl_surface.frame(&qh, wl_surface.clone());
    }
    if let Some(r) = state.renderer.as_mut() {
        r.render();
    }
    conn.flush().ok();

    loop {
        // Block the thread until the compositor sends us any event (including
        // the frame callback).  Zero CPU while occluded.
        event_queue.blocking_dispatch(&mut state).context("event dispatch")?;
        conn.flush().ok();

        if state.exit {
            break;
        }

        if !state.draw_ready {
            continue;
        }
        state.draw_ready = false;

        // Optional fps cap: sleep the remainder of the frame budget.
        if let Some(dur) = frame_duration {
            let elapsed = state.last_frame.elapsed();
            if elapsed < dur {
                std::thread::sleep(dur - elapsed);
            }
        }
        state.last_frame = Instant::now();

        // Register the NEXT frame callback before present() so that the
        // request is included in the same wl_surface.commit that wgpu
        // issues inside present().
        {
            let wl_surface = state.layer.as_ref().unwrap().wl_surface();
            wl_surface.frame(&qh, wl_surface.clone());
        }

        if let Some(r) = state.renderer.as_mut() {
            r.render();
        }

        conn.flush().ok();
    }

    Ok(())
}

/// Try to find the wl_output whose name matches `target`.
/// If `target` is `"all"` or `"*"`, returns `None` (compositor picks).
fn find_output(output_state: &OutputState, target: &str) -> Option<wl_output::WlOutput> {
    if target == "all" || target == "*" {
        return None;
    }
    for output in output_state.outputs() {
        if let Some(info) = output_state.info(&output) {
            if let Some(name) = &info.name {
                if name == target {
                    log::info!("matched output: {}", name);
                    return Some(output);
                }
            }
        }
    }
    log::warn!("output '{}' not found; using default", target);
    None
}

// ---------------------------------------------------------------------------
// Trait implementations for AppState
// ---------------------------------------------------------------------------

impl CompositorHandler for AppState {
    fn scale_factor_changed(
        &mut self, _: &Connection, _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface, _: i32,
    ) {}

    fn transform_changed(
        &mut self, _: &Connection, _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface, _: wl_output::Transform,
    ) {}

    fn frame(
        &mut self, _: &Connection, _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface, _: u32,
    ) {
        self.draw_ready = true;
    }

    fn surface_enter(
        &mut self, _: &Connection, _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface, _: &wl_output::WlOutput,
    ) {}

    fn surface_leave(
        &mut self, _: &Connection, _: &QueueHandle<Self>,
        _: &wl_surface::WlSurface, _: &wl_output::WlOutput,
    ) {}
}

impl OutputHandler for AppState {
    fn output_state(&mut self) -> &mut OutputState {
        &mut self.output_state
    }

    fn new_output(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_output::WlOutput) {}
    fn update_output(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_output::WlOutput) {}
    fn output_destroyed(&mut self, _: &Connection, _: &QueueHandle<Self>, _: wl_output::WlOutput) {}
}

impl LayerShellHandler for AppState {
    fn closed(&mut self, _: &Connection, _: &QueueHandle<Self>, _: &LayerSurface) {
        log::info!("layer surface closed");
        self.exit = true;
    }

    fn configure(
        &mut self,
        _conn: &Connection,
        _qh: &QueueHandle<Self>,
        _layer: &LayerSurface,
        configure: LayerSurfaceConfigure,
        _serial: u32,
    ) {
        let (w, h) = configure.new_size;
        self.width = if w == 0 { 1920 } else { w };
        self.height = if h == 0 { 1080 } else { h };

        if !self.configured {
            self.configured = true;
        } else if let Some(r) = self.renderer.as_mut() {
            r.resize(self.width, self.height);
        }

        log::debug!("configure: {}x{}", self.width, self.height);
    }
}

// Delegation macros
delegate_compositor!(AppState);
delegate_output!(AppState);
delegate_layer!(AppState);
delegate_registry!(AppState);

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