1use serde::{Deserialize, Serialize};
2use tokio::process::Command;
3
4use crate::ffprobe_path;
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct ProbeResult {
8 pub format: FormatInfo,
9 pub streams: Vec<StreamInfo>,
10}
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct FormatInfo {
14 pub filename: String,
15 pub format_name: String,
16 pub format_long_name: String,
17 pub duration: f64, pub size: i64, pub bit_rate: i64, pub probe_score: i32,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct StreamInfo {
25 pub index: i32,
26 pub codec_name: String,
27 pub codec_long_name: String,
28 pub codec_type: String, pub profile: String,
30 pub width: i32,
31 pub height: i32,
32 pub pix_fmt: String,
33 pub level: i32,
34 pub field_order: String,
35 pub color_range: String,
36 pub color_space: String,
37 pub color_transfer: String,
38 pub color_primaries: String,
39 pub duration: f64, pub bit_rate: i64, pub nb_frames: i32,
42 pub r_frame_rate: String, pub avg_frame_rate: String, pub sample_rate: i32, pub channels: i32, pub channel_layout: String, pub bits_per_raw_sample: i32,
48}
49
50impl StreamInfo {
51 pub fn validate(&self) -> anyhow::Result<()> {
52 if self.codec_type != "video" {
53 anyhow::bail!("not a video stream (type={})", self.codec_type);
54 }
55 if self.width <= 0 || self.height <= 0 {
56 anyhow::bail!("invalid dimensions: {}x{}", self.width, self.height);
57 }
58 Ok(())
59 }
60
61 pub fn fps(&self) -> f64 {
63 parse_rational(&self.r_frame_rate)
64 }
65
66 pub fn resolution_str(&self) -> String {
67 if self.width == 0 || self.height == 0 {
68 String::new()
69 } else {
70 format!("{}x{}", self.width, self.height)
71 }
72 }
73
74 pub fn hdr_kind(&self) -> Option<&'static str> {
75 let transfer = self.color_transfer.to_ascii_lowercase();
76 if transfer == "smpte2084" {
77 return Some("PQ");
78 }
79 if transfer == "arib-std-b67" {
80 return Some("HLG");
81 }
82
83 let primaries_bt2020 = self.color_primaries.eq_ignore_ascii_case("bt2020");
84 let high_bit_depth = self.bits_per_raw_sample >= 10
85 || self.pix_fmt.contains("10")
86 || self.pix_fmt.contains("12")
87 || self.pix_fmt.contains("16");
88 if primaries_bt2020 && high_bit_depth {
89 return Some("BT.2020");
90 }
91
92 None
93 }
94
95 pub fn is_hdr(&self) -> bool {
96 self.hdr_kind().is_some()
97 }
98}
99
100impl FormatInfo {
101 pub fn duration_secs(&self) -> std::time::Duration {
102 std::time::Duration::from_secs_f64(self.duration)
103 }
104}
105
106impl ProbeResult {
107 pub fn video_stream(&self) -> Option<&StreamInfo> {
108 self.streams.iter().find(|s| s.codec_type == "video")
109 }
110
111 pub fn audio_stream(&self) -> Option<&StreamInfo> {
112 self.streams.iter().find(|s| s.codec_type == "audio")
113 }
114}
115
116pub async fn probe(path: &str) -> anyhow::Result<ProbeResult> {
118 let args = ["-v", "error", "-print_format", "json", "-show_format", "-show_streams", path];
119
120 let output = Command::new(ffprobe_path())
121 .args(args)
122 .stderr(std::process::Stdio::piped())
123 .output()
124 .await?;
125
126 if !output.status.success() {
127 let stderr = String::from_utf8_lossy(&output.stderr);
128 anyhow::bail!("ffprobe failed for {path}: {stderr}");
129 }
130
131 let raw: ProbeJsonRaw = serde_json::from_slice(&output.stdout)
132 .map_err(|e| anyhow::anyhow!("failed to parse ffprobe output: {e}"))?;
133
134 Ok(convert_probe(raw))
135}
136
137#[derive(Deserialize)]
139struct ProbeJsonRaw {
140 format: ProbeFormatRaw,
141 streams: Vec<ProbeStreamRaw>,
142}
143
144#[derive(Deserialize)]
145struct ProbeFormatRaw {
146 #[serde(default)]
147 filename: String,
148 #[serde(default)]
149 format_name: String,
150 #[serde(default)]
151 format_long_name: String,
152 #[serde(default)]
153 duration: String,
154 #[serde(default)]
155 size: String,
156 #[serde(default)]
157 bit_rate: String,
158 #[serde(default)]
159 probe_score: i32,
160}
161
162#[derive(Deserialize)]
163struct ProbeStreamRaw {
164 #[serde(default)]
165 index: i32,
166 #[serde(default)]
167 codec_name: String,
168 #[serde(default)]
169 codec_long_name: String,
170 #[serde(default)]
171 codec_type: String,
172 #[serde(default)]
173 profile: String,
174 #[serde(default)]
175 width: i32,
176 #[serde(default)]
177 height: i32,
178 #[serde(default)]
179 pix_fmt: String,
180 #[serde(default)]
181 level: i32,
182 #[serde(default)]
183 field_order: String,
184 #[serde(default)]
185 color_range: String,
186 #[serde(default)]
187 color_space: String,
188 #[serde(default)]
189 color_transfer: String,
190 #[serde(default)]
191 color_primaries: String,
192 #[serde(default)]
193 duration: String,
194 #[serde(default)]
195 bit_rate: String,
196 #[serde(default)]
197 nb_frames: String,
198 #[serde(default)]
199 r_frame_rate: String,
200 #[serde(default)]
201 avg_frame_rate: String,
202 #[serde(default)]
203 sample_rate: String,
204 #[serde(default)]
205 channels: i32,
206 #[serde(default)]
207 channel_layout: String,
208 #[serde(default)]
209 bits_per_raw_sample: String,
210}
211
212fn convert_probe(raw: ProbeJsonRaw) -> ProbeResult {
213 let format = FormatInfo {
214 filename: raw.format.filename,
215 format_name: raw.format.format_name,
216 format_long_name: raw.format.format_long_name,
217 duration: raw.format.duration.parse().unwrap_or(0.0),
218 size: raw.format.size.parse().unwrap_or(0),
219 bit_rate: raw.format.bit_rate.parse().unwrap_or(0),
220 probe_score: raw.format.probe_score,
221 };
222
223 let streams = raw
224 .streams
225 .into_iter()
226 .map(|s| StreamInfo {
227 index: s.index,
228 codec_name: s.codec_name,
229 codec_long_name: s.codec_long_name,
230 codec_type: s.codec_type,
231 profile: s.profile,
232 width: s.width,
233 height: s.height,
234 pix_fmt: s.pix_fmt,
235 level: s.level,
236 field_order: s.field_order,
237 color_range: s.color_range,
238 color_space: s.color_space,
239 color_transfer: s.color_transfer,
240 color_primaries: s.color_primaries,
241 duration: s.duration.parse().unwrap_or(0.0),
242 bit_rate: s.bit_rate.parse().unwrap_or(0),
243 nb_frames: s.nb_frames.parse().unwrap_or(0),
244 r_frame_rate: s.r_frame_rate,
245 avg_frame_rate: s.avg_frame_rate,
246 sample_rate: s.sample_rate.parse().unwrap_or(0),
247 channels: s.channels,
248 channel_layout: s.channel_layout,
249 bits_per_raw_sample: s.bits_per_raw_sample.parse().unwrap_or(0),
250 })
251 .collect();
252
253 ProbeResult { format, streams }
254}
255
256fn parse_rational(s: &str) -> f64 {
257 if let Some((num_s, den_s)) = s.split_once('/') {
258 let num: f64 = num_s.parse().unwrap_or(0.0);
259 let den: f64 = den_s.parse().unwrap_or(0.0);
260 if den != 0.0 { num / den } else { 0.0 }
261 } else {
262 s.parse().unwrap_or(0.0)
263 }
264}
265
266#[cfg(test)]
267mod tests {
268 use super::*;
269
270 fn video_stream() -> StreamInfo {
271 StreamInfo {
272 index: 0,
273 codec_name: "h264".into(),
274 codec_long_name: String::new(),
275 codec_type: "video".into(),
276 profile: String::new(),
277 width: 1920,
278 height: 1080,
279 pix_fmt: "yuv420p".into(),
280 level: 0,
281 field_order: String::new(),
282 color_range: String::new(),
283 color_space: String::new(),
284 color_transfer: String::new(),
285 color_primaries: String::new(),
286 duration: 0.0,
287 bit_rate: 0,
288 nb_frames: 0,
289 r_frame_rate: "24/1".into(),
290 avg_frame_rate: "24/1".into(),
291 sample_rate: 0,
292 channels: 0,
293 channel_layout: String::new(),
294 bits_per_raw_sample: 8,
295 }
296 }
297
298 #[test]
299 fn test_hdr_kind_detects_pq() {
300 let mut stream = video_stream();
301 stream.color_transfer = "smpte2084".into();
302 assert_eq!(stream.hdr_kind(), Some("PQ"));
303 assert!(stream.is_hdr());
304 }
305
306 #[test]
307 fn test_hdr_kind_detects_hlg() {
308 let mut stream = video_stream();
309 stream.color_transfer = "arib-std-b67".into();
310 assert_eq!(stream.hdr_kind(), Some("HLG"));
311 }
312
313 #[test]
314 fn test_hdr_kind_detects_bt2020_high_bit_depth() {
315 let mut stream = video_stream();
316 stream.color_primaries = "bt2020".into();
317 stream.pix_fmt = "yuv420p10le".into();
318 assert_eq!(stream.hdr_kind(), Some("BT.2020"));
319 }
320
321 #[test]
322 fn test_hdr_kind_ignores_sdr() {
323 assert_eq!(video_stream().hdr_kind(), None);
324 }
325}