arib-cli 0.3.0

use std::fs;
use std::io::{BufRead, Read, Write, stdout};
use std::path::{Component, Path, PathBuf};

use anyhow::{Context, bail};
use arib::demux::Packet;
use arib::ts::carousel::Module;
use clap::Args;
use flate2::read::ZlibDecoder;
use tracing::warn;

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

/// Writes out the files of the data broadcasting under a directory, by the paths the
/// applications refer to them by, as they come round the carousel. A scrambled stream needs
/// descrambling first.
///
/// MPEG-2 TS writes each resource of a module to `<component tag>/<module ID>/<name>`, as a BML
/// document refers to it, and a module not in the entity format to `<component tag>/<module ID>`.
#[derive(Clone, Debug, Args)]
pub struct Options {
    #[command(flatten)]
    input: InputArgs,

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

    /// Reads the data broadcasting of the service rather than of the first one. MPEG-2 TS only.
    #[arg(long, value_parser = crate::print::parse_id)]
    service: Option<u16>,
}

/// The compression type of zlib, as the operational guidelines give it.
const ZLIB: u8 = 0x00;

pub fn run(options: &Options) -> anyhow::Result<()> {
    let input = options.input.open()?;
    match input.format {
        Format::Ts => read_ts(input.reader, options),
        Format::Mmt => read_mmt(input.reader, &options.output),
    }
}

fn read_ts(reader: impl BufRead, options: &Options) -> anyhow::Result<()> {
    let mut demuxer = arib::ts::demux::Demuxer::new(reader);
    if let Some(service_id) = options.service {
        demuxer = demuxer.for_service(service_id);
    }

    let mut carousel = arib::ts::carousel::Carousel::default();
    for packet in packets(demuxer) {
        let Some(module) = carousel.push(&packet?) else {
            continue;
        };
        if let Err(error) = write_module(&options.output, &module) {
            warn!(
                module = format!("{:02x}/{:04x}", module.component_tag, module.module_id),
                "{error:#}"
            );
        }
    }

    Ok(())
}

fn write_module(output: &Path, module: &Module) -> anyhow::Result<()> {
    let directory = format!("{:02x}/{:04x}", module.component_tag, module.module_id);
    let data = decompress(module.info.compression(), &module.data)?;

    match arib::ts::carousel::read_resources(&data)? {
        Some(resources) => {
            for resource in resources {
                let path = format!("{directory}/{}", resource.name);
                write(output, &path, &resource.data)?;
                print(&path, &resource.content_type, module.version)?;
            }
        }
        None => {
            write(output, &directory, &data)?;
            print(&directory, "", module.version)?;
        }
    }

    Ok(())
}

fn read_mmt(reader: impl BufRead, output: &Path) -> anyhow::Result<()> {
    let mut carousel = arib::mmt::application::Carousel::default();
    for packet in packets(arib::mmt::demux::Demuxer::new(reader)) {
        let Some(file) = carousel.push(&packet?) else {
            continue;
        };
        let written = decompress(file.compression, &file.data)
            .and_then(|data| write(output, &file.path, &data));
        if let Err(error) = written {
            warn!(path = file.path, "{error:#}");
            continue;
        }
        print(&file.path, &file.media_type, file.version)?;
    }

    Ok(())
}

/// The packets of the demultiplexer, warning about what could not be read and going on, but
/// stopping at the stream failing to be read.
fn packets(
    demuxer: impl Iterator<Item = arib::Result<Packet>>,
) -> impl Iterator<Item = anyhow::Result<Packet>> {
    demuxer.filter_map(|packet| match packet {
        Ok(packet) => Some(Ok(packet)),
        Err(arib::Error::Io(error)) => Some(Err(error.into())),
        Err(error) => {
            warn!(%error, "Could not read a packet");
            None
        }
    })
}

fn decompress(compression: Option<(u8, u32)>, data: &[u8]) -> anyhow::Result<Vec<u8>> {
    match compression {
        None => Ok(data.to_vec()),
        Some((ZLIB, original_size)) => {
            let mut decompressed = Vec::with_capacity(original_size as usize);
            ZlibDecoder::new(data)
                .read_to_end(&mut decompressed)
                .context("could not decompress")?;
            Ok(decompressed)
        }
        Some((compression_type, _)) => bail!("unknown compression type {compression_type:#04X}"),
    }
}

fn write(output: &Path, path: &str, data: &[u8]) -> anyhow::Result<()> {
    // The path comes from the broadcast, which is not to write anywhere but under the directory.
    let relative = Path::new(path);
    if !relative
        .components()
        .all(|component| matches!(component, Component::Normal(_)))
    {
        bail!("{path}: not a relative path under the directory");
    }

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

    Ok(())
}

fn print(path: &str, media_type: &str, version: u8) -> anyhow::Result<()> {
    writeln!(stdout(), "{path}\t{media_type}\tv{version}")?;
    Ok(())
}