use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use anyhow::{Context, Result};
use pinray::{SourceId, VideoCaptureTarget};
use tracing::field::{Field, Visit};
use tracing_subscriber::Layer;
use tracing_subscriber::layer::SubscriberExt;
use super::{Capture, session};
fn token_path(slot: u32) -> Option<PathBuf> {
let name = if slot <= 1 {
"restore_token".to_owned()
} else {
format!("restore_token_{slot}")
};
Some(crate::prefs::state_dir()?.join(name))
}
fn load_token(slot: u32) -> Option<String> {
let path = token_path(slot)?;
let token = std::fs::read_to_string(&path)
.ok()
.or_else(|| {
if slot > 1 {
return None;
}
let legacy = path.parent()?.parent()?.join("screencap/restore_token");
std::fs::read_to_string(legacy).ok()
})?;
let token = token.trim().to_owned();
(!token.is_empty()).then_some(token)
}
fn store_token(slot: u32, token: &str) {
let Some(path) = token_path(slot) else { return };
if let Some(dir) = path.parent() {
let _ = std::fs::create_dir_all(dir);
}
if let Err(err) = std::fs::write(&path, token) {
eprintln!("warning: could not persist portal restore token: {err}");
}
}
#[derive(Default)]
struct TokenVisitor {
token: Option<String>,
}
impl Visit for TokenVisitor {
fn record_str(&mut self, field: &Field, value: &str) {
if field.name() == "restore_token" {
self.token = Some(value.to_owned());
}
}
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
if field.name() == "restore_token" {
self.token = Some(format!("{value:?}").trim_matches('"').to_owned());
}
}
}
struct TokenCatcher {
slot: Arc<Mutex<Option<String>>>,
}
impl<S: tracing::Subscriber> Layer<S> for TokenCatcher {
fn on_event(
&self,
event: &tracing::Event<'_>,
_ctx: tracing_subscriber::layer::Context<'_, S>,
) {
let mut visitor = TokenVisitor::default();
event.record(&mut visitor);
if let Some(token) = visitor.token {
*self.slot.lock().expect("token slot poisoned") = Some(token);
}
}
}
pub fn capture_screenshot(embed_cursor: bool, pick_screen: bool, slot: u32) -> Result<Capture> {
let caught = Arc::new(Mutex::new(None));
let subscriber = tracing_subscriber::registry().with(TokenCatcher { slot: caught.clone() });
let _guard = tracing::subscriber::set_default(subscriber);
let mut builder = session::builder(embed_cursor)
.video_target(VideoCaptureTarget::Display(SourceId::new("auto")));
let mut used_token = false;
if !pick_screen && let Some(token) = load_token(slot) {
builder = builder.restore_token(token);
used_token = true;
}
let mut session = builder.build().context("building capture session")?;
session.start().context(
"starting capture (is the app running with access to the session D-Bus and PipeWire?)",
)?;
let settle = if used_token {
Duration::from_millis(150)
} else {
Duration::from_millis(600)
};
let image = session::take_frame(&mut session, settle);
session.stop().ok();
if let Some(token) = caught.lock().expect("token slot poisoned").take() {
store_token(slot, &token);
}
Ok(Capture { image: image?, display: None })
}