1use crate::{
  2    PredictionProvider,
  3    example::{ActualCursor, Example},
  4    format_prompt::TeacherPrompt,
  5    repair,
  6};
  7use anyhow::{Context as _, Result};
  8use edit_prediction::example_spec::encode_cursor_in_patch;
  9use zeta_prompt::{CURSOR_MARKER, ZetaFormat, parse_zeta2_model_output};
 10
 11pub fn run_parse_output(example: &mut Example) -> Result<()> {
 12    example
 13        .prompt_inputs
 14        .as_ref()
 15        .context("prompt_inputs required")?;
 16
 17    let to_parse: Vec<_> = example
 18        .predictions
 19        .iter()
 20        .enumerate()
 21        .filter(|(_, p)| !p.actual_output.is_empty())
 22        .map(|(ix, p)| (ix, p.actual_output.clone(), p.provider))
 23        .collect();
 24
 25    for (ix, actual_output, provider) in to_parse {
 26        let (actual_patch, actual_cursor) =
 27            parse_prediction_output(example, &actual_output, provider)?;
 28        example.predictions[ix].actual_patch = Some(actual_patch);
 29        example.predictions[ix].actual_cursor = actual_cursor;
 30    }
 31
 32    Ok(())
 33}
 34
 35pub fn parse_prediction_output(
 36    example: &Example,
 37    actual_output: &str,
 38    provider: PredictionProvider,
 39) -> Result<(String, Option<ActualCursor>)> {
 40    match provider {
 41        PredictionProvider::Teacher(_) | PredictionProvider::TeacherNonBatching(_) => {
 42            TeacherPrompt::parse(example, actual_output)
 43        }
 44        PredictionProvider::Zeta2(version) => parse_zeta2_output(example, actual_output, version),
 45        PredictionProvider::Repair => repair::parse(example, actual_output),
 46        _ => anyhow::bail!(
 47            "parse-output only supports Teacher and Zeta2 providers, got {:?}",
 48            provider
 49        ),
 50    }
 51}
 52
 53fn parse_zeta2_output(
 54    example: &Example,
 55    actual_output: &str,
 56    format: ZetaFormat,
 57) -> Result<(String, Option<ActualCursor>)> {
 58    let prompt_inputs = example
 59        .prompt_inputs
 60        .as_ref()
 61        .context("prompt_inputs required")?;
 62
 63    let parsed = parse_zeta2_model_output(actual_output, format, prompt_inputs)?;
 64    let range_in_excerpt = parsed.range_in_excerpt;
 65
 66    let excerpt = prompt_inputs.cursor_excerpt.as_ref();
 67    let old_text = excerpt[range_in_excerpt.clone()].to_string();
 68    let mut new_text = parsed.new_editable_region;
 69
 70    let cursor_offset = if let Some(offset) = new_text.find(CURSOR_MARKER) {
 71        new_text.replace_range(offset..offset + CURSOR_MARKER.len(), "");
 72        Some(offset)
 73    } else {
 74        None
 75    };
 76
 77    // Normalize trailing newlines for diff generation
 78    let mut old_text_normalized = old_text;
 79    if !new_text.is_empty() && !new_text.ends_with('\n') {
 80        new_text.push('\n');
 81    }
 82    if !old_text_normalized.is_empty() && !old_text_normalized.ends_with('\n') {
 83        old_text_normalized.push('\n');
 84    }
 85
 86    let editable_region_offset = range_in_excerpt.start;
 87    let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
 88    let editable_region_lines = old_text_normalized.lines().count() as u32;
 89
 90    let diff = language::unified_diff_with_context(
 91        &old_text_normalized,
 92        &new_text,
 93        editable_region_start_line as u32,
 94        editable_region_start_line as u32,
 95        editable_region_lines,
 96    );
 97
 98    let formatted_diff = format!(
 99        "--- a/{path}\n+++ b/{path}\n{diff}",
100        path = example.spec.cursor_path.to_string_lossy(),
101    );
102
103    let formatted_diff = encode_cursor_in_patch(&formatted_diff, cursor_offset);
104
105    let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
106        ActualCursor::from_editable_region(
107            &example.spec.cursor_path,
108            editable_region_cursor_offset,
109            &new_text,
110            excerpt,
111            editable_region_offset,
112            editable_region_start_line,
113        )
114    });
115
116    Ok((formatted_diff, actual_cursor))
117}