//! One-shot KWin adapter. The installed executable, not Electron or the AppImage
//! mount, is the identity authorized by KDE's desktop-entry permission check.
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    fs::OpenOptions,
    io::{Read, Write},
    os::{fd::AsFd, unix::net::UnixStream},
    path::Path,
    sync::mpsc,
    time::Duration,
};
use zbus::{
    blocking::{Connection, Proxy, connection::Builder},
    zvariant::{Fd, OwnedValue, Value},
};

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
const SCREENSHOT: &str = "org.kde.KWin.ScreenShot2";
const MAX_BYTES: usize = 128 * 1024 * 1024;
const DEADLINE: Duration = Duration::from_secs(5);

mod feedback;

#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
struct Window {
    id: String,
    title: String,
    app_name: String,
    app_identifier: String,
    process_id: u32,
    bounds: Bounds,
    client_bounds: Bounds,
}

impl Window {
    fn matches_after_capture(&self, current: &Self) -> bool {
        self.id == current.id
            && self.process_id == current.process_id
            && self.app_identifier == current.app_identifier
            && self.app_name == current.app_name
            && self.bounds == current.bounds
            && self.client_bounds == current.client_bounds
            && capture_title(&self.title) == capture_title(&current.title)
    }
}

/// Match the same leading CLI spinner accepted by the Wayland AT-SPI window matcher.
fn capture_title(title: &str) -> &str {
    let title = title.trim();
    let mut chars = title.chars();
    if matches!(
        chars.next(),
        Some('⠋' | '⠙' | '⠹' | '⠸' | '⠼' | '⠴' | '⠦' | '⠧' | '⠇' | '⠏')
    ) && chars.clone().next().is_some_and(char::is_whitespace)
    {
        let label = chars.as_str().trim_start();
        if !label.is_empty() {
            return label;
        }
    }
    title
}

#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
struct Bounds {
    x: i32,
    y: i32,
    width: u32,
    height: u32,
}

struct ScriptReply {
    owner: String,
    result: mpsc::SyncSender<String>,
}
#[zbus::interface(name = "com.t3tools.KdeCapture")]
impl ScriptReply {
    fn reply(
        &self,
        value: &str,
        #[zbus(header)] header: zbus::message::Header<'_>,
    ) -> zbus::fdo::Result<()> {
        if header.sender().map(|name| name.as_str()) != Some(self.owner.as_str()) {
            return Err(zbus::fdo::Error::AccessDenied("Not KWin".into()));
        }
        if value.len() > 64 * 1024 {
            return Err(zbus::fdo::Error::InvalidArgs(
                "Oversized window metadata".into(),
            ));
        }
        let _ = self.result.try_send(value.to_owned());
        Ok(())
    }
}

/// Scripts exist only for this operation. Unload even when evaluation/callback fails.
fn script(connection: &Connection, directory: &Path, body: &str) -> Result<String> {
    let dbus = zbus::blocking::fdo::DBusProxy::new(connection)?;
    let owner = dbus.get_name_owner("org.kde.KWin".try_into()?)?.to_string();
    let (send, receive) = mpsc::sync_channel(1);
    connection.object_server().at(
        "/com/t3tools/KdeCapture",
        ScriptReply {
            owner,
            result: send,
        },
    )?;
    let path = directory.join("window.js");
    let name = format!("t3-capture-{}", std::process::id());
    let destination = serde_json::to_string(
        connection
            .unique_name()
            .ok_or("Missing bus identity")?
            .as_str(),
    )?;
    let source = format!(
        "function reply(value) {{ callDBus({destination}, '/com/t3tools/KdeCapture', 'com.t3tools.KdeCapture', 'Reply', JSON.stringify(value)); }}\ntry {{ {body} }} catch (error) {{ reply({{error: String(error)}}); }}"
    );
    std::fs::write(&path, source)?;
    let scripting = Proxy::new(
        connection,
        "org.kde.KWin",
        "/Scripting",
        "org.kde.kwin.Scripting",
    )?;
    let result = (|| {
        let id: i32 = scripting.call(
            "loadScript",
            &(path.to_str().ok_or("Invalid script path")?, &name),
        )?;
        if id < 0 {
            return Err("KWin could not load the capture helper script".into());
        }
        let result = (|| {
            let object = format!("/Scripting/Script{id}");
            let loaded = Proxy::new(
                connection,
                "org.kde.KWin",
                object.as_str(),
                "org.kde.kwin.Script",
            )?;
            let _: () = loaded.call("run", &())?;
            Ok(receive.recv_timeout(DEADLINE)?)
        })();
        let _: std::result::Result<bool, _> = scripting.call("unloadScript", &(&name,));
        result
    })();
    connection
        .object_server()
        .remove::<ScriptReply, _>("/com/t3tools/KdeCapture")?;
    let _ = std::fs::remove_file(path);
    result
}

fn window(connection: &Connection, directory: &Path, id: Option<&str>) -> Result<Option<Window>> {
    let selector = match id {
        Some(id) => format!(
            "workspace.windowList().find(w => String(w.internalId) === {})",
            serde_json::to_string(id)?
        ),
        None => "workspace.activeWindow".into(),
    };
    let value = script(
        connection,
        directory,
        &format!(
            "function bounds(rect) {{ return {{ x: Math.round(rect.x), y: Math.round(rect.y), width: Math.round(rect.width), height: Math.round(rect.height) }}; }}
             const w = {selector}; reply(!w ? null : {{ id: String(w.internalId), title: w.caption, appName: w.resourceClass || w.desktopFileName || 'Application', appIdentifier: w.desktopFileName || w.resourceClass || '', processId: w.pid, bounds: bounds(w.frameGeometry), clientBounds: bounds(w.clientGeometry) }});"
        ),
    )?;
    Ok(serde_json::from_str(&value)?)
}

fn screenshot_proxy(connection: &Connection) -> zbus::Result<Proxy<'_>> {
    Proxy::new(
        connection,
        SCREENSHOT,
        "/org/kde/KWin/ScreenShot2",
        SCREENSHOT,
    )
}

/// InvalidWindow follows the authorization check and cannot produce a screenshot.
fn check(connection: &Connection) -> Result<()> {
    let proxy = screenshot_proxy(connection)?;
    if proxy.get_property::<u32>("Version")? < 2 {
        return Err("KDE's screenshot API is too old".into());
    }
    let sink = OpenOptions::new().write(true).open("/dev/null")?;
    let reply: zbus::Result<HashMap<String, OwnedValue>> = proxy.call(
        "CaptureWindow",
        &(
            "t3-permission-check-not-a-window",
            HashMap::<&str, Value<'_>>::new(),
            Fd::from(sink.as_fd()),
        ),
    );
    match reply {
        Err(zbus::Error::MethodError(name, _, _))
            if name.as_str() == "org.kde.KWin.ScreenShot2.Error.InvalidWindow" =>
        {
            Ok(())
        }
        Err(error) => Err(error.into()),
        Ok(_) => Err("KWin returned an unexpected permission-check result".into()),
    }
}

fn uint(metadata: &HashMap<String, OwnedValue>, key: &str) -> Result<usize> {
    Ok(u32::try_from(
        metadata
            .get(key)
            .ok_or_else(|| format!("Missing screenshot {key}"))?,
    )? as usize)
}

/// KWin supplies Qt ARGB32/RGB32 or RGBA8888 pixels. Refuse unknown formats.
fn rgba(raw: &[u8], width: usize, height: usize, stride: usize, format: usize) -> Result<Vec<u8>> {
    let row = width
        .checked_mul(4)
        .ok_or("Invalid screenshot dimensions")?;
    let bytes = stride
        .checked_mul(height)
        .ok_or("Invalid screenshot dimensions")?;
    if width == 0 || height == 0 || stride < row || bytes > MAX_BYTES || raw.len() != bytes {
        return Err("Invalid or oversized screenshot dimensions".into());
    }
    if ![4, 5, 6, 17, 18].contains(&format) {
        return Err("Unsupported KDE screenshot pixel format".into());
    }
    let mut output = Vec::with_capacity(row * height);
    for y in 0..height {
        for pixel in raw[y * stride..y * stride + row].chunks_exact(4) {
            let (r, g, b, a) = if format >= 17 {
                (pixel[0], pixel[1], pixel[2], pixel[3])
            } else {
                let argb = u32::from_ne_bytes(pixel.try_into()?);
                (
                    (argb >> 16) as u8,
                    (argb >> 8) as u8,
                    argb as u8,
                    if format == 4 { 255 } else { (argb >> 24) as u8 },
                )
            };
            let straight = |v: u8| {
                if a == 0 {
                    0
                } else {
                    ((v as u32 * 255 + a as u32 / 2) / a as u32).min(255) as u8
                }
            };
            if format == 6 || format == 18 {
                output.extend_from_slice(&[straight(r), straight(g), straight(b), a]);
            } else {
                output.extend_from_slice(&[r, g, b, a]);
            }
        }
    }
    Ok(output)
}

fn capture(connection: &Connection, directory: &Path) -> Result<Option<Window>> {
    // Lack of scripting only loses identity/text; never guess which app was captured.
    let before = window(connection, directory, None).ok().flatten();
    let proxy = screenshot_proxy(connection)?;
    let (mut reader, writer) = UnixStream::pair()?;
    reader.set_read_timeout(Some(DEADLINE))?;
    let options = HashMap::from([
        ("include-decoration", Value::from(true)),
        ("include-shadow", Value::from(false)),
        ("include-cursor", Value::from(false)),
        ("native-resolution", Value::from(true)),
    ]);
    let metadata: HashMap<String, OwnedValue> = if let Some(window) = &before {
        proxy.call(
            "CaptureWindow",
            &(&window.id, options, Fd::from(writer.as_fd())),
        )?
    } else {
        proxy.call("CaptureActiveWindow", &(options, Fd::from(writer.as_fd())))?
    };
    drop(writer);
    if <&str>::try_from(metadata.get("type").ok_or("Missing screenshot type")?)? != "raw" {
        return Err("Unsupported KDE screenshot type".into());
    }
    let width = uint(&metadata, "width")?;
    let height = uint(&metadata, "height")?;
    let stride = metadata
        .get("stride")
        .map(|_| uint(&metadata, "stride"))
        .transpose()?
        .unwrap_or(width.checked_mul(4).ok_or("Invalid width")?);
    let bytes = stride
        .checked_mul(height)
        .filter(|bytes| *bytes <= MAX_BYTES)
        .ok_or("Screenshot too large")?;
    let mut raw = vec![0; bytes];
    reader.read_exact(&mut raw)?;
    let pixels = rgba(&raw, width, height, stride, uint(&metadata, "format")?)?;
    let image = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(directory.join("capture.png"))?;
    let mut encoder = png::Encoder::new(image, width as u32, height as u32);
    encoder.set_color(png::ColorType::Rgba);
    encoder.set_depth(png::BitDepth::Eight);
    let mut image = encoder.write_header()?;
    image.write_image_data(&pixels)?;
    image.finish()?;
    let identity = before.filter(|before| {
        window(connection, directory, Some(&before.id))
            .ok()
            .flatten()
            .is_some_and(|current| before.matches_after_capture(&current))
    });
    Ok(identity)
}

fn run() -> Result<()> {
    let args: Vec<String> = std::env::args().collect();
    let connection = Builder::session()?.method_timeout(DEADLINE).build()?;
    match args.get(1).map(String::as_str) {
        Some("check") => {
            check(&connection)?;
            println!(
                "{}",
                serde_json::json!({"ready": true, "feedbackAvailable": feedback::supported(&connection)})
            );
        }
        Some("capture") if args.len() == 3 => {
            let window = capture(&connection, Path::new(&args[2]))?;
            println!("{}", serde_json::json!({"window": window}));
        }
        Some("feedback") if args.len() == 4 => {
            feedback::run(&connection, Path::new(&args[2]), &args[3])?;
        }
        Some("activate") if args.len() == 5 => {
            let pid: u32 = args[3].parse()?;
            let title = serde_json::to_string(&args[4])?;
            let value = script(
                &connection,
                Path::new(&args[2]),
                &format!(
                    "const targetPid = {pid}; const targetTitle = {title};\n{}",
                    include_str!("activate.js")
                ),
            )?;
            if serde_json::from_str::<serde_json::Value>(&value)?["activated"] != true {
                return Err("KDE could not identify the T3 Code window to activate".into());
            }
            println!("{{\"activated\":true}}");
        }
        _ => {
            return Err(
                "Expected check, capture <directory>, feedback <directory> <options>, or activate <directory> <pid> <title>"
                    .into(),
            );
        }
    }
    Ok(())
}

#[cfg(test)]
mod transport_tests;

fn main() {
    if let Err(error) = run() {
        let _ = writeln!(std::io::stderr(), "{error}");
        std::process::exit(1);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn window() -> Window {
        Window {
            id: "{window-1}".into(),
            title: "⠋ Document".into(),
            app_name: "Editor".into(),
            app_identifier: "editor".into(),
            process_id: 123,
            bounds: Bounds {
                x: 0,
                y: 0,
                width: 800,
                height: 600,
            },
            client_bounds: Bounds {
                x: 0,
                y: 29,
                width: 800,
                height: 571,
            },
        }
    }

    #[test]
    fn accepts_only_a_leading_cli_spinner_with_a_nonempty_label() {
        let before = window();
        for title in [
            "⠋ Document",
            "⠙ Document",
            "⠹ Document",
            "⠸ Document",
            "⠼ Document",
            "⠴ Document",
            "⠦ Document",
            "⠧ Document",
            "⠇ Document",
            "⠏ Document",
            "Document",
        ] {
            let mut after = before.clone();
            after.title = title.into();
            assert!(before.matches_after_capture(&after), "{title}");
            assert!(after.matches_after_capture(&before), "{title}");
        }
        for title in ["⠙ Other", "⠙Document", "⠁ Document", "Document ⠙", "⠙", ""] {
            let mut after = before.clone();
            after.title = title.into();
            assert!(!before.matches_after_capture(&after), "{title}");
        }
        assert_ne!(capture_title("⠋"), capture_title("⠙"));
    }

    #[test]
    fn spinner_normalization_does_not_weaken_window_identity_checks() {
        let before = window();
        for field in [
            "id",
            "process_id",
            "app_identifier",
            "app_name",
            "bounds",
            "client_bounds",
        ] {
            let mut after = before.clone();
            after.title = "⠙ Document".into();
            match field {
                "id" => after.id = "{window-2}".into(),
                "process_id" => after.process_id = 456,
                "app_identifier" => after.app_identifier = "other".into(),
                "app_name" => after.app_name = "Other".into(),
                "bounds" => after.bounds.width = 400,
                "client_bounds" => after.client_bounds.height = 500,
                _ => unreachable!(),
            }
            assert!(!before.matches_after_capture(&after), "{field}");
        }
    }

    #[test]
    fn decodes_native_argb_and_row_padding() {
        let raw = [0xff112233_u32.to_ne_bytes(), [0; 4]].concat();
        assert_eq!(rgba(&raw, 1, 1, 8, 5).unwrap(), [0x11, 0x22, 0x33, 255]);
    }
    #[test]
    fn unpremultiplies_transparency() {
        assert_eq!(
            rgba(&[64, 32, 16, 128], 1, 1, 4, 18).unwrap(),
            [128, 64, 32, 128]
        );
        assert_eq!(rgba(&[0, 0, 0, 0], 1, 1, 4, 18).unwrap(), [0, 0, 0, 0]);
    }
    #[test]
    fn rejects_invalid_dimensions_truncation_and_formats() {
        for (width, height, stride, format) in [
            (0, 1, 4, 5),
            (2, 1, 4, 5),
            (1, 2, 4, 5),
            (1, 1, 4, 99),
            (usize::MAX, 1, 4, 5),
        ] {
            assert!(rgba(&[0; 4], width, height, stride, format).is_err());
        }
    }
}