refactor import script

3 files changed, 222 insertions, 136 deletions

diff --git a/remuxer/src/import/mod.rs b/remuxer/src/import/mod.rs
index 7b0b371..4c5275e 100644
--- a/remuxer/src/import/mod.rs
+++ b/remuxer/src/import/mod.rs
@@ -3,29 +3,24 @@
     which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
     Copyright (C) 2023 metamuffin <metamuffin.org>
 */
+pub mod seek_index;
+
 use anyhow::{anyhow, bail, Result};
-use jellycommon::{BlockIndex, LocalTrack, SeekIndex, SourceTrack, SourceTrackKind};
+use jellycommon::{LocalTrack, SourceTrack, SourceTrackKind};
 use jellymatroska::{
-    block::Block,
     matroska::MatroskaTag,
     read::EbmlReader,
     unflatten::{IterWithPos, Unflat, Unflatten},
 };
-use log::{debug, error, info, trace, warn};
-use std::{collections::BTreeMap, path::PathBuf};
+use log::{debug, error, info, warn};
+use std::path::PathBuf;
 
-pub fn import_read(
+pub fn import_metadata(
     path: &PathBuf,
     input: &mut EbmlReader,
-) -> Result<(
-    Vec<SourceTrack>,
-    Vec<LocalTrack>,
-    BTreeMap<u64, SeekIndex>,
-    f64,
-)> {
+) -> Result<(Vec<SourceTrack>, Vec<LocalTrack>, f64)> {
     let mut iteminfo = Vec::new();
     let mut private = Vec::new();
-    let mut seek_index = BTreeMap::new();
     let mut dur = None;
     while let Some(item) = input.next() {
         let item = match item {
@@ -60,7 +55,6 @@ pub fn import_read(
                     &mut children,
                     &mut iteminfo,
                     &mut private,
-                    &mut seek_index,
                 )?);
                 info!("segment end");
             }
@@ -68,7 +62,7 @@ pub fn import_read(
         }
     }
 
-    Ok((iteminfo, private, seek_index, dur.unwrap_or(0.)))
+    Ok((iteminfo, private, dur.unwrap_or(0.)))
 }
 
 fn import_read_segment(
@@ -76,7 +70,6 @@ fn import_read_segment(
     segment: &mut Unflatten,
     iteminfo: &mut Vec<SourceTrack>,
     private: &mut Vec<LocalTrack>,
-    seek_index: &mut BTreeMap<u64, SeekIndex>,
 ) -> Result<Option<f64>> {
     let (mut timestamp_scale, mut duration) = (None, None);
 
@@ -193,71 +186,7 @@ fn import_read_segment(
                     }
                 }
             }
-            MatroskaTag::Cluster(_) => {
-                let mut children = children.unwrap();
-                let mut pts = 0;
-                let mut position = children.position();
-
-                loop {
-                    if let Some(Ok(Unflat { children, item, .. })) = children.n() {
-                        match item {
-                            MatroskaTag::Timestamp(ts) => pts = ts,
-                            MatroskaTag::BlockGroup(_) => {
-                                debug!("group");
-                                let mut children = children.unwrap();
-                                // let position = children.position(); //? TODO where should this point to? cluster or block? // probably block
-                                while let Some(Ok(Unflat {
-                                    children: _,
-                                    item,
-                                    position,
-                                })) = children.n()
-                                {
-                                    match item {
-                                        MatroskaTag::Block(ref buf) => {
-                                            let block = Block::parse(buf)?;
-                                            debug!(
-                                                "block: track={} tso={}",
-                                                block.track, block.timestamp_off
-                                            );
-                                            seek_index
-                                                .entry(block.track)
-                                                .or_insert(SeekIndex { blocks: vec![] })
-                                                .blocks
-                                                .push(BlockIndex {
-                                                    pts: pts + block.timestamp_off as u64,
-                                                    source_off: position,
-                                                    size: block.data.len(),
-                                                });
-                                        }
-                                        _ => trace!("{item:?}"),
-                                    }
-                                }
-                            }
-                            MatroskaTag::SimpleBlock(buf) => {
-                                let block = Block::parse(&buf)?;
-                                debug!(
-                                    "simple block: track={} tso={}",
-                                    block.track, block.timestamp_off
-                                );
-                                debug!("{pts} {}", block.timestamp_off);
-                                seek_index
-                                    .entry(block.track)
-                                    .or_insert(SeekIndex { blocks: vec![] })
-                                    .blocks
-                                    .push(BlockIndex {
-                                        pts: (pts as i64 + block.timestamp_off as i64) as u64,
-                                        source_off: position,
-                                        size: block.data.len(),
-                                    });
-                            }
-                            _ => debug!("(rsc) tag ignored: {item:?}"),
-                        }
-                    } else {
-                        break;
-                    }
-                    position = children.position();
-                }
-            }
+            MatroskaTag::Cluster(_) => {}
             _ => debug!("(rs) tag ignored: {item:?}"),
         };
     }
diff --git a/remuxer/src/import/seek_index.rs b/remuxer/src/import/seek_index.rs
new file mode 100644
index 0000000..363a12f
--- /dev/null
+++ b/remuxer/src/import/seek_index.rs
@@ -0,0 +1,116 @@
+use anyhow::Result;
+use jellycommon::{BlockIndex, SeekIndex};
+use jellymatroska::{
+    block::Block,
+    read::EbmlReader,
+    unflatten::{IterWithPos, Unflat, Unflatten},
+    MatroskaTag,
+};
+use log::{debug, info, trace, warn};
+use std::collections::BTreeMap;
+
+pub fn import_seek_index(input: &mut EbmlReader) -> Result<BTreeMap<u64, SeekIndex>> {
+    let mut seek_index = BTreeMap::new();
+    while let Some(item) = input.next() {
+        let item = match item {
+            Ok(item) => item,
+            Err(e) => {
+                warn!("{e}");
+                break;
+            }
+        };
+        match item {
+            MatroskaTag::Segment(_) => {
+                info!("segment start");
+                let mut children = Unflatten::new_with_end(input, item);
+                import_seek_index_segment(&mut children, &mut seek_index)?;
+                info!("segment end");
+            }
+            _ => debug!("(r) tag ignored: {item:?}"),
+        }
+    }
+    Ok(seek_index)
+}
+
+fn import_seek_index_segment(
+    segment: &mut Unflatten,
+    seek_index: &mut BTreeMap<u64, SeekIndex>,
+) -> Result<()> {
+    while let Some(Ok(Unflat { children, item, .. })) = segment.n() {
+        match item {
+            MatroskaTag::SeekHead(_) => {}
+            MatroskaTag::Info(_) => {}
+            MatroskaTag::Tags(_) => {}
+            MatroskaTag::Cues(_) => {}
+            MatroskaTag::Chapters(_) => {}
+            MatroskaTag::Tracks(_) => {}
+            MatroskaTag::Cluster(_) => {
+                let mut children = children.unwrap();
+                let mut pts = 0;
+                let mut position = children.position();
+
+                loop {
+                    if let Some(Ok(Unflat { children, item, .. })) = children.n() {
+                        match item {
+                            MatroskaTag::Timestamp(ts) => pts = ts,
+                            MatroskaTag::BlockGroup(_) => {
+                                debug!("group");
+                                let mut children = children.unwrap();
+                                // let position = children.position(); //? TODO where should this point to? cluster or block? // probably block
+                                while let Some(Ok(Unflat {
+                                    children: _,
+                                    item,
+                                    position,
+                                })) = children.n()
+                                {
+                                    match item {
+                                        MatroskaTag::Block(ref buf) => {
+                                            let block = Block::parse(buf)?;
+                                            debug!(
+                                                "block: track={} tso={}",
+                                                block.track, block.timestamp_off
+                                            );
+                                            seek_index
+                                                .entry(block.track)
+                                                .or_insert(SeekIndex { blocks: vec![] })
+                                                .blocks
+                                                .push(BlockIndex {
+                                                    pts: pts + block.timestamp_off as u64,
+                                                    source_off: position,
+                                                    size: block.data.len(),
+                                                });
+                                        }
+                                        _ => trace!("{item:?}"),
+                                    }
+                                }
+                            }
+                            MatroskaTag::SimpleBlock(buf) => {
+                                let block = Block::parse(&buf)?;
+                                debug!(
+                                    "simple block: track={} tso={}",
+                                    block.track, block.timestamp_off
+                                );
+                                debug!("{pts} {}", block.timestamp_off);
+                                seek_index
+                                    .entry(block.track)
+                                    .or_insert(SeekIndex { blocks: vec![] })
+                                    .blocks
+                                    .push(BlockIndex {
+                                        pts: (pts as i64 + block.timestamp_off as i64) as u64,
+                                        source_off: position,
+                                        size: block.data.len(),
+                                    });
+                            }
+                            _ => debug!("(rsc) tag ignored: {item:?}"),
+                        }
+                    } else {
+                        break;
+                    }
+                    position = children.position();
+                }
+            }
+            _ => debug!("(rs) tag ignored: {item:?}"),
+        };
+    }
+    Ok(())
+}
diff --git a/tools/src/bin/import.rs b/tools/src/bin/import.rs
index 471a669..eada6ad 100644
--- a/tools/src/bin/import.rs
+++ b/tools/src/bin/import.rs
@@ -5,12 +5,13 @@
 */
 use anyhow::Context;
 use clap::{Parser, Subcommand};
-use jellycommon::{AssetLocation, MediaInfo, Node, NodeKind, NodePrivate, NodePublic};
+use jellycommon::{AssetLocation, MediaInfo, MediaSource, Node, NodeKind, NodePrivate, NodePublic};
 use jellymatroska::read::EbmlReader;
-use jellyremuxer::import::import_read;
+use jellyremuxer::import::{import_metadata, seek_index::import_seek_index};
 use jellytools::tmdb::{tmdb_details, tmdb_image, tmdb_search};
 use log::info;
 use std::{
+    collections::BTreeMap,
     fs::{remove_file, File},
     io::stdin,
     path::PathBuf,
@@ -29,9 +30,10 @@ struct Args {
 enum Action {
     Create {
         path: PathBuf,
-        title: Option<String>,
+        #[arg(short = 't', long)]
+        tmdb_search: Option<String>,
         #[arg(short = 'T', long)]
-        tmdb: Option<String>,
+        tmdb_id: Option<String>,
         #[arg(long)]
         copy: bool,
         #[arg(long)]
@@ -69,21 +71,20 @@ fn main() -> anyhow::Result<()> {
     match args.action {
         Action::Create {
             path,
-            title,
-            tmdb: id,
+            tmdb_id,
+            tmdb_search,
             input,
             series,
             copy,
             r#move,
         } => {
-            assert!(series || input.is_some(), "series or input required");
-            let kind = if series { "tv" } else { "movie" };
-            let key = std::env::var("TMDB_API_KEY").context("tmdb api key required")?;
-            let id = if let Some(id) = id {
-                id.parse().unwrap()
+            let tmdb_kind = if series { "tv" } else { "movie" };
+            let tmdb_key = std::env::var("TMDB_API_KEY").context("tmdb api key required")?;
+            let tmdb_id = if let Some(id) = tmdb_id {
+                Some(id.parse().unwrap())
             } else {
-                let title = title.as_ref().unwrap();
-                let results = tmdb_search(kind, title, &key)?;
+                let title = tmdb_search.as_ref().unwrap();
+                let results = crate::tmdb_search(tmdb_kind, title, &tmdb_key)?;
                 info!("results:");
                 for (i, r) in results.results.iter().enumerate() {
                     info!(
@@ -104,43 +105,67 @@ fn main() -> anyhow::Result<()> {
                 } else {
                     0
                 };
-                results.results[res_index].id
+                Some(results.results[res_index].id)
             };
 
-            let details = tmdb_details(kind, id, &key).context("fetching details")?;
-            let ident = make_ident(details.title.as_ref().or(details.name.as_ref()).unwrap());
+            let tmdb_details = tmdb_id.map(|id| {
+                let td = tmdb_details(tmdb_kind, id, &tmdb_key)
+                    .context("fetching details")
+                    .unwrap();
+                if td.title.is_some() {
+                    info!("is this correct? [y/n]");
+                    if stdin().lines().next().unwrap().unwrap() != "y" {
+                        exit(0)
+                    }
+                }
+                td
+            });
+
+            let title = tmdb_details
+                .as_ref()
+                .map(|d| d.title.clone().or(d.name.clone()))
+                .flatten()
+                .unwrap();
+            let ident = make_ident(&title);
             let path = path.join(&ident);
-            std::fs::create_dir_all(&path)?;
 
-            let poster = details
-                .poster_path
-                .map(|p| {
-                    let pu = path.join("poster.jpeg");
-                    let mut f = File::create(&pu)?;
-                    tmdb_image(&p, &mut f)?;
-                    Ok::<_, anyhow::Error>(pu)
-                })
-                .transpose()?;
-            let backdrop = details
-                .backdrop_path
-                .map(|p| {
-                    let pu = path.join("backdrop.jpeg");
-                    let mut f = File::create(&pu)?;
-                    tmdb_image(&p, &mut f)?;
-                    Ok::<_, anyhow::Error>(pu)
+            let poster = tmdb_details
+                .as_ref()
+                .map(|d| {
+                    d.poster_path
+                        .as_ref()
+                        .map(|p| {
+                            let pu = path.join("poster.jpeg");
+                            let mut f = File::create(&pu)?;
+                            tmdb_image(&p, &mut f)?;
+                            Ok::<_, anyhow::Error>(pu)
+                        })
+                        .transpose()
                 })
-                .transpose()?;
+                .transpose()?
+                .flatten();
 
-            if title.is_some() {
-                info!("is this correct? [y/n]");
-                if stdin().lines().next().unwrap().unwrap() != "y" {
-                    exit(0)
-                }
-            }
+            let backdrop = tmdb_details
+                .as_ref()
+                .map(|d| {
+                    d.backdrop_path
+                        .as_ref()
+                        .map(|p| {
+                            let pu = path.join("backdrop.jpeg");
+                            let mut f = File::create(&pu)?;
+                            tmdb_image(&p, &mut f)?;
+                            Ok::<_, anyhow::Error>(pu)
+                        })
+                        .transpose()
+                })
+                .transpose()?
+                .flatten();
 
             let kind;
             let media;
             let source;
+            let mut seek_index = BTreeMap::new();
+            let mut source_path_e = None;
 
             if let Some(input) = input {
                 let source_path = path.join(format!("source.mkv"));
@@ -154,25 +179,24 @@ fn main() -> anyhow::Result<()> {
                     }
                     std::os::unix::fs::symlink(&input, &source_path)?;
                 }
-                let input = File::open(&source_path).unwrap();
-                let mut input = EbmlReader::new(input);
-                let (tracks, local_tracks, seek_index, duration) =
-                    import_read(&source_path.to_path_buf(), &mut input)?;
 
-                for (tn, index) in seek_index {
-                    info!("writing index {tn} with {} blocks", index.blocks.len());
-                    bincode::encode_into_std_write(
-                        index,
-                        &mut File::create(source_path.with_extension(&format!("si.{tn}")))?,
-                        bincode::config::standard(),
-                    )?;
-                }
+                let (tracks, local_tracks, duration) = {
+                    let input = File::open(&source_path).unwrap();
+                    let mut input = EbmlReader::new(input);
+                    import_metadata(&source_path.to_path_buf(), &mut input)?
+                };
+                seek_index = {
+                    let input = File::open(&source_path).unwrap();
+                    let mut input = EbmlReader::new(input);
+                    import_seek_index(&mut input)?
+                };
 
                 kind = NodeKind::Movie;
                 media = Some(MediaInfo { duration, tracks });
-                source = Some(jellycommon::MediaSource::Local {
+                source = Some(MediaSource::Local {
                     tracks: local_tracks,
                 });
+                source_path_e = Some(source_path)
             } else {
                 kind = NodeKind::Series;
                 media = None;
@@ -189,9 +213,12 @@ fn main() -> anyhow::Result<()> {
                 public: NodePublic {
                     parent: None,
                     federated: None,
-                    description: Some(details.overview),
-                    tagline: details.tagline,
-                    title: details.title.clone().or(details.name.clone()).unwrap(),
+                    description: tmdb_details.as_ref().map(|d| d.overview.to_owned()),
+                    tagline: tmdb_details
+                        .as_ref()
+                        .map(|d| d.tagline.to_owned())
+                        .flatten(),
+                    title,
                     index: None,
                     kind,
                     children: Vec::new(),
@@ -202,6 +229,20 @@ fn main() -> anyhow::Result<()> {
             if args.dry {
                 println!("{node:?}")
             } else {
+                std::fs::create_dir_all(&path)?;
+                for (tn, index) in seek_index {
+                    info!("writing index {tn} with {} blocks", index.blocks.len());
+                    bincode::encode_into_std_write(
+                        index,
+                        &mut File::create(
+                            source_path_e
+                                .as_ref()
+                                .unwrap()
+                                .with_extension(&format!("si.{tn}")),
+                        )?,
+                        bincode::config::standard(),
+                    )?;
+                }
                 let f = File::create(path.join(if series {
                     "directory.json".to_string()
                 } else {