sceptre-cli 0.7.0

//! Command definitions and dispatch.

use std::io::Write;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;

use anyhow::{Context, Result};
use clap::{CommandFactory, Parser, Subcommand};
use sceptre::{OcrConfig, ReadOptions, Reader};

use crate::output::{self, OutputFormat};
use crate::overrides::OcrOverrides;
use crate::timing::{StageTimer, TimingsReport};

/// CRAFT + gen2 CRNN optical character recognition over ONNX.
#[derive(Parser)]
#[command(name = "sceptre", version, about)]
pub struct Cli {
    #[command(subcommand)]
    command: Commands,

    /// Log level: error, warn, info, debug, or trace.
    #[arg(long, global = true, default_value = "warn", env = "EASYOCR_LOG")]
    log_level: String,
}

/// Model-management actions.
#[derive(Subcommand)]
pub enum ModelsAction {
    /// List the known models and their cache status.
    List {
        #[command(flatten)]
        overrides: OcrOverrides,
        /// Cover every supported language instead of the configured ones.
        #[arg(long, conflicts_with = "lang")]
        all: bool,
        /// Output format.
        #[arg(long, value_enum, default_value_t = OutputFormat::Text)]
        format: OutputFormat,
    },
    /// Download the models for the configured languages.
    Download {
        #[command(flatten)]
        overrides: OcrOverrides,
        /// Download every supported language, not just the configured ones.
        #[arg(long, conflicts_with = "lang")]
        all: bool,
        /// Output format.
        #[arg(long, value_enum, default_value_t = OutputFormat::Text)]
        format: OutputFormat,
    },
}

#[derive(Subcommand)]
enum Commands {
    /// Run the full OCR pipeline over one or more images.
    Run {
        /// Paths to the input images (one or more); a single `Reader` is reused across them.
        #[arg(required = true, num_args = 1..)]
        images: Vec<PathBuf>,
        #[command(flatten)]
        overrides: OcrOverrides,
        /// Output format.
        #[arg(long, value_enum, default_value_t = OutputFormat::Text)]
        format: OutputFormat,
        /// Emit only the recognized text, omitting confidence and box detail.
        #[arg(long)]
        no_detail: bool,
        /// Report a per-stage timing breakdown (load/detect/recognize) on stderr, and in
        /// the `--format json` payload.
        #[arg(long)]
        timings: bool,
    },
    /// Detect text regions only, without recognition.
    Detect {
        /// Path to the input image.
        image: PathBuf,
        #[command(flatten)]
        overrides: OcrOverrides,
        /// Output format.
        #[arg(long, value_enum, default_value_t = OutputFormat::Text)]
        format: OutputFormat,
    },
    /// Recognize text in a pre-cropped line image.
    Recognize {
        /// Path to the cropped line image.
        image: PathBuf,
        #[command(flatten)]
        overrides: OcrOverrides,
        /// Output format.
        #[arg(long, value_enum, default_value_t = OutputFormat::Text)]
        format: OutputFormat,
    },
    /// List or download models.
    Models {
        #[command(subcommand)]
        action: ModelsAction,
    },
    /// Report the runtime this build would execute on, plus the model pins.
    Env {
        #[command(flatten)]
        overrides: OcrOverrides,
        /// Output format.
        #[arg(long, value_enum, default_value_t = OutputFormat::Text)]
        format: OutputFormat,
    },
    /// Print shell completions to stdout.
    Completions {
        /// Target shell.
        shell: clap_complete::Shell,
    },
    /// Run the MCP stdio server.
    #[cfg(feature = "mcp")]
    Mcp {
        #[command(flatten)]
        overrides: OcrOverrides,
    },
}

/// Build an `OcrConfig` from defaults with the CLI overrides applied.
fn config_from(overrides: &OcrOverrides) -> OcrConfig {
    let mut config = OcrConfig::default();
    overrides.apply(&mut config);
    config
}

/// Build a `Reader` from the CLI overrides.
fn build_reader(overrides: &OcrOverrides) -> Result<Reader> {
    Reader::builder()
        .config(config_from(overrides))
        .build()
        .context("building the OCR reader")
}

/// Build a `Reader` wired to a stage timer for the `--timings` breakdown.
fn build_reader_timed(overrides: &OcrOverrides, timer: Arc<StageTimer>) -> Result<Reader> {
    Reader::builder()
        .config(config_from(overrides))
        .progress(timer)
        .build()
        .context("building the OCR reader")
}

/// A color-aware stdout writer that strips escapes on non-TTY / `NO_COLOR`.
fn stdout() -> anstream::Stdout {
    anstream::stdout()
}

/// A color-aware stderr writer that strips escapes on non-TTY / `NO_COLOR`.
fn stderr() -> anstream::Stderr {
    anstream::stderr()
}

/// Run the full OCR pipeline over one or more images, reusing a single `Reader`.
///
/// A single image keeps the historical single-result output; two or more images
/// switch to batch rendering where a per-image failure is recorded and the run
/// continues, exiting non-zero if any image failed.
fn run_ocr(
    images: Vec<PathBuf>,
    overrides: OcrOverrides,
    format: OutputFormat,
    no_detail: bool,
    timings: bool,
) -> Result<()> {
    let timer = timings.then(|| Arc::new(StageTimer::new()));
    let reader = match &timer {
        Some(timer) => build_reader_timed(&overrides, timer.clone())?,
        None => build_reader(&overrides)?,
    };
    let options = ReadOptions { detail: !no_detail };
    let detail = !no_detail;
    let started = Instant::now();

    if let [image] = images.as_slice() {
        let result = reader
            .readtext(image, &options)
            .with_context(|| format!("running OCR over {image:?}"))?;
        let timings = report_timings(&timer, started);
        output::render_result(&result, format, detail, timings, &mut stdout()).context("writing OCR results")?;
        return Ok(());
    }

    let mut outcomes: Vec<(PathBuf, output::ImageOutcome)> = Vec::with_capacity(images.len());
    let mut failures = 0usize;
    for image in &images {
        match reader.readtext(image, &options) {
            Ok(result) => outcomes.push((image.clone(), Ok(result))),
            Err(error) => {
                failures += 1;
                let message = format!("{error:#}");
                // Diagnostic on stderr; the failure is also recorded in the structured stdout output. ~keep
                let _ = writeln!(stderr(), "error: {}: {message}", image.display());
                outcomes.push((image.clone(), Err(message)));
            }
        }
    }
    let timings = report_timings(&timer, started);
    output::render_batch(&outcomes, format, detail, timings, &mut stdout()).context("writing OCR results")?;
    if failures > 0 {
        anyhow::bail!("{failures} of {} image(s) failed", images.len());
    }
    Ok(())
}

/// Print the stage-timing breakdown to stderr when `--timings` installed a timer.
///
/// Returns the same breakdown so the caller can fold it into a `--format json` payload.
fn report_timings(timer: &Option<Arc<StageTimer>>, started: Instant) -> Option<TimingsReport> {
    let timer = timer.as_ref()?;
    let breakdown = timer.breakdown(started, Instant::now());
    let _ = writeln!(stderr(), "{}", crate::timing::render(&breakdown));
    Some(TimingsReport::from(&breakdown))
}

/// Detect text regions and render their quads.
fn run_detect(image: PathBuf, overrides: OcrOverrides, format: OutputFormat) -> Result<()> {
    let reader = build_reader(&overrides)?;
    let image_data = sceptre::Image::from_path(&image).with_context(|| format!("loading {image:?}"))?;
    let quads = reader
        .detect(&image_data, &ReadOptions::default())
        .with_context(|| format!("detecting text regions in {image:?}"))?;
    output::render_quads(&quads, format, &mut stdout()).context("writing detected regions")?;
    Ok(())
}

/// Recognize a single cropped line and render it.
fn run_recognize(image: PathBuf, overrides: OcrOverrides, format: OutputFormat) -> Result<()> {
    let reader = build_reader(&overrides)?;
    let image_data = sceptre::Image::from_path(&image).with_context(|| format!("loading {image:?}"))?;
    let line = reader
        .recognize_line(&image_data, &ReadOptions::default())
        .with_context(|| format!("recognizing text in {image:?}"))?;
    output::render_line(&line, format, true, &mut stdout()).context("writing recognized line")?;
    Ok(())
}

/// Build the `models` config, expanding `--all` to every supported language.
fn models_config(overrides: &OcrOverrides, all: bool) -> OcrConfig {
    let mut config = config_from(overrides);
    if all {
        config.model.languages = crate::overrides::every_language();
    }
    config
}

/// Dispatch a `models` subcommand.
fn run_models(action: ModelsAction) -> Result<()> {
    match action {
        ModelsAction::List { overrides, all, format } => {
            let config = models_config(&overrides, all);
            let models = sceptre::model_manifest(&config).context("building the model manifest")?;
            output::render_models(&models, format, &mut stdout()).context("writing the model list")?;
        }
        ModelsAction::Download { overrides, all, format } => {
            let config = models_config(&overrides, all);
            let models = sceptre::download_models(&config).context("downloading models")?;
            output::render_models(&models, format, &mut stdout()).context("writing the model list")?;
        }
    }
    Ok(())
}

/// Report the runtime and the model pins for the overridden configuration.
///
/// Probing goes through the overridden configuration so `--accelerator coreml`
/// answers for CoreML rather than for the default; with no overrides this is
/// exactly [`sceptre::runtime_info`], which reads the same default configuration.
///
/// The model pins describe the binary rather than one run, so without an explicit
/// `--lang` every supported language is listed: a report that records this block
/// once may have exercised any of them.
fn run_env(overrides: OcrOverrides, format: OutputFormat) -> Result<()> {
    let config = models_config(&overrides, !overrides.has_languages());
    let runtime = sceptre::runtime_info_for(&config.model).context("describing the runtime")?;
    let models = sceptre::model_descriptors(&config).context("resolving the model registry")?;
    output::render_environment(&runtime, &models, format, &mut stdout()).context("writing the environment report")?;
    Ok(())
}

impl Cli {
    /// Initialize tracing to stderr from the configured log level.
    pub fn init_tracing(&self) {
        use tracing_subscriber::EnvFilter;
        let filter = EnvFilter::try_new(&self.log_level).unwrap_or_else(|_| EnvFilter::new("warn"));
        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .with_writer(std::io::stderr)
            .init();
    }

    /// Dispatch the selected command.
    pub fn run(self) -> Result<()> {
        match self.command {
            Commands::Run {
                images,
                overrides,
                format,
                no_detail,
                timings,
            } => run_ocr(images, overrides, format, no_detail, timings),
            Commands::Detect {
                image,
                overrides,
                format,
            } => run_detect(image, overrides, format),
            Commands::Recognize {
                image,
                overrides,
                format,
            } => run_recognize(image, overrides, format),
            Commands::Models { action } => run_models(action),
            Commands::Env { overrides, format } => run_env(overrides, format),
            Commands::Completions { shell } => {
                let mut command = Cli::command();
                clap_complete::generate(shell, &mut command, "sceptre", &mut std::io::stdout());
                Ok(())
            }
            #[cfg(feature = "mcp")]
            Commands::Mcp { overrides } => {
                let reader = build_reader(&overrides)?;
                sceptre::mcp::serve(reader).context("running the MCP server")
            }
        }
    }
}