ppsspp/Tools/langtool/src/main.rs

use std::io;

use std::collections::{BTreeMap, BTreeSet};

mod section;
use section::{SAME_COMMENT, Section, line_value, marked_same, split_comment};

mod inifile;
use inifile::IniFile;

mod ai;
mod chatgpt;
mod claude;
use clap::Parser;

mod util;
mod validate;

use crate::{
    ai::{Ai, Provider},
    section::split_line,
    util::ask_letter,
};

#[derive(Parser, Debug)]
struct Args {
    #[command(subcommand)]
    cmd: Command,
    #[arg(short, long)]
    dry_run: bool,
    #[arg(short, long)]
    verbose: bool,
    /// Which AI service to use for translation. Defaults to whichever one we have an API key for.
    #[arg(short, long, value_enum)]
    provider: Option<Provider>,
    /// Model to use, see ai.rs for examples. Defaults to a reasonable one for the provider.
    #[arg(short, long)]
    model: Option<String>,
}

#[derive(Parser, Debug)]
enum Command {
    CopyMissingLines {
        #[arg(short, long)]
        dont_comment_missing: bool,
    },
    ListUnknownLines {},
    CommentUnknownLines {},
    RemoveUnknownLines {},
    AddNewKey {
        section: String,
        key: String,
    },
    AddNewKeyAI {
        section: String,
        key: String,
        extra: Option<String>,
        #[arg(short, long)]
        overwrite_translated: bool,
    },
    AddNewKeyValueAI {
        section: String,
        key: String,
        value: String,
        extra: Option<String>,
        #[arg(short, long)]
        overwrite_translated: bool,
    },
    AddNewKeyValue {
        section: String,
        key: String,
        value: String,
    },
    MoveKey {
        old: String,
        new: String,
        key: String,
    },
    CopyKey {
        old_section: String,
        new_section: String,
        key: String,
    },
    DupeKey {
        section: String,
        old: String,
        new: String,
    },
    RenameKey {
        section: String,
        old: String,
        new: String,
    },
    SortSection {
        section: String,
    },
    RemoveKey {
        section: String,
        key: String,
    },
    GetNewKeys,
    ImportSingle {
        filename: String,
        section: String,
        key: String,
    },
    CheckRefKeys,
    FixupRefKeys, // This was too big a job.
    FinishLanguageWithAI {
        language: String,
        section: Option<String>,
    },
    RemoveLinebreaks {
        section: String,
        key: String,
    },
    RemoveAmpersands {
        section: String,
        key: String,
    },
    ApplyRegex {
        section: String,
        key: String,
        pattern: String,
        replacement: Option<String>,
    },
    SplitKey {
        section: String,
        key: String,
    },
    /// Check every translated string against the English one, see validate.rs.
    Validate,
}

fn copy_missing_lines(
    reference_ini: &IniFile,
    target_ini: &mut IniFile,
    comment_missing: bool,
) -> io::Result<()> {
    for reference_section in &reference_ini.sections {
        // Insert any missing full sections.
        if !target_ini.insert_section_if_missing(reference_section) {
            if let Some(target_section) = target_ini.get_section_mut(&reference_section.name) {
                for line in &reference_section.lines {
                    target_section.insert_line_if_missing(line);
                }

                //target_section.remove_lines_if_not_in(reference_section);
                if comment_missing {
                    target_section.comment_out_lines_if_not_in(reference_section);
                }
            }
        } else {
            // Note: insert_section_if_missing will copy the entire section,
            // no need to loop over the lines here.
            println!("Inserted missing section: {}", reference_section.name);
        }
    }
    Ok(())
}

enum UnknownLineAction {
    Remove,
    Comment,
    Log,
}

fn deal_with_unknown_lines(
    reference_ini: &IniFile,
    target_ini: &mut IniFile,
    action: UnknownLineAction,
) -> io::Result<()> {
    for reference_section in &reference_ini.sections {
        if let Some(target_section) = target_ini.get_section_mut(&reference_section.name) {
            match action {
                UnknownLineAction::Remove => {
                    target_section.remove_lines_if_not_in(reference_section)
                }
                UnknownLineAction::Comment => {
                    target_section.comment_out_lines_if_not_in(reference_section)
                }
                UnknownLineAction::Log => {
                    let unknown_lines = target_section.get_lines_not_in(reference_section);
                    if !unknown_lines.is_empty() {
                        println!("Unknown lines in section [{}]:", target_section.name);
                        for line in unknown_lines {
                            println!("  {}", line);
                        }
                    }
                }
            }
        }
    }
    Ok(())
}

fn print_keys_if_not_in(
    reference_ini: &IniFile,
    target_ini: &mut IniFile,
    header: &str,
) -> io::Result<()> {
    for reference_section in &reference_ini.sections {
        if let Some(target_section) = target_ini.get_section_mut(&reference_section.name) {
            let keys = target_section.get_keys_if_not_in(reference_section);
            if !keys.is_empty() {
                println!("{} ({})", reference_section.name, header);
                for key in &keys {
                    println!("- {key}");
                }
            }
        }
    }
    Ok(())
}

fn move_key(target_ini: &mut IniFile, old: &str, new: &str, key: &str) -> io::Result<()> {
    if let Some(old_section) = target_ini.get_section_mut(old) {
        if let Some(line) = old_section.remove_line(key) {
            if let Some(new_section) = target_ini.get_section_mut(new) {
                new_section.insert_line_if_missing(&line);
            } else {
                println!("No new section {new}");
            }
        } else {
            println!("No key {key} in section {old}");
        }
    } else {
        println!("No old section {old}");
    }
    Ok(())
}

fn copy_key(target_ini: &mut IniFile, old: &str, new: &str, key: &str) -> io::Result<()> {
    if let Some(old_section) = target_ini.get_section_mut(old) {
        if let Some(line) = old_section.get_line(key) {
            if let Some(new_section) = target_ini.get_section_mut(new) {
                new_section.insert_line_if_missing(&line);
            } else {
                println!("No new section {new}");
            }
        } else {
            println!("No key {key} in section {old}");
        }
    } else {
        println!("No old section {old}");
    }
    Ok(())
}

fn remove_key(target_ini: &mut IniFile, section: &str, key: &str) -> io::Result<()> {
    if let Some(old_section) = target_ini.get_section_mut(section) {
        old_section.remove_line(key);
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn remove_linebreaks(target_ini: &mut IniFile, section: &str, key: &str) -> io::Result<()> {
    if let Some(old_section) = target_ini.get_section_mut(section) {
        old_section.remove_linebreaks(key);
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn remove_ampersands(target_ini: &mut IniFile, section: &str, key: &str) -> io::Result<()> {
    if let Some(old_section) = target_ini.get_section_mut(section) {
        old_section.remove_ampersands(key);
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn add_new_key(
    target_ini: &mut IniFile,
    section: &str,
    key: &str,
    value: &str,
    overwrite_translated: bool,
) -> io::Result<()> {
    if let Some(section) = target_ini.get_section_mut(section) {
        if !overwrite_translated {
            if let Some(existing_value) = section.get_value(key)
                && existing_value != key
            {
                // This one was already translated. Skip it.
                println!(
                    "Key '{key}' already has a translated value '{existing_value}', skipping."
                );
                return Ok(());
            }
        }
        section.insert_line_if_missing(&format!("{key} = {value}"));
    } else {
        println!("No section {section}");
    }
    Ok(())
}

/// Runs the validation checks over a whole language file, printing what it finds.
/// Returns the number of problems, so the caller can set the exit code.
fn validate_ini(target_ini: &IniFile, reference_ini: &IniFile, language: &str) -> usize {
    let mut count = 0;
    for section in &target_ini.sections {
        let Some(ref_section) = reference_ini.get_section(&section.name) else {
            continue;
        };
        for line in &section.lines {
            let Some((key, _)) = split_line(line) else {
                continue;
            };
            let (Some(value), Some(ref_value)) = (section.get_value(key), ref_section.get_value(key))
            else {
                continue;
            };
            for issue in validate::check(&ref_value, &value) {
                println!("{language} [{}] {key}: {issue}", section.name);
                println!("      en: {ref_value}");
                println!("   {language}: {value}");
                count += 1;
            }
        }
    }
    count
}

fn check_keys(target_ini: &IniFile) -> io::Result<()> {
    for section in &target_ini.sections {
        let mut mismatches = Vec::new();

        if section.name == "DesktopUI" {
            // ignore the ampersands for now
            continue;
        }

        for line in &section.lines {
            if let Some((key, value)) = split_line(line) {
                if key != value {
                    mismatches.push((key, value));
                }
            }
        }

        if !mismatches.is_empty() {
            println!("[{}]", section.name);
            for (key, value) in mismatches {
                println!("  {key} != {value}");
            }
        }
    }
    Ok(())
}

fn fixup_keys(target_ini: &IniFile, dry_run: bool) -> io::Result<()> {
    for section in &target_ini.sections {
        let mut mismatches = Vec::new();

        if section.name == "DesktopUI"
            || section.name == "MappableControls"
            || section.name == "PostShaders"
        {
            // ignore the ampersands for now, also mappable controls, we don't want to change those strings.
            continue;
        }

        for line in &section.lines {
            if let Some((key, value)) = split_line(line)
                && key != value
            {
                mismatches.push((key, value));
            }
        }

        if !mismatches.is_empty() {
            println!("[{}]", section.name);
            for (key, value) in mismatches {
                if (key.len() as i32 - value.len() as i32).abs() > 15 {
                    println!("  (skipping {key} = {value} (probably an alias))");
                    continue;
                }
                if value.contains(r"\n") {
                    println!("  (skipping {key} = {value} (line break))");
                    continue;
                }
                if value.contains("×") || value.contains("\"") {
                    println!("  (skipping {key} = {value} (symbol))");
                    continue;
                }
                if key.contains("ardboard") {
                    println!("  (skipping {key} = {value} (cardboard))");
                    continue;
                }
                if key.contains("translators") {
                    continue;
                }

                let _ = cli_clipboard::set_contents(format!("\"{key}\""));

                match ask_letter(&format!("  '{key}' != '{value}' ? >\n"), "ynrd") {
                    'y' => execute_command(
                        Command::RenameKey {
                            section: section.name.clone(),
                            old: key.to_string(),
                            new: value.to_string(),
                        },
                        None,
                        dry_run,
                        false,
                    ),
                    'r' => {
                        println!("reverse fixup not supported yet");
                    }
                    'q' => {
                        println!("Cancelled! Quitting.");
                        return Ok(());
                    }
                    'd' => execute_command(
                        Command::RemoveKey {
                            section: section.name.clone(),
                            key: key.to_string(),
                        },
                        None,
                        dry_run,
                        false,
                    ),

                    'n' => {}
                    _ => {
                        println!("Invalid response, ignoring.");
                    }
                }
            }
        }
    }
    Ok(())
}

fn finish_language_with_ai(
    target_language: &str,
    target_ini: &mut IniFile,
    ref_ini: &IniFile,
    section: Option<&str>,
    ai: &Ai,
    dry_run: bool,
) -> anyhow::Result<()> {
    println!("Finishing language with {}", ai.description());
    println!(
        "Step 1: Compare all strings in the section with the matching strings from the reference."
    );

    let sections: Vec<Section> = if let Some(section_name) = section {
        vec![target_ini.get_section(section_name).unwrap().clone()]
    } else {
        target_ini.sections.to_vec()
    };

    let base_prompt = format!(
        "Please translate the below list of strings from US English to {target_language}.
    After the strings to translate, there are related already-translated strings that may help for context.
    Note that the strings are UI strings for my PSP emulator application.
    Also, please output similarly to the input, with section headers and key=value pairs. The section name
    is not to be translated.

    Here are the strings to translate:
    "
    );
    let suffix = "    Do not output any text before or after the list of translated strings, do not ask followups.
    'undo state' means a saved state that's been saved so that the last save state operation can be undone.
    DO NOT translate strings like DDMMYYYY, MMDDYYYY and similar technical letters and designations. Not even
    translating the individual letters, they need to be kept as-is.
    'JIT using IR' should be interpreted as 'JIT, with IR'.
    Don't translate strings about 'load undo state' or 'save undo state', also not about savestate slots.
    IMPORTANT! 'Notification screen position' means the position on the screen where notifications are displayed,
    not the position of a 'notification screen', no such thing.
    %1 is a placeholder for a number or word, do not change it, just make sure it ends up in the right location.
    A 'driver manager' is a built-in tool to manage drivers, not a human boss. Same goes for other types of manager.
    The '=' at the end of the lines to translate is not part of the translation keys.
    ";

    for section in sections {
        let Some(ref_section) = ref_ini.get_section(&section.name) else {
            println!("Section '{}' not found in reference file", section.name);
            continue;
        };
        let mut alias_map = BTreeMap::new();
        let mut alias_inverse_map: BTreeMap<String, Vec<&str>> = BTreeMap::new();
        for line in &ref_section.lines {
            if let Some((key, value)) = split_line(line) {
                // We actually process almost everything here, we could check for case but we don't
                // since the aliased case is better.
                if key != value {
                    println!("Saving alias: {key} = {value}");
                    alias_map.insert(key, value.to_string());
                    // Several keys can share the same English string (like "Texture Filter" and
                    // "Texture Filtering" in [Graphics]), so this can't be a plain value -> key
                    // map, we'd silently lose all but the last of them.
                    alias_inverse_map
                        .entry(value.to_string())
                        .or_default()
                        .push(key);
                }
            }
        }

        // When just testing aliases.
        // return Ok(());

        let mut untranslated_keys = vec![];
        let mut translated_keys = vec![];
        for line in &section.lines {
            if let Some((key, raw_value)) = split_line(line) {
                let (value, comment) = split_comment(raw_value);
                if let Some(ref_value) = ref_section.get_value(key) {
                    if marked_same(comment) {
                        // Settled: this one stays as the English string, don't ask about it again.
                        continue;
                    }
                    if value == ref_value {
                        // Key not translated.
                        // However, we need to reject some things that the AI likes to mishandle.
                        if value.to_uppercase() == value {
                            println!(
                                "Skipping untranslated key '{}' with uppercase value '{}'",
                                key, value
                            );
                            continue;
                        }
                        untranslated_keys.push((key, ref_value));
                    } else {
                        translated_keys.push((key, value));
                    }
                } else {
                    println!(
                        "Key '{}' not found in reference section '{}'",
                        key, ref_section.name
                    );
                }
            }
        }

        println!(
            "[{}]: Found {} untranslated keys",
            section.name,
            untranslated_keys.len()
        );
        if untranslated_keys.is_empty() {
            continue;
        }

        for (key, ref_value) in &untranslated_keys {
            println!(" - '{} (ref: '{}')", key, ref_value);
        }

        // Here you would call the AI to translate the keys.
        let section_prompt = format!(
            "{base_prompt}\n\n[{}]\n{}\n\n\n\nBelow are the already translated strings for context, don't re-translate these:\n\n{}\n\n{}",
            section.name,
            untranslated_keys
                .iter()
                .map(|(k, _v)| format!("{} = ", alias_map.get(k).unwrap_or(&k.to_string())))
                // Deduped, since keys that share an English string produce the same line.
                .collect::<BTreeSet<String>>()
                .into_iter()
                .collect::<Vec<String>>()
                .join("\n"),
            translated_keys
                .iter()
                .map(|(k, v)| format!("{} = {}", k, v))
                .collect::<Vec<String>>()
                .join("\n"),
            suffix
        );
        println!("[{}] AI prompt:\n{}", section.name, section_prompt);
        if !dry_run {
            println!("Running AI translation...");
            let response = ai
                .chat(&section_prompt)
                .map_err(|e| anyhow::anyhow!("chat failed: {e}"))?;
            println!("AI response:\n{}", response);
            // Now we just need to merge the AI response into the target_ini.
            let parsed_response = IniFile::parse_string(&response)
                .map_err(|e| anyhow::anyhow!("Failed to parse AI response: {e}"))?;
            if parsed_response.sections.is_empty() {
                println!("No sections found in AI response! bad!");
            }
            let target_section = target_ini.get_section_mut(&section.name).unwrap();
            for parsed_section in parsed_response.sections {
                if parsed_section.name == section.name {
                    println!("Merging AI response for section '{}'", parsed_section.name);
                    for line in &parsed_section.lines {
                        if let Some((key, value)) = split_line(line) {
                            // Put the key through the inverse alias map. Can give us more than one
                            // key, if they share the same English string - they all get the same
                            // translation, which is what we want.
                            let itself = vec![key];
                            let original_keys = alias_inverse_map.get(key).unwrap_or(&itself);
                            for original_key in original_keys {
                                let ref_value =
                                    ref_section.get_value(original_key).unwrap_or_default();
                                let issues = validate::check_ai_translation(&ref_value, value);
                                if issues == vec![validate::Issue::Untranslated] {
                                    // The AI handed the English string back unchanged, which for
                                    // things like "Vsync" or "Ad Hoc multiplayer" is the right
                                    // answer. Write that down so we stop asking every run.
                                    println!("Marking '{original_key}' as not needing translation");
                                    target_section.set_value(
                                        original_key,
                                        value,
                                        Some(SAME_COMMENT),
                                    );
                                    continue;
                                }
                                if !issues.is_empty() {
                                    for issue in issues {
                                        println!("Rejecting '{original_key}' = '{value}': {issue}");
                                    }
                                    continue;
                                }
                                print!("Updating '{}': {}", original_key, value);
                                if key != *original_key {
                                    println!(" ({})", key);
                                } else {
                                    println!();
                                }
                                if !target_section.set_value(
                                    original_key,
                                    value,
                                    Some("AI translated"),
                                ) {
                                    println!("Failed to update '{}'", original_key);
                                }
                            }
                        }
                    }
                } else {
                    println!("Mismatched section name '{}'", parsed_section.name);
                }
            }
        }
    }

    Ok(())
}

fn rename_key(target_ini: &mut IniFile, section: &str, old: &str, new: &str) -> io::Result<()> {
    if let Some(section) = target_ini.get_section_mut(section) {
        section.rename_key(old, new);
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn apply_regex(
    target_ini: &mut IniFile,
    section: &str,
    key: &str,
    pattern: &str,
    replacement: &str,
) -> io::Result<()> {
    if let Some(section) = target_ini.get_section_mut(section) {
        section.apply_regex(key, pattern, replacement);
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn split_key(target_ini: &mut IniFile, section: &str, key: &str) -> io::Result<()> {
    if let Some(section) = target_ini.get_section_mut(section) {
        section.split_key(key);
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn dupe_key(target_ini: &mut IniFile, section: &str, old: &str, new: &str) -> io::Result<()> {
    if let Some(section) = target_ini.get_section_mut(section) {
        section.dupe_key(old, new);
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn sort_section(target_ini: &mut IniFile, section: &str) -> io::Result<()> {
    if let Some(section) = target_ini.get_section_mut(section) {
        section.sort();
    } else {
        println!("No section {section}");
    }
    Ok(())
}

fn generate_prompt(filenames: &[String], section: &str, value: &str, extra: &str) -> String {
    let languages = filenames
        .iter()
        .map(|filename| filename.split_once(".").unwrap().0)
        .collect::<Vec<&str>>()
        .join(", ");

    let base_str = format!("Please translate '{value}' from US English to all of these languages: {languages}.
    Output in json format, a single dictionary, key=value. Include en_US first (the original string).
    For context, the string will be in the translation section '{section}', and these strings are UI strings for my PSP emulator application.
    {extra}
    Keep the strings relatively short, try not to let them become more than 60% longer than the original string unless necessary.
    Do not output any text before or after the list of translated strings, do not ask followups.
    ");

    base_str
}

fn parse_response(response: &str) -> Option<BTreeMap<String, String>> {
    // Try to find JSON in the response (it might have other text around it)
    let response = response.trim();

    // Look for JSON object boundaries
    let start = response.find('{')?;
    let end = response.rfind('}')? + 1;
    let json_str = &response[start..end];

    // Parse the JSON into a BTreeMap
    match serde_json::from_str::<BTreeMap<String, serde_json::Value>>(json_str) {
        Ok(json_map) => {
            let mut result = BTreeMap::new();
            for (key, value) in json_map {
                // Convert JSON values to strings
                let string_value = match value {
                    serde_json::Value::String(s) => s,
                    _ => value.to_string().trim_matches('"').to_string(),
                };
                result.insert(key, string_value);
            }
            Some(result)
        }
        Err(e) => {
            eprintln!("Failed to parse JSON response: {}", e);
            eprintln!("JSON string was: {}", json_str);
            None
        }
    }
}

fn main() {
    let opt = Args::parse();

    let ai = Ai::create(opt.provider, opt.model);

    // TODO: Grab extra arguments from opt somehow.
    // let args: Vec<String> = vec![]; //std::env::args().skip(1).collect();
    execute_command(opt.cmd, ai.as_ref(), opt.dry_run, opt.verbose);
}

fn execute_command(cmd: Command, ai: Option<&Ai>, dry_run: bool, verbose: bool) {
    let root = "../../assets/lang";
    let reference_ini_filename = "en_US.ini";

    let reference_ini = IniFile::parse_file(&format!("{root}/{reference_ini_filename}")).unwrap();

    let mut filenames = Vec::new();
    if filenames.is_empty() {
        // Grab them all.
        for path in std::fs::read_dir(root).unwrap() {
            let path = path.unwrap();
            let filename = path.file_name();
            let filename = filename.to_string_lossy();
            if !filename.ends_with(".ini") {
                continue;
            }
            filenames.push(path.file_name().to_string_lossy().to_string());
        }
    }

    let mut single_ini_section: Option<Section> = None;
    if let Command::ImportSingle {
        filename,
        section,
        key: _,
    } = &cmd
    {
        if let Ok(single_ini) = IniFile::parse_file(filename) {
            if let Some(single_section) = single_ini.get_section("Single") {
                single_ini_section = Some(single_section.clone());
            } else {
                println!("No section {section} in {filename}");
            }
        } else {
            println!("Failed to parse {filename}");
            return;
        }
    }

    if let Command::FinishLanguageWithAI { language, section } = &cmd {
        if let Some(ai) = &ai {
            let target_ini_filename = format!("{root}/{language}.ini");
            let mut target_ini = IniFile::parse_file(&target_ini_filename).unwrap();
            finish_language_with_ai(
                language,
                &mut target_ini,
                &reference_ini,
                section.as_deref(),
                ai,
                dry_run,
            )
            .unwrap();
            if !dry_run {
                println!("Writing modified file for target language: {}", language);
                target_ini.write().unwrap();
            }
        } else {
            println!("FinishLanguageWithAI: AI key not set, skipping AI command.");
        }
        return;
    }

    // This is a bit ugly, but we need to generate the AI response before processing files.
    let ai_response = if let Command::AddNewKeyAI {
        section,
        key,
        extra,
        overwrite_translated: _,
    } = &cmd
    {
        match generate_ai_response(ai, &filenames, section, key, extra) {
            Some(value) => value,
            None => return,
        }
    } else if let Command::AddNewKeyValueAI {
        section,
        key: _, // We don't need the key here, it's used later when writing to the ini file.
        value,
        extra,
        overwrite_translated: _,
    } = &cmd
    {
        match generate_ai_response(ai, &filenames, section, value, extra) {
            Some(value) => value,
            None => return,
        }
    } else {
        None
    };

    let mut issue_count = 0;

    for filename in &filenames {
        let reference_ini = &reference_ini;
        if filename == "langtool" {
            // Get this from cargo run for some reason.
            continue;
        }
        let target_ini_filename = format!("{root}/{filename}");
        if verbose {
            println!("Langtool processing {target_ini_filename}");
        }

        let mut target_ini = IniFile::parse_file(&target_ini_filename).unwrap();

        let is_reference = filename == reference_ini_filename;

        match cmd {
            Command::ApplyRegex {
                ref section,
                ref key,
                ref pattern,
                ref replacement,
            } => {
                apply_regex(
                    &mut target_ini,
                    section,
                    key,
                    pattern,
                    replacement.as_ref().unwrap_or(&"".to_string()),
                )
                .unwrap();
            }
            Command::SplitKey {
                ref section,
                ref key,
            } => {
                split_key(&mut target_ini, section, key).unwrap();
            }
            Command::Validate => {
                if !is_reference {
                    let language = filename.split_once('.').unwrap().0;
                    issue_count += validate_ini(&target_ini, reference_ini, language);
                }
            }
            Command::FinishLanguageWithAI {
                language: _,
                section: _,
            } => {}
            Command::FixupRefKeys => {
                if is_reference {
                    fixup_keys(&target_ini, dry_run).unwrap();
                }
            }
            Command::CheckRefKeys => {
                if is_reference {
                    check_keys(&target_ini).unwrap();
                }
            }
            Command::CopyMissingLines {
                dont_comment_missing,
            } => {
                copy_missing_lines(reference_ini, &mut target_ini, !dont_comment_missing).unwrap();
            }
            Command::CommentUnknownLines {} => {
                if !is_reference {
                    deal_with_unknown_lines(
                        reference_ini,
                        &mut target_ini,
                        UnknownLineAction::Comment,
                    )
                    .unwrap();
                }
            }
            Command::RemoveUnknownLines {} => {
                if !is_reference {
                    deal_with_unknown_lines(
                        reference_ini,
                        &mut target_ini,
                        UnknownLineAction::Remove,
                    )
                    .unwrap();
                }
            }
            Command::ListUnknownLines {} => {
                if !is_reference {
                    deal_with_unknown_lines(reference_ini, &mut target_ini, UnknownLineAction::Log)
                        .unwrap();
                }
            }
            Command::GetNewKeys => {
                if !is_reference {
                    print_keys_if_not_in(reference_ini, &mut target_ini, &target_ini_filename)
                        .unwrap();
                }
            }
            Command::SortSection { ref section } => sort_section(&mut target_ini, section).unwrap(),
            Command::RenameKey {
                ref section,
                ref old,
                ref new,
            } => rename_key(&mut target_ini, section, old, new).unwrap(),
            Command::AddNewKey {
                ref section,
                ref key,
            } => add_new_key(&mut target_ini, section, key, key, false).unwrap(),
            Command::AddNewKeyAI {
                ref section,
                ref key,
                ref extra,
                overwrite_translated,
            } => {
                if !is_reference {
                    let lang = filename.split_once('.').unwrap().0;
                    if let Some(ai_response) = &ai_response {
                        // Process it.
                        if let Some(translated_string) = ai_response.get(lang) {
                            println!("{lang}:");
                            add_new_key(
                                &mut target_ini,
                                section,
                                key,
                                &format!("{translated_string} # AI translated"),
                                overwrite_translated,
                            )
                            .unwrap();
                        } else {
                            println!("No usable translation for {lang}, skipping it.");
                            continue;
                        }
                    }
                } else {
                    if ai_response.is_some() {
                        let _ = extra;
                        add_new_key(&mut target_ini, section, key, key, overwrite_translated)
                            .unwrap();
                    }
                }
            }
            Command::AddNewKeyValueAI {
                ref section,
                ref key,
                ref value, // was translated above
                ref extra,
                overwrite_translated,
            } => {
                if !is_reference {
                    let lang = filename.split_once('.').unwrap().0;
                    if let Some(ai_response) = &ai_response {
                        // Process it.
                        if let Some(translated_string) = ai_response.get(lang) {
                            println!("{lang}:");
                            add_new_key(
                                &mut target_ini,
                                section,
                                key,
                                &format!("{translated_string} # AI translated"),
                                overwrite_translated,
                            )
                            .unwrap();
                        } else {
                            println!("No usable translation for {lang}, skipping it.");
                            continue;
                        }
                    }
                } else {
                    if ai_response.is_some() {
                        let _ = extra;
                        add_new_key(&mut target_ini, section, key, value, overwrite_translated)
                            .unwrap();
                    }
                }
            }
            Command::AddNewKeyValue {
                ref section,
                ref key,
                ref value,
            } => add_new_key(&mut target_ini, section, key, value, false).unwrap(),
            Command::MoveKey {
                ref old,
                ref new,
                ref key,
            } => {
                move_key(&mut target_ini, old, new, key).unwrap();
            }
            Command::CopyKey {
                // Copies between sections
                ref old_section,
                ref new_section,
                ref key,
            } => {
                copy_key(&mut target_ini, old_section, new_section, key).unwrap();
            }
            Command::DupeKey {
                ref section,
                ref old,
                ref new,
            } => {
                dupe_key(&mut target_ini, section, old, new).unwrap();
            }
            Command::RemoveKey {
                ref section,
                ref key,
            } => {
                remove_key(&mut target_ini, section, key).unwrap();
            }
            Command::RemoveLinebreaks {
                ref section,
                ref key,
            } => {
                remove_linebreaks(&mut target_ini, section, key).unwrap();
            }
            Command::RemoveAmpersands {
                ref section,
                ref key,
            } => {
                remove_ampersands(&mut target_ini, section, key).unwrap();
            }
            Command::ImportSingle {
                filename: _,
                ref section,
                ref key,
            } => {
                let lang_id = filename.strip_suffix(".ini").unwrap();
                if let Some(single_section) = &single_ini_section {
                    let english = single_section
                        .get_line("en_US")
                        .and_then(|line| line_value(&line).map(|value| value.to_string()));
                    if let Some(target_section) = target_ini.get_section_mut(section) {
                        if let Some(single_line) = single_section.get_line(lang_id) {
                            if let Some(value) = line_value(&single_line) {
                                // en_US is the string the others were translated from, not a
                                // translation, so it gets no comment. A language where the
                                // translation is just the English string gets marked as such, so
                                // nothing tries to translate it again later.
                                let comment = if is_reference {
                                    None
                                } else if Some(value) == english.as_deref() {
                                    Some(SAME_COMMENT)
                                } else {
                                    Some("AI translated")
                                };
                                println!(
                                    "Inserting value {value} for key {key} in section {section} in {target_ini_filename}"
                                );
                                let line = match comment {
                                    Some(comment) => format!("{key} = {value} # {comment}"),
                                    None => format!("{key} = {value}"),
                                };
                                if !target_section.insert_line_if_missing(&line) {
                                    // Didn't insert it, so it exists. We need to replace it.
                                    target_section.set_value(key, value, comment);
                                }
                            }
                        } else {
                            println!("No lang_id {lang_id} in single section");
                        }
                    } else {
                        println!("No section {section} in {target_ini_filename}");
                    }
                } else {
                    println!("No section {section} in {filename}");
                }
            }
        }

        if !dry_run {
            target_ini.write().unwrap();
        }
    }

    if let Command::Validate = cmd {
        println!("Found {issue_count} problems.");
        if issue_count > 0 {
            // Makes this usable as a check in scripts.
            std::process::exit(1);
        }
    }

    // Don't need to additionally process reference here like we did before, this is done
    // in the main loop.
}

fn generate_ai_response(
    ai: Option<&Ai>,
    filenames: &[String],
    section: &str,
    key: &str,
    extra: &Option<String>,
) -> Option<Option<BTreeMap<String, String>>> {
    let prompt = generate_prompt(
        filenames,
        section,
        key,
        &extra.clone().unwrap_or("".to_string()),
    );
    println!("generated prompt:\n{prompt}");
    Some(if let Some(ai) = &ai {
        println!("Using {} for translation...", ai.description());
        let response = ai
            .chat(&prompt)
            .map_err(|e| anyhow::anyhow!("chat failed: {e}"))
            .unwrap();
        println!("AI response: {response}");
        if let Some(mut parsed) = parse_response(&response) {
            println!("Parsed: {:?}", parsed);

            // The English string we asked for a translation of is in the prompt as 'key'.
            parsed.retain(|language, translation| {
                let issues = if language == "en_US" {
                    validate::check(key, translation)
                } else {
                    validate::check_ai_translation(key, translation)
                };
                for issue in &issues {
                    println!("Rejecting {language} translation '{translation}': {issue}");
                }
                issues.is_empty()
            });

            if parsed.len() < filenames.len() {
                println!(
                    "Not enough languages generated! {} vs {}",
                    parsed.len(),
                    filenames.len()
                );
            }

            Some(parsed)
        } else {
            println!("Failed to parse AI response, not doing anything.");
            return None;
        }
    } else {
        println!("AI key not set, skipping AI command.");
        return None;
    })
}