arib-cli 0.2.0

Reads the signalling of ARIB broadcasts and descrambles them, as an example of the arib crate
use std::fs;
use std::io::{Write, stdout};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};

use arib::demux::Packet;
use arib::logo::{Logo, Logos, png_size};
use clap::Args;
use tracing::warn;

use crate::input::{Format, InputArgs};

/// Writes out the logos of the services under a directory, as PNGs with the palette they are
/// broadcast without filled in, to `<network>/<logo ID>_<logo type>.png`, and prints each: its
/// path, its size, its version and the services showing it.
#[derive(Clone, Debug, Args)]
pub struct Options {
    #[command(flatten)]
    input: InputArgs,

    /// The directory to write the logos under.
    #[arg(short, long)]
    output: PathBuf,

    /// Seconds to read the stream for, if it does not end before.
    #[arg(long, default_value_t = 60)]
    timeout: u64,
}

pub fn run(options: &Options) -> anyhow::Result<()> {
    let input = options.input.open()?;
    let deadline = Instant::now() + Duration::from_secs(options.timeout);
    match input.format {
        Format::Ts => read(
            arib::ts::demux::Demuxer::new(input.reader),
            &options.output,
            deadline,
        ),
        Format::Mmt => read(
            arib::mmt::demux::Demuxer::new(input.reader),
            &options.output,
            deadline,
        ),
    }
}

fn read(
    demuxer: impl Iterator<Item = arib::Result<Packet>>,
    output: &Path,
    deadline: Instant,
) -> anyhow::Result<()> {
    let mut logos = Logos::default();
    for packet in demuxer {
        let packet = match packet {
            Ok(packet) => packet,
            Err(arib::Error::Io(error)) => return Err(error.into()),
            Err(error) => {
                warn!(%error, "Could not read a packet");
                continue;
            }
        };
        for logo in logos.push(&packet) {
            write(output, &logo)?;
        }
        if Instant::now() >= deadline {
            break;
        }
    }

    Ok(())
}

fn write(output: &Path, logo: &Logo) -> anyhow::Result<()> {
    let path = format!(
        "{:04x}/{:03x}_{}.png",
        logo.original_network_id, logo.logo_id, logo.logo_type
    );
    let Some(png) = logo.png() else {
        warn!(path, "Ignoring a broken logo");
        return Ok(());
    };

    let file = output.join(&path);
    if let Some(parent) = file.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(file, &png)?;

    let (width, height) = png_size(&png).unwrap_or_default();
    let version = logo
        .logo_version
        .map_or_else(|| "-".to_owned(), |version| format!("v{version}"));
    let services: Vec<_> = logo
        .services
        .iter()
        .map(|service| format!("{:#06X}", service.service_id))
        .collect();
    writeln!(
        stdout(),
        "{path}\t{width}x{height}\t{version}\t{}",
        services.join(" ")
    )?;

    Ok(())
}