use std::collections::{BTreeMap, HashMap};
use std::path::Path;
use serde::{Deserialize, Serialize};
use crate::error::TranslatorError;
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum SubtitleFormat {
Srt,
Vtt,
Ass,
}
impl SubtitleFormat {
pub fn detect_from_path(path: &Path) -> Option<Self> {
let extension = path.extension()?.to_str()?;
Self::parse_name(extension)
}
pub fn detect_from_content(source: &str) -> Option<Self> {
let trimmed = source.trim_start_matches('\u{feff}').trim_start();
if trimmed.starts_with("WEBVTT") {
return Some(Self::Vtt);
}
if trimmed.contains("[Script Info]") || trimmed.contains("[V4+ Styles]") {
return Some(Self::Ass);
}
if trimmed.contains("-->") {
return Some(Self::Srt);
}
None
}
pub fn parse_name(value: &str) -> Option<Self> {
match value.trim().to_ascii_lowercase().as_str() {
"srt" => Some(Self::Srt),
"vtt" => Some(Self::Vtt),
"ass" | "ssa" => Some(Self::Ass),
_ => None,
}
}
pub fn extension(self) -> &'static str {
match self {
Self::Srt => "srt",
Self::Vtt => "vtt",
Self::Ass => "ass",
}
}
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum CueKind {
Dialogue,
Karaoke,
Song,
}
impl CueKind {
pub fn default_disposition(self) -> CueDisposition {
match self {
Self::Dialogue => CueDisposition::Translate,
Self::Karaoke | Self::Song => CueDisposition::Preserve,
}
}
pub fn is_translatable(self) -> bool {
self.default_disposition().is_translatable()
}
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum CueDisposition {
Translate,
Preserve,
Review,
}
impl CueDisposition {
pub fn is_translatable(self) -> bool {
matches!(self, Self::Translate)
}
pub fn parse_name(value: &str) -> Option<Self> {
match value.trim().to_ascii_lowercase().as_str() {
"translate" => Some(Self::Translate),
"preserve" => Some(Self::Preserve),
"review" => Some(Self::Review),
_ => None,
}
}
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ThinkingMode {
Off,
On,
Auto,
}
impl ThinkingMode {
pub fn parse_name(value: &str) -> Option<Self> {
match value.trim().to_ascii_lowercase().as_str() {
"off" => Some(Self::Off),
"on" => Some(Self::On),
"auto" => Some(Self::Auto),
_ => None,
}
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default)]
pub struct AssClassificationPolicy {
pub karaoke_policy: CueDisposition,
pub explicit_song_policy: CueDisposition,
pub inferred_song_policy: CueDisposition,
pub style_markers: Vec<String>,
pub effect_markers: Vec<String>,
pub name_markers: Vec<String>,
pub enable_inferred_song_detection: bool,
pub min_inferred_song_run_length: usize,
}
impl Default for AssClassificationPolicy {
fn default() -> Self {
let default_markers = vec![
"karaoke".to_owned(),
"kfx".to_owned(),
"song".to_owned(),
"songs".to_owned(),
"lyric".to_owned(),
"lyrics".to_owned(),
"opening".to_owned(),
"ending".to_owned(),
"insert".to_owned(),
"music".to_owned(),
"romaji".to_owned(),
"op".to_owned(),
"ed".to_owned(),
];
Self {
karaoke_policy: CueDisposition::Preserve,
explicit_song_policy: CueDisposition::Preserve,
inferred_song_policy: CueDisposition::Review,
style_markers: default_markers.clone(),
effect_markers: default_markers.clone(),
name_markers: default_markers,
enable_inferred_song_detection: true,
min_inferred_song_run_length: 4,
}
}
}
impl AssClassificationPolicy {
pub fn validate(&self) -> Result<(), TranslatorError> {
if self.min_inferred_song_run_length == 0 {
return Err(TranslatorError::InvalidConfig(
"min_inferred_song_run_length must be greater than zero".to_owned(),
));
}
Ok(())
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct CueClassification {
kind: CueKind,
disposition: CueDisposition,
confidence: u8,
reason: Option<String>,
}
impl CueClassification {
pub fn new(
kind: CueKind,
disposition: CueDisposition,
confidence: u8,
reason: Option<String>,
) -> Self {
Self {
kind,
disposition,
confidence,
reason,
}
}
pub fn kind(&self) -> CueKind {
self.kind
}
pub fn disposition(&self) -> CueDisposition {
self.disposition
}
pub fn confidence(&self) -> u8 {
self.confidence
}
pub fn reason(&self) -> Option<&str> {
self.reason.as_deref()
}
}
impl Default for CueClassification {
fn default() -> Self {
Self {
kind: CueKind::Dialogue,
disposition: CueDisposition::Translate,
confidence: 100,
reason: None,
}
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SubtitleClassificationEntry {
pub cue_id: String,
pub kind: CueKind,
pub disposition: CueDisposition,
pub confidence: u8,
pub reason: Option<String>,
pub start: String,
pub end: String,
pub text_preview: String,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SubtitleClassificationSummary {
pub total_cues: usize,
pub translatable_cues: usize,
pub preserved_cues: usize,
pub review_cues: usize,
pub dialogue_cues: usize,
pub karaoke_cues: usize,
pub song_cues: usize,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SubtitleClassificationReport {
pub format: SubtitleFormat,
pub summary: SubtitleClassificationSummary,
pub entries: Vec<SubtitleClassificationEntry>,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SubtitleCue {
id: String,
identifier: Option<String>,
start: String,
end: String,
settings: Option<String>,
text: String,
attributes: BTreeMap<String, String>,
classification: CueClassification,
}
impl SubtitleCue {
pub fn new(
id: impl Into<String>,
identifier: Option<String>,
start: impl Into<String>,
end: impl Into<String>,
settings: Option<String>,
text: impl Into<String>,
attributes: BTreeMap<String, String>,
) -> Self {
Self::new_with_classification(
id,
identifier,
start,
end,
settings,
text,
attributes,
CueClassification::default(),
)
}
pub fn new_with_kind(
id: impl Into<String>,
identifier: Option<String>,
start: impl Into<String>,
end: impl Into<String>,
settings: Option<String>,
text: impl Into<String>,
attributes: BTreeMap<String, String>,
kind: CueKind,
) -> Self {
Self::new_with_classification(
id,
identifier,
start,
end,
settings,
text,
attributes,
CueClassification::new(kind, kind.default_disposition(), 100, None),
)
}
pub fn new_with_classification(
id: impl Into<String>,
identifier: Option<String>,
start: impl Into<String>,
end: impl Into<String>,
settings: Option<String>,
text: impl Into<String>,
attributes: BTreeMap<String, String>,
classification: CueClassification,
) -> Self {
Self {
id: id.into(),
identifier,
start: start.into(),
end: end.into(),
settings,
text: text.into(),
attributes,
classification,
}
}
pub fn id(&self) -> &str {
&self.id
}
pub fn identifier(&self) -> Option<&str> {
self.identifier.as_deref()
}
pub fn start(&self) -> &str {
&self.start
}
pub fn end(&self) -> &str {
&self.end
}
pub fn settings(&self) -> Option<&str> {
self.settings.as_deref()
}
pub fn text(&self) -> &str {
&self.text
}
pub fn attributes(&self) -> &BTreeMap<String, String> {
&self.attributes
}
pub fn kind(&self) -> CueKind {
self.classification.kind()
}
pub fn disposition(&self) -> CueDisposition {
self.classification.disposition()
}
pub fn classification(&self) -> &CueClassification {
&self.classification
}
pub fn classification_reason(&self) -> Option<&str> {
self.classification.reason()
}
pub fn classification_confidence(&self) -> u8 {
self.classification.confidence()
}
pub fn is_translatable(&self) -> bool {
self.disposition().is_translatable()
}
pub fn text_len(&self) -> usize {
self.text.chars().count()
}
pub(crate) fn set_text(&mut self, text: String) {
self.text = text;
}
pub(crate) fn set_classification(&mut self, classification: CueClassification) {
self.classification = classification;
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SubtitleDocument {
format: SubtitleFormat,
cues: Vec<SubtitleCue>,
pub(crate) render_plan: RenderPlan,
}
impl SubtitleDocument {
pub(crate) fn from_parts(
format: SubtitleFormat,
cues: Vec<SubtitleCue>,
render_plan: RenderPlan,
) -> Self {
Self {
format,
cues,
render_plan,
}
}
pub fn format(&self) -> SubtitleFormat {
self.format
}
pub fn cues(&self) -> &[SubtitleCue] {
&self.cues
}
pub fn cue_count(&self) -> usize {
self.cues.len()
}
pub fn translatable_cue_count(&self) -> usize {
self.cues.iter().filter(|cue| cue.is_translatable()).count()
}
pub fn preserved_cue_count(&self) -> usize {
self.cues
.iter()
.filter(|cue| cue.disposition() == CueDisposition::Preserve)
.count()
}
pub fn review_cue_count(&self) -> usize {
self.cues
.iter()
.filter(|cue| cue.disposition() == CueDisposition::Review)
.count()
}
pub(crate) fn cues_mut(&mut self) -> &mut [SubtitleCue] {
&mut self.cues
}
pub(crate) fn translated_with(
&self,
replacements: &HashMap<String, String>,
) -> Result<Self, TranslatorError> {
let mut translated = self.clone();
for cue in &mut translated.cues {
if !cue.is_translatable() {
continue;
}
let replacement = replacements.get(cue.id()).ok_or_else(|| {
TranslatorError::Validation(format!("missing translated text for cue {}", cue.id()))
})?;
cue.set_text(replacement.clone());
}
Ok(translated)
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum RenderPlan {
Srt,
Vtt { blocks: Vec<RenderBlock> },
Ass { lines: Vec<AssRenderBlock> },
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum RenderBlock {
Raw(String),
Cue(usize),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum AssRenderBlock {
Raw(String),
Cue(AssCueRenderTemplate),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct AssCueRenderTemplate {
pub cue_index: usize,
pub prefix: String,
pub suffix: String,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct TranslationOptions {
pub source_language: Option<String>,
pub target_language: String,
pub max_batch_items: usize,
pub max_batch_characters: usize,
pub max_parallel_batches: usize,
pub ass_classification_policy: AssClassificationPolicy,
pub system_prompt: Option<String>,
}
impl Default for TranslationOptions {
fn default() -> Self {
Self {
source_language: None,
target_language: "English".to_owned(),
max_batch_items: 8,
max_batch_characters: 4_000,
max_parallel_batches: 5,
ass_classification_policy: AssClassificationPolicy::default(),
system_prompt: None,
}
}
}
impl TranslationOptions {
pub fn validate(&self) -> Result<(), TranslatorError> {
if self.target_language.trim().is_empty() {
return Err(TranslatorError::InvalidConfig(
"target_language must not be empty".to_owned(),
));
}
if self.max_batch_items == 0 {
return Err(TranslatorError::InvalidConfig(
"max_batch_items must be greater than zero".to_owned(),
));
}
if self.max_batch_characters == 0 {
return Err(TranslatorError::InvalidConfig(
"max_batch_characters must be greater than zero".to_owned(),
));
}
if self.max_parallel_batches == 0 {
return Err(TranslatorError::InvalidConfig(
"max_parallel_batches must be greater than zero".to_owned(),
));
}
self.ass_classification_policy.validate()?;
Ok(())
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default)]
pub struct ProviderConfig {
pub base_url: String,
pub model: String,
pub api_key: Option<String>,
pub timeout_seconds: u64,
pub max_retries: u32,
pub temperature: f32,
pub thinking_mode: ThinkingMode,
pub custom_headers: BTreeMap<String, String>,
}
impl Default for ProviderConfig {
fn default() -> Self {
Self {
base_url: "https://api.openai.com/v1".to_owned(),
model: "gpt-4o-mini".to_owned(),
api_key: None,
timeout_seconds: 60,
max_retries: 2,
temperature: 0.2,
thinking_mode: ThinkingMode::Off,
custom_headers: BTreeMap::new(),
}
}
}
impl ProviderConfig {
pub fn validate(&self) -> Result<(), TranslatorError> {
if self.base_url.trim().is_empty() {
return Err(TranslatorError::InvalidConfig(
"base_url must not be empty".to_owned(),
));
}
if self.model.trim().is_empty() {
return Err(TranslatorError::InvalidConfig(
"model must not be empty".to_owned(),
));
}
if self.timeout_seconds == 0 {
return Err(TranslatorError::InvalidConfig(
"timeout_seconds must be greater than zero".to_owned(),
));
}
Ok(())
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TranslationResult {
document: SubtitleDocument,
rendered: String,
batches: usize,
classification_report: Option<SubtitleClassificationReport>,
}
impl TranslationResult {
pub(crate) fn new(
document: SubtitleDocument,
rendered: String,
batches: usize,
classification_report: Option<SubtitleClassificationReport>,
) -> Self {
Self {
document,
rendered,
batches,
classification_report,
}
}
pub fn document(&self) -> &SubtitleDocument {
&self.document
}
pub fn rendered(&self) -> &str {
&self.rendered
}
pub fn batches(&self) -> usize {
self.batches
}
pub fn classification_report(&self) -> Option<&SubtitleClassificationReport> {
self.classification_report.as_ref()
}
}