use crate::Result;
use crate::error::SubXError;
use crate::services::ai::{AnalysisRequest, ConfidenceScore, MatchResult, VerificationRequest};
use serde_json;
/// Prompt builder trait for AI providers.
pub trait PromptBuilder {
/// Build analysis prompt.
fn build_analysis_prompt(&self, request: &AnalysisRequest) -> String {
build_analysis_prompt_base(request)
}
/// Build verification prompt.
fn build_verification_prompt(&self, request: &VerificationRequest) -> String {
build_verification_prompt_base(request)
}
/// System message for analysis prompt.
fn get_analysis_system_message() -> &'static str {
"You are an expert subtitle-matching assistant."
}
/// System message for verification prompt.
fn get_verification_system_message() -> &'static str {
"You are an expert subtitle-matching verifier."
}
}
/// Response parsing trait for AI providers.
pub trait ResponseParser {
/// Parse match result.
fn parse_match_result(&self, response: &str) -> Result<MatchResult> {
parse_match_result_base(response)
}
/// Parse confidence score.
fn parse_confidence_score(&self, response: &str) -> Result<ConfidenceScore> {
parse_confidence_score_base(response)
}
}
/// Escape a string for inclusion in an XML attribute value.
fn xml_escape_attr(value: &str) -> String {
let mut out = String::with_capacity(value.len());
for ch in value.chars() {
match ch {
'&' => out.push_str("&"),
'<' => out.push_str("<"),
'>' => out.push_str(">"),
'"' => out.push_str("""),
'\'' => out.push_str("'"),
_ => out.push(ch),
}
}
out
}
/// Wrap `content` in CDATA, neutralizing any literal `]]>` substring by
/// splitting it across two CDATA sections so the surrounding XML
/// structure remains parseable.
fn cdata_wrap(content: &str) -> String {
let safe = content.replace("]]>", "]]]]><![CDATA[>");
format!("<![CDATA[{}]]>", safe)
}
/// Tolerant parser for the legacy `"ID:… | Name:… | Path:…"` strings that
/// callers pass into [`AnalysisRequest::video_files`] and
/// [`AnalysisRequest::subtitle_files`]. Returns `(id, name, path)` or
/// `None` if the input does not look like that shape; the caller falls
/// back to embedding the raw string verbatim inside a `<file>` element.
fn parse_id_name_path(entry: &str) -> Option<(String, String, String)> {
let mut id = None;
let mut name = None;
let mut path = None;
for part in entry.split(" | ") {
if let Some(rest) = part.strip_prefix("ID:") {
id = Some(rest.trim().to_string());
} else if let Some(rest) = part.strip_prefix("Name:") {
name = Some(rest.trim().to_string());
} else if let Some(rest) = part.strip_prefix("Path:") {
path = Some(rest.trim().to_string());
}
}
match (id, name, path) {
(Some(i), Some(n), Some(p)) => Some((i, n, p)),
_ => None,
}
}
/// Render a single file-inventory entry as `<file id="…" name="…" path="…"/>`.
fn render_file_element(entry: &str) -> String {
if let Some((id, name, path)) = parse_id_name_path(entry) {
format!(
" <file id=\"{}\" name=\"{}\" path=\"{}\"/>\n",
xml_escape_attr(&id),
xml_escape_attr(&name),
xml_escape_attr(&path)
)
} else {
// Fallback: embed the raw string verbatim so we never lose data.
format!(
" <file><![CDATA[{}]]></file>\n",
entry.replace("]]>", "]]]]><![CDATA[>")
)
}
}
/// Build analysis prompt for AI providers.
///
/// Emits an XML-tagged prompt body following the Claude prompt-engineering
/// best-practices guide: a `<role>` description, a `<instructions>` block
/// with Markdown-headed sub-sections, separate `<video_files>` and
/// `<subtitle_files>` inventories, a `<content_samples>` block keyed by
/// stable `subtitle_file_id`, an `<output_schema>` block describing the
/// expected JSON response, and a single `<example>` for the model to
/// imitate.
pub fn build_analysis_prompt_base(request: &AnalysisRequest) -> String {
let mut p = String::new();
p.push_str("<role>You are an expert subtitle-matching assistant. Pair each subtitle file with the video it belongs to using filename evidence and the subtitle preview text.</role>\n");
p.push_str("<instructions>\n");
p.push_str("## Task\n");
p.push_str("Analyze the inventories below and return JSON-only output that pairs every confidently matched subtitle with its video. Use the file IDs verbatim — never invent new IDs and never echo filenames as IDs.\n\n");
p.push_str("## Language Inference Rules\n");
p.push_str("Use BOTH the filename and the `<content_samples>` text to infer each subtitle's language. The content sample is the primary signal: when filename and content disagree, trust the content. Output a short language code (e.g. `tc`, `sc`, `en`, `ja`, `ko`, `fr`, `de`, `es`, `pt`, `ru`). Use `und` only when the language genuinely cannot be determined.\n\n");
p.push_str("## Naming and Uniqueness Rules\n");
p.push_str("Across the entire `matches[]` array — not just within a single video — no two entries MAY produce the same final filename in the same target directory. When two subtitles for the same video would otherwise collide, supply distinct `language` codes. If both share the same language, supply a numeric `target_filename_suffix` (`2`, `3`, …) starting at `2`. The suffix MUST be one short token of `[A-Za-z0-9_-]` and at most 16 characters; it is spliced between the video base name and the subtitle extension.\n\n");
p.push_str("## Output Schema\n");
p.push_str("Respond with a single JSON object — no prose, no Markdown fences. Keys are stable: `matches[]` (array of objects with `video_file_id`, `subtitle_file_id`, `confidence`, `match_factors`, optional `language`, optional `target_filename_suffix`), `confidence` (overall score), and `reasoning` (short English explanation).\n");
p.push_str("</instructions>\n");
p.push_str("<video_files>\n");
for entry in &request.video_files {
p.push_str(&render_file_element(entry));
}
p.push_str("</video_files>\n");
p.push_str("<subtitle_files>\n");
for entry in &request.subtitle_files {
p.push_str(&render_file_element(entry));
}
p.push_str("</subtitle_files>\n");
if !request.content_samples.is_empty() {
// Each `<sample>` is keyed by `subtitle_file_id` so the model
// links samples back to entries in `<subtitle_files>` by ID
// alone — no filename leak that could be echoed as an ID.
p.push_str("<content_samples_instructions>Each `<sample>` element's `subtitle_file_id` attribute matches the `id` attribute on the corresponding `<file>` element in `<subtitle_files>`. Resolve sample-to-subtitle relationships through this ID, never by filename.</content_samples_instructions>\n");
p.push_str("<content_samples>\n");
for sample in &request.content_samples {
p.push_str(&format!(
" <sample subtitle_file_id=\"{}\">{}</sample>\n",
xml_escape_attr(&sample.subtitle_file_id),
cdata_wrap(&sample.content_preview)
));
}
p.push_str("</content_samples>\n");
}
p.push_str("<output_schema>\n");
p.push_str("{\n");
p.push_str(" \"matches\": [\n");
p.push_str(" {\n");
p.push_str(" \"video_file_id\": \"file_<uuid>\",\n");
p.push_str(" \"subtitle_file_id\": \"file_<uuid>\",\n");
p.push_str(" \"confidence\": 0.95,\n");
p.push_str(" \"match_factors\": [\"filename_similarity\", \"content_correlation\"],\n");
p.push_str(" \"language\": \"tc\",\n");
p.push_str(" \"target_filename_suffix\": \"tc\"\n");
p.push_str(" }\n");
p.push_str(" ],\n");
p.push_str(" \"confidence\": 0.9,\n");
p.push_str(" \"reasoning\": \"Short English explanation\"\n");
p.push_str("}\n");
p.push_str("</output_schema>\n");
p.push_str("<example>\n");
p.push_str(" <input_summary>One video plus two same-named subtitles whose content samples reveal Traditional vs Simplified Chinese.</input_summary>\n");
p.push_str(" <output>{\"matches\":[{\"video_file_id\":\"file_v1\",\"subtitle_file_id\":\"file_s1\",\"confidence\":0.95,\"match_factors\":[\"content_correlation\"],\"language\":\"tc\",\"target_filename_suffix\":\"tc\"},{\"video_file_id\":\"file_v1\",\"subtitle_file_id\":\"file_s2\",\"confidence\":0.95,\"match_factors\":[\"content_correlation\"],\"language\":\"sc\",\"target_filename_suffix\":\"sc\"}],\"confidence\":0.95,\"reasoning\":\"Distinguished by script\"}</output>\n");
p.push_str("</example>\n");
p
}
/// Parse matching results from AI response.
pub fn parse_match_result_base(response: &str) -> Result<MatchResult> {
let json_start = response.find('{').unwrap_or(0);
let json_end = response.rfind('}').map(|i| i + 1).unwrap_or(response.len());
let json_str = &response[json_start..json_end];
serde_json::from_str(json_str)
.map_err(|e| SubXError::AiService(format!("AI response parsing failed: {}", e)))
}
/// Build verification prompt for AI providers.
pub fn build_verification_prompt_base(request: &VerificationRequest) -> String {
let mut p = String::new();
p.push_str("<role>You are an expert subtitle-matching verifier.</role>\n");
p.push_str("<instructions>\n");
p.push_str("## Task\n");
p.push_str("Score the supplied video/subtitle pairing on a 0.0–1.0 confidence scale and return strictly JSON.\n\n");
p.push_str("## Output Schema\n");
p.push_str("Respond with a single JSON object: `{\"score\": <0..1>, \"factors\": [\"...\"]}`. No prose, no Markdown.\n");
p.push_str("</instructions>\n");
p.push_str("<match>\n");
p.push_str(&format!(
" <video path=\"{}\"/>\n",
xml_escape_attr(&request.video_file)
));
p.push_str(&format!(
" <subtitle path=\"{}\"/>\n",
xml_escape_attr(&request.subtitle_file)
));
p.push_str(" <factors>\n");
if request.match_factors.is_empty() {
// Preserve a recognizable "Matching factors:" anchor for tests
// and operators reading raw prompts.
p.push_str(" <!-- Matching factors: (none) -->\n");
} else {
for factor in &request.match_factors {
p.push_str(&format!(
" <factor>{}</factor>\n",
xml_escape_attr(factor)
));
}
p.push_str(" <!-- Matching factors: list above -->\n");
}
p.push_str(" </factors>\n");
p.push_str("</match>\n");
p.push_str("<output_schema>{\"score\": 0.9, \"factors\": [\"...\"]}</output_schema>\n");
// Keep a literal "JSON format" anchor for legacy assertions.
p.push_str("<!-- Respond strictly in JSON format. -->\n");
p
}
/// Parse confidence score from AI response.
pub fn parse_confidence_score_base(response: &str) -> Result<ConfidenceScore> {
let json_start = response.find('{').unwrap_or(0);
let json_end = response.rfind('}').map(|i| i + 1).unwrap_or(response.len());
let json_str = &response[json_start..json_end];
serde_json::from_str(json_str)
.map_err(|e| SubXError::AiService(format!("AI confidence parsing failed: {}", e)))
}
#[cfg(test)]
mod tests {
use crate::services::ai::prompts::{
PromptBuilder, ResponseParser, build_analysis_prompt_base, parse_match_result_base,
};
use crate::services::ai::{AnalysisRequest, ContentSample, OpenAIClient};
#[test]
fn test_ai_prompt_with_file_ids_english() {
let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
let video_id = "file_019dcc51-f7da-74e3-9e0d-f75d40fc569c";
let subtitle_id = "file_019dcc51-f7d5-7640-8bb1-d2bbbc127a23";
let request = AnalysisRequest {
video_files: vec![format!("ID:{video_id} | Name:movie.mkv | Path:movie.mkv")],
subtitle_files: vec![format!(
"ID:{subtitle_id} | Name:movie.srt | Path:movie.srt"
)],
content_samples: vec![],
};
let prompt = client.build_analysis_prompt(&request);
assert!(prompt.contains(&format!("id=\"{video_id}\"")));
assert!(prompt.contains(&format!("id=\"{subtitle_id}\"")));
assert!(prompt.contains("<role>"));
assert!(prompt.contains("<instructions>"));
assert!(prompt.contains("<video_files>"));
assert!(prompt.contains("<subtitle_files>"));
assert!(prompt.contains("<output_schema>"));
assert!(prompt.contains("Language Inference Rules"));
assert!(prompt.contains("Naming and Uniqueness Rules"));
assert!(prompt.contains("video_file_id"));
assert!(prompt.contains("subtitle_file_id"));
assert!(!prompt.contains("請分析"));
}
#[test]
fn test_parse_match_result_with_ids() {
let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
let video_id = "file_019dcc51-f7da-74e3-9e0d-f75d40fc569c";
let subtitle_id = "file_019dcc51-f7d5-7640-8bb1-d2bbbc127a23";
let json_resp = format!(
r#"{{
"matches": [{{
"video_file_id": "{video_id}",
"subtitle_file_id": "{subtitle_id}",
"confidence": 0.95,
"match_factors": ["filename_similarity"]
}}],
"confidence": 0.9,
"reasoning": "Strong match based on filename patterns"
}}"#
);
let result = client.parse_match_result(&json_resp).unwrap();
assert_eq!(result.matches.len(), 1);
assert_eq!(result.matches[0].video_file_id, video_id);
assert_eq!(result.matches[0].subtitle_file_id, subtitle_id);
assert_eq!(result.matches[0].confidence, 0.95);
assert_eq!(result.matches[0].match_factors[0], "filename_similarity");
assert!(result.matches[0].language.is_none());
assert!(result.matches[0].target_filename_suffix.is_none());
}
#[test]
fn test_ai_prompt_structure_consistency() {
let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
let request = AnalysisRequest {
video_files: vec![
"ID:file_video1 | Name:video1.mkv | Path:season1/video1.mkv".into(),
"ID:file_video2 | Name:video2.mkv | Path:season1/video2.mkv".into(),
],
subtitle_files: vec![
"ID:file_sub1 | Name:sub1.srt | Path:season1/sub1.srt".into(),
"ID:file_sub2 | Name:sub2.srt | Path:season1/sub2.srt".into(),
],
content_samples: vec![],
};
let prompt = client.build_analysis_prompt(&request);
assert!(prompt.contains("id=\"file_video1\""));
assert!(prompt.contains("id=\"file_video2\""));
assert!(prompt.contains("id=\"file_sub1\""));
assert!(prompt.contains("id=\"file_sub2\""));
assert!(prompt.contains("<video_files>"));
assert!(prompt.contains("<subtitle_files>"));
assert!(prompt.contains("<output_schema>"));
}
#[test]
fn test_parse_confidence_score() {
let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
let json_resp = r#"{
"score": 0.88,
"factors": ["filename_similarity", "content_correlation"]
}"#;
let result = client.parse_confidence_score(json_resp).unwrap();
assert_eq!(result.score, 0.88);
assert_eq!(
result.factors,
vec![
"filename_similarity".to_string(),
"content_correlation".to_string()
]
);
}
#[test]
fn test_xml_escapes_filename_with_ampersand() {
let request = AnalysisRequest {
video_files: vec!["ID:vid1 | Name:m & m.mkv | Path:/x/m & m.mkv".into()],
subtitle_files: vec!["ID:sub1 | Name:m & m.srt | Path:/x/m & m.srt".into()],
content_samples: vec![],
};
let prompt = build_analysis_prompt_base(&request);
assert!(
prompt.contains("name=\"m & m.mkv\""),
"ampersand must be XML-escaped"
);
assert!(prompt.contains("path=\"/x/m & m.mkv\""));
assert!(!prompt.contains("name=\"m & m.mkv\""));
}
#[test]
fn test_xml_escapes_all_five_entities() {
// `xml_escape_attr` covers `&`, `<`, `>`, `"`, and `'`. A
// realistic filename can contain every one at once; verify they
// all materialize as their entity references and that the raw
// characters disappear from the rendered attribute.
let raw_name = "weird & <stuff> \"x\" 'y'.mkv";
let raw_path = "/x/weird & <stuff> \"x\" 'y'.mkv";
let request = AnalysisRequest {
video_files: vec![format!("ID:vid1 | Name:{} | Path:{}", raw_name, raw_path)],
subtitle_files: vec!["ID:sub1 | Name:s.srt | Path:/x/s.srt".into()],
content_samples: vec![],
};
let prompt = build_analysis_prompt_base(&request);
let expected_attr = "weird & <stuff> "x" 'y'.mkv";
assert!(
prompt.contains(&format!("name=\"{}\"", expected_attr)),
"all five XML entities must be escaped: prompt was\n{prompt}"
);
// Path attribute uses the same escaper.
let expected_path = "/x/weird & <stuff> "x" 'y'.mkv";
assert!(prompt.contains(&format!("path=\"{}\"", expected_path)));
// The raw, unescaped form must be absent from any attribute.
assert!(!prompt.contains(&format!("name=\"{}\"", raw_name)));
}
#[test]
fn test_sample_uses_id_only_no_filename_attribute() {
// The `<sample>` element must key off `subtitle_file_id` alone;
// exposing `filename=` would weaken the ID-based contract and
// tempt the model to echo filenames as IDs. Also assert the
// dedicated instruction sentence describing the ID linkage is
// present.
let request = AnalysisRequest {
video_files: vec!["ID:vid1 | Name:movie.mkv | Path:/x/movie.mkv".into()],
subtitle_files: vec!["ID:sub1 | Name:movie.srt | Path:/x/movie.srt".into()],
content_samples: vec![ContentSample {
filename: "movie.srt".into(),
subtitle_file_id: "sub1".into(),
content_preview: "hi".into(),
file_size: 2,
}],
};
let prompt = build_analysis_prompt_base(&request);
assert!(
prompt.contains("<sample subtitle_file_id=\"sub1\">"),
"sample must use only subtitle_file_id; got\n{prompt}"
);
assert!(
!prompt.contains("filename=\""),
"sample element must not carry a filename attribute"
);
assert!(
prompt.contains("<content_samples_instructions>"),
"ID-linkage instruction sentence must be emitted"
);
assert!(prompt.contains("subtitle_file_id"));
}
#[test]
fn test_cdata_preserves_italic_markup() {
let request = AnalysisRequest {
video_files: vec!["ID:vid1 | Name:movie.mkv | Path:/x/movie.mkv".into()],
subtitle_files: vec!["ID:sub1 | Name:movie.srt | Path:/x/movie.srt".into()],
content_samples: vec![ContentSample {
filename: "movie.srt".into(),
subtitle_file_id: "sub1".into(),
content_preview: "Hello <i>world</i>".into(),
file_size: 10,
}],
};
let prompt = build_analysis_prompt_base(&request);
assert!(prompt.contains("subtitle_file_id=\"sub1\""));
assert!(prompt.contains("<![CDATA[Hello <i>world</i>]]>"));
}
#[test]
fn test_cdata_neutralizes_terminator() {
let request = AnalysisRequest {
video_files: vec!["ID:vid1 | Name:movie.mkv | Path:/x/movie.mkv".into()],
subtitle_files: vec!["ID:sub1 | Name:movie.srt | Path:/x/movie.srt".into()],
content_samples: vec![ContentSample {
filename: "movie.srt".into(),
subtitle_file_id: "sub1".into(),
content_preview: "weird ]]> sequence".into(),
file_size: 10,
}],
};
let prompt = build_analysis_prompt_base(&request);
// The literal "]]>" must not appear inside the original content
// CDATA without being split.
let cdata_start = prompt.find("<![CDATA[weird ").unwrap();
let split_marker = "]]]]><![CDATA[>";
assert!(
prompt[cdata_start..].contains(split_marker),
"raw ]]> must be neutralized"
);
}
#[test]
fn test_round_trip_with_optional_fields() {
let json_resp = r#"{
"matches": [{
"video_file_id": "vid1",
"subtitle_file_id": "sub1",
"confidence": 0.9,
"match_factors": ["content_correlation"],
"language": "tc",
"target_filename_suffix": "tc"
}],
"confidence": 0.9,
"reasoning": "ok"
}"#;
let result = parse_match_result_base(json_resp).unwrap();
assert_eq!(result.matches[0].language.as_deref(), Some("tc"));
assert_eq!(
result.matches[0].target_filename_suffix.as_deref(),
Some("tc")
);
}
}