arib-cli 0.2.0

use std::collections::{BTreeMap, BTreeSet};
use std::io::{BufRead, Write, stdout};
use std::time::Duration;

use anyhow::bail;
use arib::demux::{MediaPacket, Packet, SignalingEvent, TrackType};
use arib::mmt::descriptor::{
    AdditionalAribSubtitleInfo, AdditionalDataComponentInfo, Descriptor, SubtitleTimeControlMode,
};
use arib::mmt::message::Message;
use arib::ts::caption::{DataGroup, DataGroupData};
use arib::ts::table::{Table, Tdt, Tot};
use chrono::{NaiveDateTime, TimeDelta};
use clap::Args;
use tracing::warn;

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

/// Prints the closed-caption and superimposition, a line for each: the time of day it is
/// presented at, in JST, and its text. A scrambled stream needs descrambling first.
#[derive(Clone, Debug, Args)]
pub struct Options {
    #[command(flatten)]
    input: InputArgs,

    /// Writes out the TTML documents as they are, one after another. MMT only.
    #[arg(long)]
    raw: bool,
}

pub fn run(options: &Options) -> anyhow::Result<()> {
    let input = options.input.open()?;
    let mut output = stdout().lock();

    match input.format {
        // The captions of MPEG-2 TS are no documents of their own, but data groups of the 8-bit
        // coding, which tell nothing apart once written one after another.
        Format::Ts if options.raw => bail!("--raw is for MMT only, whose captions are TTML"),
        Format::Ts => read_ts(input.reader, &mut output),
        Format::Mmt => read_mmt(input.reader, options.raw, &mut output),
    }
}

/// Times the captions by their PTS on the clock of the TOT and the TDT, through the PCR, which
/// unlike the video is never scrambled.
fn read_ts(reader: impl BufRead, output: &mut impl Write) -> anyhow::Result<()> {
    let mut demuxer = arib::ts::demux::Demuxer::new(reader);
    let mut captions = BTreeSet::new();
    let mut pcr = None;
    // The time of day the TOT or the TDT told last, and the PCR then.
    let mut clock: Option<(NaiveDateTime, Option<f64>)> = None;

    loop {
        let packet = demuxer.read_packet()?;
        // ponytail: any PCR is taken, which is off if the services keep clocks of their own.
        pcr = packet
            .as_ref()
            .and_then(|packet| packet.pcr())
            .map(|pcr| pcr as f64 / 90_000.0)
            .or(pcr);

        for packet in packets(demuxer.take_packets()) {
            match packet? {
                Packet::Media(MediaPacket::Track {
                    track_id,
                    ty: TrackType::AribCaption,
                }) => {
                    captions.insert(track_id);
                }
                Packet::Signaling(SignalingEvent::Ts {
                    table:
                        Table::Tot(Tot {
                            jst_time: Some(time),
                            ..
                        })
                        | Table::Tdt(Tdt {
                            jst_time: Some(time),
                            ..
                        }),
                    ..
                }) => clock = Some((time, pcr)),
                Packet::Media(MediaPacket::Sample {
                    track_id,
                    data,
                    pts,
                    ..
                }) if captions.contains(&track_id) => {
                    let group = match DataGroup::read_pes_data(data) {
                        Ok(group) => group,
                        Err(error) => {
                            warn!(%error, "Ignoring a malformed caption");
                            continue;
                        }
                    };
                    let DataGroupData::Statement { statement, .. } = group.data else {
                        continue;
                    };

                    // ponytail: the TOT tells whole seconds only, so the time is up to a second
                    // off; the STM of the real time mode and the clock wrapping round are not
                    // followed.
                    let time = clock.map(|(time, anchor)| match pts.zip(anchor) {
                        Some((pts, anchor)) => {
                            time + TimeDelta::milliseconds(((pts - anchor) * 1000.0).round() as i64)
                        }
                        // Superimposition comes without a PTS, to be shown as soon as it arrives.
                        None => time,
                    });
                    print_caption(output, time.map(When::At), &statement.text())?;
                }
                _ => {}
            }
        }

        if packet.is_none() {
            return Ok(());
        }
    }
}

/// Times the paragraphs of the TTML documents as the MH-data component descriptor of their asset
/// says.
fn read_mmt(reader: impl BufRead, raw: bool, output: &mut impl Write) -> anyhow::Result<()> {
    let mut tracks = BTreeMap::<u16, Option<AdditionalAribSubtitleInfo>>::new();

    for packet in packets(arib::mmt::demux::Demuxer::new(reader)) {
        match packet? {
            Packet::Media(MediaPacket::Track {
                track_id,
                ty: TrackType::Ttml,
            }) => {
                tracks.entry(track_id).or_default();
            }
            Packet::Signaling(SignalingEvent::Mmt(Message::Pa(message))) => {
                let assets = message.tables.iter().flat_map(|table| match table {
                    arib::mmt::table::Table::Mpt(mpt) => mpt.assets.as_slice(),
                    _ => &[],
                });
                for asset in assets {
                    let Some(info) = asset
                        .locations
                        .last()
                        .and_then(|location| tracks.get_mut(&location.packet_id()?))
                    else {
                        continue;
                    };
                    for descriptor in &asset.asset_descriptors {
                        if let Descriptor::MhDataComponent(descriptor) = descriptor
                            && let AdditionalDataComponentInfo::Subtitle(subtitle) =
                                &descriptor.additional_data_component_info
                        {
                            *info = Some(subtitle.clone());
                        }
                    }
                }
            }
            Packet::Media(MediaPacket::Sample {
                track_id,
                data,
                pts,
                ..
            }) => {
                let Some(info) = tracks.get(&track_id) else {
                    continue;
                };
                if raw {
                    output.write_all(&data)?;
                    if !data.ends_with(b"\n") {
                        output.write_all(b"\n")?;
                    }
                    output.flush()?;
                    continue;
                }

                let cues = match std::str::from_utf8(&data)
                    .map_err(std::io::Error::other)
                    .and_then(arib::mmt::ttml::read_cues)
                {
                    Ok(cues) => cues,
                    Err(error) => {
                        warn!(%error, "Ignoring a malformed TTML document");
                        continue;
                    }
                };

                // ponytail: the times from the start of the event, from the NPT and of the day
                // in UTC are printed as they are, not put on the clock.
                let origin = info
                    .as_ref()
                    .and_then(|info| match info.time_control_mode {
                        SubtitleTimeControlMode::ReferenceStartTime => {
                            info.reference_start_time.map(|time| {
                                (time >> 32) as f64 + (time as u32) as f64 / 2_f64.powi(32)
                            })
                        }
                        SubtitleTimeControlMode::MpuTimestamp => pts,
                        _ => None,
                    })
                    .map(ntp_to_jst);
                for cue in cues {
                    let when = match (origin, cue.begin) {
                        (Some(origin), begin) => Some(When::At(origin + begin.unwrap_or_default())),
                        (None, Some(begin)) => Some(When::After(begin)),
                        (None, None) => None,
                    };
                    print_caption(output, when, &cue.text)?;
                }
            }
            _ => {}
        }
    }

    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
        }
    })
}

/// When a caption is presented: at a time of day, or only a time from where it counts from, which
/// the stream does not tell.
enum When {
    At(NaiveDateTime),
    After(Duration),
}

fn print_caption(output: &mut impl Write, when: Option<When>, text: &str) -> anyhow::Result<()> {
    let text = text
        .split('\n')
        .map(str::trim)
        .collect::<Vec<_>>()
        .join(" ");
    let text = text.trim();
    if text.is_empty() {
        return Ok(());
    }

    let when = match when {
        Some(When::At(time)) => time.format("%Y-%m-%d %H:%M:%S%.3f").to_string(),
        Some(When::After(time)) => {
            let millis = time.as_millis();
            format!(
                "+{:02}:{:02}:{:02}.{:03}",
                millis / 3_600_000,
                millis / 60_000 % 60,
                millis / 1000 % 60,
                millis % 1000,
            )
        }
        None => "-".to_owned(),
    };
    writeln!(output, "{when}\t{text}")?;
    output.flush()?;

    Ok(())
}

/// The time of day in JST of an NTP time in seconds.
fn ntp_to_jst(seconds: f64) -> NaiveDateTime {
    const NTP_TO_UNIX: f64 = 2_208_988_800.0;
    const JST: TimeDelta = TimeDelta::hours(9);

    let millis = ((seconds - NTP_TO_UNIX) * 1000.0).round() as i64;
    chrono::DateTime::from_timestamp_millis(millis)
        .unwrap_or_default()
        .naive_utc()
        + JST
}