mod backlight;
mod logging;
mod logind;
mod power;

use clap::{Parser, ValueEnum};
use hyprwire::client;
use hyprwire_protocols_sessiond::client::session_backlight::v1::session_backlight_v1;
use hyprwire_protocols_sessiond::client::session_backlight::v1::session_backlight_v1::{
    session_backlight_device, session_backlight_manager,
};
use hyprwire_protocols_sessiond::client::session_core::v1::session_core_v1;
use hyprwire_protocols_sessiond::client::session_core::v1::session_core_v1::{
    session_core_manager, session_core_session,
};
use hyprwire_protocols_sessiond::client::session_power::v1::session_power_v1::{
    self, session_power_manager,
};
use rustix::event::{PollFd, PollFlags, Timespec, poll};
use std::os::fd::OwnedFd;
use std::{env, process, sync};
use tokio::io::unix::AsyncFd;
use tokio::sync::{mpsc, oneshot};
use zbus::fdo;

#[derive(Parser)]
struct Cli {
    #[arg(long, value_enum, default_value_t = LogTarget::Stderr)]
    log_target: LogTarget,
}

#[derive(Clone, Copy, ValueEnum)]
enum LogTarget {
    Stderr,
    Syslog,
}

pub(crate) struct BacklightRequest {
    subsystem: String,
    name: String,
    brightness: u32,
    response: oneshot::Sender<fdo::Result<()>>,
}

#[derive(Clone, Copy)]
pub(crate) enum PowerOperation {
    Reboot,
    Poweroff,
    Suspend,
    Hibernate,
}

#[allow(clippy::struct_excessive_bools)]
#[derive(Clone, Copy, Default)]
pub(crate) struct PowerCapabilities {
    reboot: bool,
    poweroff: bool,
    suspend: bool,
    hibernate: bool,
}

impl PowerCapabilities {
    const fn allows(self, operation: PowerOperation) -> bool {
        match operation {
            PowerOperation::Reboot => self.reboot,
            PowerOperation::Poweroff => self.poweroff,
            PowerOperation::Suspend => self.suspend,
            PowerOperation::Hibernate => self.hibernate,
        }
    }
}

pub(crate) struct PowerRequest {
    pub operation: PowerOperation,
    pub response: oneshot::Sender<fdo::Result<()>>,
}

struct Session {
    id: u32,
    object: session_core_session::SessionCoreSession,
    local: bool,
    uid: u32,
    active: session_core_v1::SessionState,
    done: bool,
}

impl Session {
    const fn new(id: u32, object: session_core_session::SessionCoreSession) -> Self {
        Self {
            id,
            object,
            local: false,
            uid: 0,
            done: false,
            active: session_core_v1::SessionState::Opening,
        }
    }
}

#[derive(Clone, Copy)]
struct SessionAccess {
    id: u32,
    local: bool,
    uid: u32,
    active: session_core_v1::SessionState,
    done: bool,
}

impl From<&Session> for SessionAccess {
    fn from(session: &Session) -> Self {
        Self {
            id: session.id,
            local: session.local,
            uid: session.uid,
            active: session.active,
            done: session.done,
        }
    }
}

struct SessionAccessState {
    sessions: Vec<SessionAccess>,
}

impl SessionAccessState {
    const fn new() -> Self {
        Self {
            sessions: Vec::new(),
        }
    }

    fn update(&mut self, session: &Session) {
        let access = SessionAccess::from(session);

        if let Some(existing) = self
            .sessions
            .iter_mut()
            .find(|existing| existing.id == session.id)
        {
            *existing = access;
        } else {
            self.sessions.push(access);
        }
    }

    fn remove(&mut self, session_id: u32) {
        self.sessions.retain(|session| session.id != session_id);
    }

    fn has_active_local_session(&self, uid: u32) -> bool {
        self.sessions.iter().any(|session| {
            session.done
                && session.local
                && session.uid == uid
                && matches!(session.active, session_core_v1::SessionState::Active)
        })
    }
}

struct SessiondProxy {
    session_core_manager: Option<session_core_manager::SessionCoreManager>,
    backlight_manager: Option<session_backlight_manager::SessionBacklightManager>,
    devices: Vec<BacklightDevice>,
    sessions: Vec<Session>,
    access: sync::Arc<sync::Mutex<SessionAccessState>>,
    pending: Option<oneshot::Sender<fdo::Result<()>>>,
    power_manager: Option<session_power_manager::SessionPowerManager>,
    power_allowed: sync::Arc<sync::Mutex<PowerCapabilities>>,
    power_pending: Option<oneshot::Sender<fdo::Result<()>>>,
}

struct BacklightDevice {
    object: session_backlight_device::SessionBacklightDevice,
    subsystem: String,
    name: String,
}

impl hyprwire::Dispatch<session_core_manager::SessionCoreManager> for SessiondProxy {
    fn event(
        &mut self,
        object: &session_core_manager::SessionCoreManager,
        event: <session_core_manager::SessionCoreManager as hyprwire::Object>::Event<'_>,
    ) {
        if let session_core_manager::Event::SessionAdded { session_id } = event
            && let Some(session_object) = object.send_get_session_by_id::<Self>(session_id)
        {
            self.sessions.push(Session::new(session_id, session_object));
        }
    }
}

impl hyprwire::Dispatch<session_core_session::SessionCoreSession> for SessiondProxy {
    fn event(
        &mut self,
        object: &session_core_session::SessionCoreSession,
        event: <session_core_session::SessionCoreSession as hyprwire::Object>::Event<'_>,
    ) {
        let Some(session) = self
            .sessions
            .iter_mut()
            .find(|session| session.object == *object)
        else {
            return;
        };

        match event {
            session_core_session::Event::State { state } => session.active = state,
            session_core_session::Event::Uid { uid } => session.uid = uid,
            session_core_session::Event::IsLocal { is_local } => session.local = is_local != 0,
            session_core_session::Event::Done => session.done = true,
            session_core_session::Event::Destroyed => {
                let session_id = session.id;
                self.sessions.retain(|session| session.id != session_id);
                if let Ok(mut access) = self.access.lock() {
                    access.remove(session_id);
                }
                return;
            }
            _ => {}
        }

        if let Ok(mut access) = self.access.lock() {
            access.update(session);
        }
    }
}

impl hyprwire::Dispatch<session_backlight_manager::SessionBacklightManager> for SessiondProxy {
    fn event(
        &mut self,
        _object: &session_backlight_manager::SessionBacklightManager,
        event: <session_backlight_manager::SessionBacklightManager as hyprwire::Object>::Event<'_>,
    ) {
        if matches!(event, session_backlight_manager::Event::Failed)
            && let Some(response) = self.pending.take()
        {
            let _ = response.send(Err(fdo::Error::Failed(
                "sessiond could not open the requested backlight device".to_string(),
            )));
        }
    }
}

impl hyprwire::Dispatch<session_backlight_device::SessionBacklightDevice> for SessiondProxy {
    fn event(
        &mut self,
        _object: &session_backlight_device::SessionBacklightDevice,
        event: <session_backlight_device::SessionBacklightDevice as hyprwire::Object>::Event<'_>,
    ) {
        match event {
            session_backlight_device::Event::Done => {
                if let Some(response) = self.pending.take() {
                    let _ = response.send(Ok(()));
                }
            }
            session_backlight_device::Event::SetBrightnessFailed => {
                if let Some(response) = self.pending.take() {
                    let _ = response.send(Err(fdo::Error::Failed(
                        "sessiond could not set the requested brightness".to_string(),
                    )));
                }
            }
            _ => {}
        }
    }
}

impl hyprwire::Dispatch<session_power_manager::SessionPowerManager> for SessiondProxy {
    fn event(
        &mut self,
        _object: &session_power_manager::SessionPowerManager,
        event: <session_power_manager::SessionPowerManager as hyprwire::Object>::Event<'_>,
    ) {
        self.handle_power_event(&event);
    }
}

impl SessiondProxy {
    fn handle_power_event(&mut self, event: &session_power_manager::Event) {
        match event {
            session_power_manager::Event::CanReboot { can_reboot } => {
                if let Ok(mut allowed) = self.power_allowed.lock() {
                    allowed.reboot = *can_reboot != 0;
                }
            }
            session_power_manager::Event::CanPoweroff { can_poweroff } => {
                if let Ok(mut allowed) = self.power_allowed.lock() {
                    allowed.poweroff = *can_poweroff != 0;
                }
            }
            session_power_manager::Event::CanSuspend { can_suspend } => {
                if let Ok(mut allowed) = self.power_allowed.lock() {
                    allowed.suspend = *can_suspend != 0;
                }
            }
            session_power_manager::Event::CanHibernate { can_hibernate } => {
                if let Ok(mut allowed) = self.power_allowed.lock() {
                    allowed.hibernate = *can_hibernate != 0;
                }
            }
            session_power_manager::Event::Done => {
                if let Some(response) = self.power_pending.take() {
                    let _ = response.send(Ok(()));
                }
            }
            session_power_manager::Event::RebootFailed => self.fail_power("reboot"),
            session_power_manager::Event::PoweroffFailed => self.fail_power("power off"),
            session_power_manager::Event::SuspendFailed => self.fail_power("suspend"),
            session_power_manager::Event::HibernateFailed => self.fail_power("hibernate"),
            session_power_manager::Event::HybridSleepFailed => self.fail_power("hybrid sleep"),
            _ => {}
        }
    }

    fn fail_power(&mut self, operation: &str) {
        if let Some(response) = self.power_pending.take() {
            let _ = response.send(Err(fdo::Error::Failed(format!(
                "sessiond {operation} request failed",
            ))));
        }
    }
}

#[tokio::main]
async fn main() -> process::ExitCode {
    let cli = Cli::parse();
    if let Err(error) = init_logging(cli.log_target) {
        eprintln!("failed to initialize logging: {error:#}");
        return process::ExitCode::FAILURE;
    }

    let status = match run().await {
        Ok(()) => process::ExitCode::SUCCESS,
        Err(error) => {
            log::error!("application failed: {error:#}");
            process::ExitCode::FAILURE
        }
    };
    log::logger().flush();
    status
}

fn init_logging(log_target: LogTarget) -> anyhow::Result<()> {
    let log_level = match env::var("LOG_LEVEL") {
        Ok(log) if log == "trace" => log::LevelFilter::Trace,
        Ok(log) if log == "debug" => log::LevelFilter::Debug,
        Ok(log) if log == "warn" => log::LevelFilter::Warn,
        Ok(log) if log == "error" => log::LevelFilter::Error,
        Ok(log) if log == "off" => log::LevelFilter::Off,
        _ => log::LevelFilter::Info,
    };

    match log_target {
        LogTarget::Stderr => {
            env_logger::Builder::new()
                .filter(Some("sessiond_login1_proxy"), log_level)
                .try_init()?;
        }
        LogTarget::Syslog => {
            logging::init_syslog(
                syslog::Facility::LOG_DAEMON,
                log_level,
                "sessiond-login1-proxy",
            )?;
        }
    }
    Ok(())
}

async fn run() -> anyhow::Result<()> {
    let sessions = sync::Arc::new(sync::Mutex::new(SessionAccessState::new()));

    let mut sessiond_proxy = SessiondProxy {
        session_core_manager: None,
        backlight_manager: None,
        devices: Vec::new(),
        sessions: Vec::new(),
        access: sync::Arc::clone(&sessions),
        pending: None,
        power_manager: None,
        power_allowed: sync::Arc::new(sync::Mutex::new(PowerCapabilities::default())),
        power_pending: None,
    };

    let mut sessiond_client = client::Client::connect("/run/sessiond/sessiond.sock")?;
    let sessiond_eq = sessiond_client.new_event_queue();

    sessiond_client.add_implementation::<session_core_v1::SessionCoreV1>();

    sessiond_eq.wait_for_handshake(&mut sessiond_proxy)?;

    sessiond_proxy.session_core_manager = Some(
        sessiond_client
            .bind::<session_core_manager::SessionCoreManager, _>(
                &sessiond_eq,
                &mut sessiond_proxy,
                1,
            )
            .map_err(|error| anyhow::anyhow!("binding session_core_v1: {error}"))?,
    );
    let mut backlight = None::<(client::Client, hyprwire::EventQueue, AsyncFd<OwnedFd>)>;
    let mut power = None::<(client::Client, hyprwire::EventQueue, AsyncFd<OwnedFd>)>;
    let mut reconnect = tokio::time::interval(std::time::Duration::from_secs(5));
    reconnect.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);

    let (requests, mut request_rx) = mpsc::channel(16);
    let (power_requests, mut power_request_rx) = mpsc::channel(4);
    let _logind = logind::Service::start(
        sessions,
        requests,
        sync::Arc::clone(&sessiond_proxy.power_allowed),
        power_requests,
    )
    .await?;

    let sessiond_fd = AsyncFd::new(sessiond_client.extract_loop_fd().try_clone_to_owned()?)?;

    loop {
        tokio::select! {
            _ = reconnect.tick() => {
                if backlight.is_none()
                    && let Ok(mut client) = client::Client::connect("/run/sessiond-backlight.sock")
                {
                    let queue = client.new_event_queue();
                    client.add_implementation::<session_backlight_v1::SessionBacklightV1>();
                    if queue.wait_for_handshake(&mut sessiond_proxy).is_ok()
                        && let Ok(manager) = client.bind::<session_backlight_manager::SessionBacklightManager, _>(
                            &queue, &mut sessiond_proxy, 1,
                        )
                        && let Ok(fd) = client.extract_loop_fd().try_clone_to_owned().and_then(AsyncFd::new)
                    {
                        sessiond_proxy.backlight_manager = Some(manager);
                        backlight = Some((client, queue, fd));
                    }
                }
                if power.is_none()
                    && let Ok(mut client) = client::Client::connect("/run/sessiond/power.sock")
                {
                    let queue = client.new_event_queue();
                    client.add_implementation::<session_power_v1::SessionPowerV1>();
                    if queue.wait_for_handshake(&mut sessiond_proxy).is_ok()
                        && let Ok(manager) = client.bind::<session_power_manager::SessionPowerManager, _>(
                            &queue, &mut sessiond_proxy, 1,
                        )
                        && let Ok(fd) = client.extract_loop_fd().try_clone_to_owned().and_then(AsyncFd::new)
                    {
                        sessiond_proxy.power_manager = Some(manager);
                        power = Some((client, queue, fd));
                    } else if let Ok(mut allowed) = sessiond_proxy.power_allowed.lock() {
                        *allowed = PowerCapabilities::default();
                    }
                }
            }
            Some(request) = request_rx.recv(), if sessiond_proxy.pending.is_none() => {
                if let Some(manager) = &sessiond_proxy.backlight_manager {
                    if let Some(device) = sessiond_proxy.devices.iter().find(|device| {
                        device.subsystem == request.subsystem && device.name == request.name
                    }) {
                        device.object.send_set_brightness(request.brightness);
                        sessiond_proxy.pending = Some(request.response);
                    } else if let Some(device) = manager.send_get_device::<SessiondProxy>(&request.subsystem, &request.name) {
                        device.send_set_brightness(request.brightness);
                        sessiond_proxy.devices.push(BacklightDevice {
                            object: device,
                            subsystem: request.subsystem,
                            name: request.name,
                        });
                        sessiond_proxy.pending = Some(request.response);
                    } else {
                        let _ = request.response.send(Err(fdo::Error::Failed("sessiond did not create a backlight device".to_string())));
                    }
                } else {
                    let _ = request.response.send(Err(fdo::Error::Failed("backlight manager is unavailable".to_string())));
                }
            }
            Some(request) = power_request_rx.recv(), if sessiond_proxy.power_pending.is_none() => {
                if let Some(manager) = &sessiond_proxy.power_manager {
                    match request.operation {
                        PowerOperation::Reboot => manager.send_reboot(),
                        PowerOperation::Poweroff => manager.send_poweroff(),
                        PowerOperation::Suspend => manager.send_suspend(),
                        PowerOperation::Hibernate => manager.send_hibernate(),
                    }
                    sessiond_proxy.power_pending = Some(request.response);
                } else {
                    let _ = request.response.send(Err(fdo::Error::Failed("sessiond power manager is unavailable".to_string())));
                }
            }
            Some((result, queue)) = async {
                if let Some((_, queue, fd)) = &backlight {
                    Some((fd.readable().await, queue))
                } else {
                    None
                }
            } => {
                if result.map_err(anyhow::Error::from).and_then(|mut guard| {
                    dispatch_ready(&mut guard, queue, &mut sessiond_proxy)
                }).is_err() {
                    backlight = None;
                    sessiond_proxy.backlight_manager = None;
                    sessiond_proxy.devices.clear();
                    if let Some(response) = sessiond_proxy.pending.take() {
                        let _ = response.send(Err(fdo::Error::Failed("backlight disconnected".to_string())));
                    }
                }
            }
            result = sessiond_fd.readable() => {
                let mut guard = result?;
                dispatch_ready(&mut guard, &sessiond_eq, &mut sessiond_proxy)?;
            }
            Some((result, queue)) = async {
                if let Some((_, queue, fd)) = &power {
                    Some((fd.readable().await, queue))
                } else {
                    None
                }
            } => {
                if result.map_err(anyhow::Error::from).and_then(|mut guard| {
                    dispatch_ready(&mut guard, queue, &mut sessiond_proxy)
                }).is_err() {
                    power = None;
                    sessiond_proxy.power_manager = None;
                    if let Ok(mut allowed) = sessiond_proxy.power_allowed.lock() {
                        *allowed = PowerCapabilities::default();
                    }
                    if let Some(response) = sessiond_proxy.power_pending.take() {
                        let _ = response.send(Err(fdo::Error::Failed("power disconnected".to_string())));
                    }
                }
            }
        }
    }
}

fn dispatch_ready<D: 'static>(
    guard: &mut tokio::io::unix::AsyncFdReadyGuard<'_, OwnedFd>,
    queue: &hyprwire::EventQueue,
    state: &mut D,
) -> anyhow::Result<()> {
    queue.dispatch_events(state, false)?;

    // Hyprwire returns Ok even when no events remain, rather than WouldBlock.
    // Only clear Tokio's cached readiness once the poller fd is actually idle;
    // otherwise retain it so the next select iteration dispatches more work.
    let mut fds = [PollFd::new(guard.get_inner(), PollFlags::IN)];
    loop {
        match poll(&mut fds, Some(&Timespec::default())) {
            Ok(0) => {
                guard.clear_ready();
                return Ok(());
            }
            Ok(_) => return Ok(()),
            Err(rustix::io::Errno::INTR) => {}
            Err(error) => return Err(error.into()),
        }
    }
}