shinkai-translator 0.1.3

mod tui;

use std::path::{Path, PathBuf};
use std::sync::Arc;

use clap::{Args, Parser, Subcommand};
use tempfile::tempdir;

use shinkai_translator::config::{AppConfig, LoadedConfig};
use shinkai_translator::formats;
use shinkai_translator::media::{
    count_subtitle_streams, detect_input_kind, extract_pgs_stream, is_supported_input_path,
    mux_translated_subtitle, prepare_video_subtitle_job,
};
use shinkai_translator::ocr::{cleanup_fragmentary_subtitle_document, debug_pgs_ocr};
use shinkai_translator::{
    AssClassificationPolicy, CueDisposition, ExternalToolConfig, MediaInputKind,
    OpenAiCompatibleProvider, PgsOcrConfig, PgsOcrLanguage, ProviderConfig,
    SelectedSubtitleStream, SubtitleFormat, ThinkingMode, TranslationOptions, Translator,
    TranslatorError,
};

#[derive(Debug, Parser)]
#[command(
    author,
    version,
    about = "Translate video subtitles by extracting, translating, and muxing subtitle tracks"
)]
struct Cli {
    #[arg(long, global = true)]
    config: Option<PathBuf>,
    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    #[command(about = "Translate using config defaults; advanced overrides remain available but hidden")]
    Translate(TranslateArgs),
    #[command(about = "Validate the selected subtitle stream using config defaults")]
    Validate(ValidateArgs),
    #[command(about = "Inspect PGS subtitle OCR at specific timestamps for debugging")]
    DebugOcr(DebugOcrArgs),
    #[command(about = "Open the interactive config editor")]
    Config,
}

#[derive(Debug, Args, Clone)]
struct TranslateArgs {
    #[arg()]
    input: PathBuf,
    #[arg(long, hide = true)]
    subtitle_input: Option<PathBuf>,
    #[arg(long, short)]
    output: Option<PathBuf>,
    #[arg(long, hide = true)]
    subtitle_output: Option<PathBuf>,
    #[arg(long, hide = true)]
    format: Option<String>,
    #[arg(long, hide = true)]
    subtitle_stream_index: Option<usize>,
    #[arg(long, hide = true)]
    source_language: Option<String>,
    #[arg(long, hide = true)]
    target_language: Option<String>,
    #[arg(long, hide = true)]
    model: Option<String>,
    #[arg(long, hide = true)]
    base_url: Option<String>,
    #[arg(long, env = "SHINKAI_TRANSLATOR_API_KEY", hide = true)]
    api_key: Option<String>,
    #[arg(long, hide = true)]
    thinking_mode: Option<String>,
    #[arg(long, hide = true)]
    max_batch_items: Option<usize>,
    #[arg(long, hide = true)]
    max_batch_characters: Option<usize>,
    #[arg(long, hide = true)]
    max_parallel_batches: Option<usize>,
    #[arg(long, hide = true)]
    timeout_seconds: Option<u64>,
    #[arg(long, hide = true)]
    max_retries: Option<u32>,
    #[arg(long, hide = true)]
    system_prompt: Option<String>,
    #[arg(long)]
    overwrite: bool,
    #[arg(long)]
    dry_run: bool,
    #[command(flatten)]
    ocr: PgsOcrArgs,
    #[command(flatten)]
    classification: AssClassificationArgs,
}

#[derive(Debug, Args)]
struct ValidateArgs {
    #[arg()]
    input: PathBuf,
    #[arg(long, hide = true)]
    subtitle_input: Option<PathBuf>,
    #[arg(long, hide = true)]
    format: Option<String>,
    #[arg(long, hide = true)]
    subtitle_stream_index: Option<usize>,
    #[command(flatten)]
    ocr: PgsOcrArgs,
    #[command(flatten)]
    classification: AssClassificationArgs,
}

#[derive(Debug, Args, Clone)]
struct PgsOcrArgs {
    #[arg(long, hide = true)]
    ocr_language: Option<String>,
    #[arg(long, hide = true)]
    ocr_model_dir: Option<PathBuf>,
}

#[derive(Debug, Args)]
struct DebugOcrArgs {
    /// Video file or PGS .sup file to debug-inspect
    #[arg()]
    input: PathBuf,
    /// Subtitle stream index to extract (when input is a video)
    #[arg(long)]
    subtitle_stream_index: Option<usize>,
    /// Center timestamp to inspect, e.g. 01:23:45 or 01:23:45.500
    #[arg(long, conflicts_with_all = ["from", "to"])]
    at: Option<String>,
    /// Start of time range to inspect, e.g. 01:23:00
    #[arg(long, requires = "to")]
    from: Option<String>,
    /// End of time range to inspect, e.g. 01:25:00
    #[arg(long, requires = "from")]
    to: Option<String>,
    /// Half-window in seconds around --at timestamp (default 30 → scans from at-30s to at+30s)
    #[arg(long, default_value = "30")]
    window: u64,
    /// Output directory for debug files (images + OCR text reports)
    #[arg(long, short, default_value = "./ocr-debug")]
    output: PathBuf,
    #[command(flatten)]
    ocr: PgsOcrArgs,
}

#[derive(Debug, Args, Clone)]
struct AssClassificationArgs {    #[arg(long, hide = true)]
    karaoke_policy: Option<String>,
    #[arg(long, hide = true)]
    explicit_song_policy: Option<String>,
    #[arg(long, hide = true)]
    inferred_song_policy: Option<String>,
    #[arg(long = "song-style-marker", hide = true)]
    style_markers: Vec<String>,
    #[arg(long = "song-effect-marker", hide = true)]
    effect_markers: Vec<String>,
    #[arg(long = "song-name-marker", hide = true)]
    name_markers: Vec<String>,
    #[arg(long, hide = true)]
    disable_inferred_song_detection: bool,
    #[arg(long, hide = true)]
    min_inferred_song_run_length: Option<usize>,
    #[arg(long = "report", alias = "classification-report", hide = true)]
    classification_report: Option<PathBuf>,
}

#[tokio::main]
async fn main() {
    if let Err(error) = run().await {
        eprintln!("error: {error}");
        std::process::exit(1);
    }
}

async fn run() -> Result<(), TranslatorError> {
    let cli = Cli::parse();
    let mut loaded_config = LoadedConfig::load(cli.config.as_deref())?;

    match cli.command {
        Command::Translate(args) => translate(args, &loaded_config.data).await,
        Command::Validate(args) => validate(args, &loaded_config.data).await,
        Command::DebugOcr(args) => debug_ocr(args, &loaded_config.data).await,
        Command::Config => tui::run_config_tui(&mut loaded_config),
    }
}

async fn translate(args: TranslateArgs, config: &AppConfig) -> Result<(), TranslatorError> {
    if args.input.is_dir() {
        if args.output.is_some() {
            return Err(TranslatorError::InvalidConfig(
                "--output cannot be used with a directory input; output paths are derived automatically"
                    .to_owned(),
            ));
        }
        let files = collect_translatable_files(&args.input);
        if files.is_empty() {
            eprintln!("no video or subtitle files found in {}", args.input.display());
            return Ok(());
        }
        let mut had_error = false;
        for path in files {
            println!("translating {} ...", path.display());
            let mut file_args = args.clone();
            file_args.input = path;
            if let Err(e) = translate_single(file_args, config).await {
                eprintln!("error: {e}");
                had_error = true;
            }
        }
        return if had_error {
            Err(TranslatorError::InvalidConfig(
                "batch translation: one or more files failed".to_owned(),
            ))
        } else {
            Ok(())
        };
    }
    translate_single(args, config).await
}

fn collect_translatable_files(dir: &Path) -> Vec<PathBuf> {
    let Ok(read_dir) = std::fs::read_dir(dir) else {
        return vec![];
    };
    let mut files: Vec<PathBuf> = read_dir
        .filter_map(|e| e.ok())
        .map(|e| e.path())
        .filter(|p| p.is_file() && is_supported_input_path(p))
        .collect();
    files.sort();
    files
}

async fn translate_single(args: TranslateArgs, config: &AppConfig) -> Result<(), TranslatorError> {
    let ass_classification_policy = args.classification.to_policy(&config.classification)?;
    let ocr_config = args.ocr.to_config(&config.ocr)?;
    let format_hint = parse_format_flag(args.format.as_deref());
    let input_kind = detect_input_kind(&args.input, format_hint);
    let provider_config = build_provider_config(&args, &config.provider)?;
    let provider = OpenAiCompatibleProvider::new(provider_config)?;
    let translator = Translator::new(Arc::new(provider));
    let options = build_translation_options(&args, &config.translation, ass_classification_policy)?;

    match input_kind {
        MediaInputKind::SubtitleFile { format_hint } => {
            if args.subtitle_input.is_some() {
                return Err(TranslatorError::InvalidConfig(
                    "--subtitle-input only applies when the input is a video container"
                        .to_owned(),
                ));
            }
            if args.subtitle_output.is_some() {
                return Err(TranslatorError::InvalidConfig(
                    "--subtitle-output only applies when the input is a video container"
                        .to_owned(),
                ));
            }
            if args.subtitle_stream_index.is_some() {
                return Err(TranslatorError::InvalidConfig(
                    "--subtitle-stream-index only applies when the input is a video container"
                        .to_owned(),
                ));
            }
            translate_subtitle_input(&args, &translator, &options, format_hint).await
        }
        MediaInputKind::Video => {
            if args.subtitle_input.is_some() && args.subtitle_stream_index.is_some() {
                return Err(TranslatorError::InvalidConfig(
                    "--subtitle-input cannot be combined with --subtitle-stream-index"
                        .to_owned(),
                ));
            }
            let tools = ExternalToolConfig::from_tool_config(&config.tools);
            translate_video_input(&args, &translator, &options, &tools, &ocr_config).await
        }
    }
}

async fn translate_subtitle_input(
    args: &TranslateArgs,
    translator: &Translator,
    options: &TranslationOptions,
    format_hint: Option<SubtitleFormat>,
) -> Result<(), TranslatorError> {
    let source = tokio::fs::read_to_string(&args.input).await?;
    let result = translator
        .translate_source(&source, format_hint, Some(&args.input), &options)
        .await?;

    maybe_write_classification_report(
        args.classification.classification_report.as_deref(),
        result.classification_report(),
    )
    .await?;

    if args.dry_run {
        eprintln!("note: --dry-run is set; translated subtitle text shown below (no output file will be written)");
        println!("{}", result.rendered());
        return Ok(());
    }

    let output_path = args.output.clone().unwrap_or_else(|| {
        default_output_path(
            &args.input,
            result.document().format(),
            &options.target_language,
        )
    });

    if output_path.exists() && !args.overwrite {
        return Err(TranslatorError::InvalidConfig(format!(
            "output file already exists: {} (use --overwrite to replace it)",
            output_path.display()
        )));
    }

    tokio::fs::write(&output_path, result.rendered()).await?;
    println!(
        "translated {} cues in {} batch(es) to {} ({} translatable, {} preserved, {} review)",
        result.document().cue_count(),
        result.batches(),
        output_path.display(),
        result.document().translatable_cue_count(),
        result.document().preserved_cue_count(),
        result.document().review_cue_count(),
    );

    Ok(())
}

async fn translate_video_input(
    args: &TranslateArgs,
    translator: &Translator,
    options: &TranslationOptions,
    tools: &ExternalToolConfig,
    ocr_config: &PgsOcrConfig,
) -> Result<(), TranslatorError> {
    let (result, stream_description, selected_stream, original_subtitle_stream_count) =
        if let Some(subtitle_input) = args.subtitle_input.as_ref() {
            let result = translator.translate_file_path(subtitle_input, options).await?;
            let stream_count = count_subtitle_streams(&args.input, tools).await?;
            (
                result,
                format!("external subtitle {}", subtitle_input.display()),
                None,
                stream_count,
            )
        } else {
            let job = prepare_video_subtitle_job(
                &args.input,
                tools,
                args.subtitle_stream_index,
                ocr_config,
                options.source_language.as_deref(),
            )
            .await?;
            let result = translator.translate_file_path(job.subtitle_path(), options).await?;
            (
                result,
                describe_selected_stream(job.selected_stream()),
                Some(job.selected_stream().clone()),
                job.original_subtitle_stream_count(),
            )
        };

    let rendered_subtitle = if selected_stream
        .as_ref()
        .is_some_and(|stream| stream.requires_ocr())
    {
        let mut cleaned_document = result.document().clone();
        cleanup_fragmentary_subtitle_document(&mut cleaned_document);
        formats::render_subtitle(&cleaned_document)?
    } else {
        result.rendered().to_owned()
    };

    maybe_write_classification_report(
        args.classification.classification_report.as_deref(),
        result.classification_report(),
    )
    .await?;

    if args.dry_run {
        eprintln!("note: --dry-run is set; translated subtitle text shown below (no output file will be written)");
        println!("{}", rendered_subtitle);
        return Ok(());
    }

    let temp_dir = tempdir()?;
    let fallback_subtitle_path = temp_dir.path().join(format!(
        "translated.{}",
        result.document().format().extension()
    ));
    let translated_subtitle_path = args
        .subtitle_output
        .clone()
        .unwrap_or_else(|| fallback_subtitle_path.clone());

    if translated_subtitle_path.exists() && !args.overwrite {
        return Err(TranslatorError::InvalidConfig(format!(
            "subtitle output file already exists: {} (use --overwrite to replace it)",
            translated_subtitle_path.display()
        )));
    }

    tokio::fs::write(&translated_subtitle_path, rendered_subtitle).await?;

    // Determine where the muxed video ends up.
    //
    // • --output <path>  → honour it explicitly; check overwrite guard against that path.
    // • default          → mux into a sibling temp file, then atomically rename it over the
    //                      original so no extra disk space is wasted for a permanent copy.
    let final_output_path: PathBuf;
    if let Some(explicit_output) = args.output.clone() {
        if explicit_output.exists() && !args.overwrite {
            return Err(TranslatorError::InvalidConfig(format!(
                "output file already exists: {} (use --overwrite to replace it)",
                explicit_output.display()
            )));
        }
        mux_translated_subtitle(
            &args.input,
            &translated_subtitle_path,
            &explicit_output,
            result.document().format(),
            selected_stream.as_ref(),
            original_subtitle_stream_count,
            &options.target_language,
            tools,
            args.overwrite,
        )
        .await?;
        final_output_path = explicit_output;
    } else {
        // Build a temp path in the same directory as the input so that the subsequent
        // rename(2) is guaranteed to be atomic (same filesystem, same mount point).
        let input_parent = args.input.parent().unwrap_or_else(|| std::path::Path::new("."));
        let input_stem = args
            .input
            .file_stem()
            .and_then(|stem| stem.to_str())
            .unwrap_or("video");
        let input_ext = args
            .input
            .extension()
            .and_then(|ext| ext.to_str())
            .unwrap_or("mkv");
        let temp_video_path = input_parent.join(format!(
            ".shinkai-translator-tmp-{}-{}.{}",
            input_stem,
            std::process::id(),
            input_ext,
        ));
        mux_translated_subtitle(
            &args.input,
            &translated_subtitle_path,
            &temp_video_path,
            result.document().format(),
            selected_stream.as_ref(),
            original_subtitle_stream_count,
            &options.target_language,
            tools,
            true, // always overwrite the temp file
        )
        .await?;
        tokio::fs::rename(&temp_video_path, &args.input).await?;
        final_output_path = args.input.clone();
    }

    if args.subtitle_output.is_some() {
        println!(
            "translated video subtitle stream {} into {} and muxed it into {} ({} translatable, {} preserved, {} review)",
            stream_description,
            translated_subtitle_path.display(),
            final_output_path.display(),
            result.document().translatable_cue_count(),
            result.document().preserved_cue_count(),
            result.document().review_cue_count(),
        );
    } else {
        println!(
            "translated video subtitle stream {} and muxed it into {} ({} translatable, {} preserved, {} review)",
            stream_description,
            final_output_path.display(),
            result.document().translatable_cue_count(),
            result.document().preserved_cue_count(),
            result.document().review_cue_count(),
        );
    }

    Ok(())
}

async fn validate(args: ValidateArgs, config: &AppConfig) -> Result<(), TranslatorError> {
    let ass_classification_policy = args.classification.to_policy(&config.classification)?;
    let ocr_config = args.ocr.to_config(&config.ocr)?;
    let format_hint = parse_format_flag(args.format.as_deref());
    let input_kind = detect_input_kind(&args.input, format_hint);

    match input_kind {
        MediaInputKind::SubtitleFile { format_hint } => {
            if args.subtitle_input.is_some() {
                return Err(TranslatorError::InvalidConfig(
                    "--subtitle-input only applies when the input is a video container"
                        .to_owned(),
                ));
            }
            if args.subtitle_stream_index.is_some() {
                return Err(TranslatorError::InvalidConfig(
                    "--subtitle-stream-index only applies when the input is a video container"
                        .to_owned(),
                ));
            }
            validate_subtitle_input(&args, &ass_classification_policy, format_hint).await
        }
        MediaInputKind::Video => {
            if args.subtitle_input.is_some() && args.subtitle_stream_index.is_some() {
                return Err(TranslatorError::InvalidConfig(
                    "--subtitle-input cannot be combined with --subtitle-stream-index"
                        .to_owned(),
                ));
            }
            let tools = ExternalToolConfig::from_tool_config(&config.tools);
            validate_video_input(&args, &ass_classification_policy, &tools, &ocr_config).await
        }
    }
}

async fn validate_subtitle_input(
    args: &ValidateArgs,
    ass_classification_policy: &AssClassificationPolicy,
    format_hint: Option<SubtitleFormat>,
) -> Result<(), TranslatorError> {
    let source = tokio::fs::read_to_string(&args.input).await?;
    let document = formats::parse_subtitle(&source, format_hint, Some(&args.input))?;
    let (classified_document, report) =
        formats::classify_document(&document, ass_classification_policy);

    maybe_write_classification_report(
        args.classification.classification_report.as_deref(),
        report.as_ref(),
    )
    .await?;

    println!(
        "valid {:?} subtitle with {} cue(s), {} translatable, {} preserved, {} review",
        classified_document.format(),
        classified_document.cue_count(),
        classified_document.translatable_cue_count(),
        classified_document.preserved_cue_count(),
        classified_document.review_cue_count()
    );
    Ok(())
}

async fn validate_video_input(
    args: &ValidateArgs,
    ass_classification_policy: &AssClassificationPolicy,
    tools: &ExternalToolConfig,
    ocr_config: &PgsOcrConfig,
) -> Result<(), TranslatorError> {
    let (document, stream_description) = if let Some(subtitle_input) = args.subtitle_input.as_ref() {
        let source = tokio::fs::read_to_string(subtitle_input).await?;
        let format_hint = SubtitleFormat::detect_from_path(subtitle_input);
        let document = formats::parse_subtitle(&source, format_hint, Some(subtitle_input))?;
        (
            document,
            format!("external subtitle {}", subtitle_input.display()),
        )
    } else {
        let job = prepare_video_subtitle_job(
            &args.input,
            tools,
            args.subtitle_stream_index,
            ocr_config,
            None,
        )
        .await?;
        let source = tokio::fs::read_to_string(job.subtitle_path()).await?;
        let document =
            formats::parse_subtitle(&source, Some(job.format()), Some(job.subtitle_path()))?;
        (document, describe_selected_stream(job.selected_stream()))
    };
    let (classified_document, report) =
        formats::classify_document(&document, ass_classification_policy);

    maybe_write_classification_report(
        args.classification.classification_report.as_deref(),
        report.as_ref(),
    )
    .await?;

    println!(
        "valid {:?} subtitle extracted from video stream {} with {} cue(s), {} translatable, {} preserved, {} review",
        classified_document.format(),
        stream_description,
        classified_document.cue_count(),
        classified_document.translatable_cue_count(),
        classified_document.preserved_cue_count(),
        classified_document.review_cue_count()
    );
    Ok(())
}

fn parse_format_flag(value: Option<&str>) -> Option<SubtitleFormat> {
    value.and_then(SubtitleFormat::parse_name)
}

fn default_output_path(input: &Path, format: SubtitleFormat, target_language: &str) -> PathBuf {
    let extension = format.extension();
    let stem = input
        .file_stem()
        .and_then(|stem| stem.to_str())
        .unwrap_or("translated");
    let target = sanitize_language_tag(target_language);

    input.with_file_name(format!("{stem}.{target}.{extension}"))
}

fn sanitize_language_tag(target_language: &str) -> String {
    let sanitized = target_language
        .chars()
        .map(|character| {
            if character.is_ascii_alphanumeric() {
                character.to_ascii_lowercase()
            } else {
                '-'
            }
        })
        .collect::<String>();

    sanitized.trim_matches('-').to_owned()
}

fn describe_selected_stream(stream: &SelectedSubtitleStream) -> String {
    let mut description = format!("#{} ({})", stream.index(), stream.codec_name());

    if let Some(language) = stream.language() {
        description.push_str(&format!(", lang={language}"));
    }

    if let Some(title) = stream.title() {
        description.push_str(&format!(", title={title}"));
    }

    description
}

impl AssClassificationArgs {
    fn to_policy(
        &self,
        defaults: &AssClassificationPolicy,
    ) -> Result<AssClassificationPolicy, TranslatorError> {
        let mut policy = defaults.clone();
        if let Some(value) = self.karaoke_policy.as_deref() {
            policy.karaoke_policy = parse_disposition(value)?;
        }
        if let Some(value) = self.explicit_song_policy.as_deref() {
            policy.explicit_song_policy = parse_disposition(value)?;
        }
        if let Some(value) = self.inferred_song_policy.as_deref() {
            policy.inferred_song_policy = parse_disposition(value)?;
        }

        if !self.style_markers.is_empty() {
            policy.style_markers = self.style_markers.clone();
        }

        if !self.effect_markers.is_empty() {
            policy.effect_markers = self.effect_markers.clone();
        }

        if !self.name_markers.is_empty() {
            policy.name_markers = self.name_markers.clone();
        }

        if self.disable_inferred_song_detection {
            policy.enable_inferred_song_detection = false;
        }

        if let Some(min_run_length) = self.min_inferred_song_run_length {
            policy.min_inferred_song_run_length = min_run_length;
        }

        policy.validate()?;
        Ok(policy)
    }
}

impl PgsOcrArgs {
    fn to_config(&self, defaults: &PgsOcrConfig) -> Result<PgsOcrConfig, TranslatorError> {
        let mut config = defaults.clone();
        if let Some(value) = self.ocr_language.as_deref() {
            config.language = PgsOcrLanguage::parse_name(value).ok_or_else(|| {
                TranslatorError::InvalidConfig(format!(
                    "invalid OCR language {:?}; expected one of auto, english, latin",
                    value
                ))
            })?;
        }
        if self.ocr_model_dir.is_some() {
            config.model_cache_dir = self.ocr_model_dir.clone();
        }
        config.validate()?;
        Ok(config)
    }
}

fn build_provider_config(
    args: &TranslateArgs,
    defaults: &ProviderConfig,
) -> Result<ProviderConfig, TranslatorError> {
    let mut provider_config = defaults.clone();

    if let Some(value) = args.base_url.as_ref() {
        provider_config.base_url = value.clone();
    }
    if let Some(value) = args.model.as_ref() {
        provider_config.model = value.clone();
    }
    if args.api_key.is_some() {
        provider_config.api_key = args.api_key.clone();
    }
    if let Some(value) = args.thinking_mode.as_deref() {
        provider_config.thinking_mode = parse_thinking_mode(value)?;
    }
    if let Some(value) = args.timeout_seconds {
        provider_config.timeout_seconds = value;
    }
    if let Some(value) = args.max_retries {
        provider_config.max_retries = value;
    }

    provider_config.validate()?;
    Ok(provider_config)
}

fn build_translation_options(
    args: &TranslateArgs,
    defaults: &shinkai_translator::TranslationDefaults,
    ass_classification_policy: AssClassificationPolicy,
) -> Result<TranslationOptions, TranslatorError> {
    let target_language = args
        .target_language
        .clone()
        .or_else(|| defaults.target_language.clone())
        .ok_or_else(|| {
            TranslatorError::InvalidConfig(
                "target language must be provided via --target-language or the config file"
                    .to_owned(),
            )
        })?;

    let options = TranslationOptions {
        source_language: args
            .source_language
            .clone()
            .or_else(|| defaults.source_language.clone()),
        target_language,
        max_batch_items: args.max_batch_items.unwrap_or(defaults.max_batch_items),
        max_batch_characters: args
            .max_batch_characters
            .unwrap_or(defaults.max_batch_characters),
        max_parallel_batches: args
            .max_parallel_batches
            .unwrap_or(defaults.max_parallel_batches),
        ass_classification_policy,
        system_prompt: args
            .system_prompt
            .clone()
            .or_else(|| defaults.system_prompt.clone()),
    };
    options.validate()?;
    Ok(options)
}

fn parse_disposition(value: &str) -> Result<CueDisposition, TranslatorError> {
    CueDisposition::parse_name(value).ok_or_else(|| {
        TranslatorError::InvalidConfig(format!(
            "invalid cue disposition {value:?}; expected one of translate, preserve, review"
        ))
    })
}

fn parse_thinking_mode(value: &str) -> Result<ThinkingMode, TranslatorError> {
    ThinkingMode::parse_name(value).ok_or_else(|| {
        TranslatorError::InvalidConfig(format!(
            "invalid thinking mode {value:?}; expected one of off, on, auto"
        ))
    })
}

async fn maybe_write_classification_report(
    path: Option<&Path>,
    report: Option<&shinkai_translator::SubtitleClassificationReport>,
) -> Result<(), TranslatorError> {
    let Some(path) = path else {
        return Ok(());
    };
    let report = report.ok_or_else(|| {
        TranslatorError::InvalidConfig(
            "classification reports are currently available only for ASS subtitles".to_owned(),
        )
    })?;
    let body = serde_json::to_string_pretty(report)?;
    tokio::fs::write(path, body).await?;
    Ok(())
}

async fn debug_ocr(args: DebugOcrArgs, config: &AppConfig) -> Result<(), TranslatorError> {
    let ocr_config = args.ocr.to_config(&config.ocr)?;
    let tools = ExternalToolConfig::from_tool_config(&config.tools);

    let is_pgs_file = args
        .input
        .extension()
        .and_then(|e| e.to_str())
        .is_some_and(|e| e.eq_ignore_ascii_case("sup") || e.eq_ignore_ascii_case("pgs"));

    // Keep extraction result alive to preserve the temp directory
    let extraction = if is_pgs_file {
        None
    } else {
        Some(extract_pgs_stream(&args.input, &tools, args.subtitle_stream_index).await?)
    };

    let stream_language: Option<String> = extraction
        .as_ref()
        .and_then(|e| e.selected_stream().language())
        .map(ToOwned::to_owned);

    let sup_path = if is_pgs_file {
        args.input.clone()
    } else {
        extraction.as_ref().unwrap().sup_path().to_path_buf()
    };

    let (from_ms, to_ms) = if let Some(at_str) = args.at.as_deref() {
        let center_ms = parse_timestamp_ms(at_str)?;
        let half_window_ms = args.window * 1_000;
        (
            Some(center_ms.saturating_sub(half_window_ms)),
            Some(center_ms.saturating_add(half_window_ms)),
        )
    } else if let (Some(from_str), Some(to_str)) = (args.from.as_deref(), args.to.as_deref()) {
        (
            Some(parse_timestamp_ms(from_str)?),
            Some(parse_timestamp_ms(to_str)?),
        )
    } else {
        (None, None)
    };

    if from_ms.is_none() && to_ms.is_none() {
        eprintln!("warning: no time range specified, processing all display sets (may be slow)");
    }

    let result = debug_pgs_ocr(
        &sup_path,
        &ocr_config,
        stream_language.as_deref(),
        from_ms,
        to_ms,
        &args.output,
    )
    .await?;

    println!(
        "debug OCR: {} display sets processed, {} with text → {}",
        result.frames_processed,
        result.frames_with_text,
        args.output.display()
    );

    Ok(())
}

fn parse_timestamp_ms(input: &str) -> Result<u64, TranslatorError> {
    let s = input.trim();
    let normalized = s.replace(',', ".");

    let (hms_part, frac_ms): (&str, u64) = if let Some(dot_pos) = normalized.rfind('.') {
        let hms = &normalized[..dot_pos];
        let frac_str = &normalized[dot_pos + 1..];
        let frac: u64 = frac_str
            .parse()
            .map_err(|_| TranslatorError::InvalidConfig(format!("invalid timestamp: {s}")))?;
        let ms = match frac_str.len() {
            1 => frac * 100,
            2 => frac * 10,
            _ => frac,
        };
        (hms, ms)
    } else {
        (normalized.as_str(), 0)
    };

    let parts: Vec<&str> = hms_part.split(':').collect();
    let (hours, minutes, seconds) = match parts.as_slice() {
        [h_str, m_str, s_str] => {
            let h: u64 = h_str
                .parse()
                .map_err(|_| TranslatorError::InvalidConfig(format!("invalid timestamp: {s}")))?;
            let m: u64 = m_str
                .parse()
                .map_err(|_| TranslatorError::InvalidConfig(format!("invalid timestamp: {s}")))?;
            let sec: u64 = s_str
                .parse()
                .map_err(|_| TranslatorError::InvalidConfig(format!("invalid timestamp: {s}")))?;
            (h, m, sec)
        }
        _ => {
            return Err(TranslatorError::InvalidConfig(format!(
                "invalid timestamp format: {s} (expected HH:MM:SS or HH:MM:SS.mmm)"
            )));
        }
    };

    if minutes >= 60 || seconds >= 60 {
        return Err(TranslatorError::InvalidConfig(format!(
            "invalid timestamp: {s} (minutes and seconds must be 0–59)"
        )));
    }

    Ok(hours * 3_600_000 + minutes * 60_000 + seconds * 1_000 + frac_ms)
}