1use crate::{PredictionProvider, example::Example, format_prompt::TeacherPrompt};
  2use anyhow::{Context as _, Result};
  3use zeta_prompt::{CURSOR_MARKER, ZetaVersion};
  4
  5pub fn run_parse_output(example: &mut Example) -> Result<()> {
  6    let provider = example
  7        .prompt
  8        .as_ref()
  9        .context("prompt required (run format-prompt first)")?
 10        .provider;
 11    example
 12        .prompt_inputs
 13        .as_ref()
 14        .context("prompt_inputs required")?;
 15
 16    let parsed_patches: Vec<_> = example
 17        .predictions
 18        .iter()
 19        .enumerate()
 20        .filter(|(_, p)| !p.actual_output.is_empty())
 21        .map(|(ix, prediction)| {
 22            let result = parse_prediction_output(example, &prediction.actual_output, provider);
 23            result.map(|(patch, cursor_offset)| (ix, patch, cursor_offset))
 24        })
 25        .collect::<Result<Vec<_>>>()?;
 26
 27    for (ix, actual_patch, actual_cursor_offset) in parsed_patches {
 28        example.predictions[ix].actual_patch = Some(actual_patch);
 29        example.predictions[ix].actual_cursor_offset = actual_cursor_offset;
 30        example.predictions[ix].provider = provider;
 31    }
 32
 33    Ok(())
 34}
 35
 36pub fn parse_prediction_output(
 37    example: &Example,
 38    actual_output: &str,
 39    provider: PredictionProvider,
 40) -> Result<(String, Option<usize>)> {
 41    match provider {
 42        PredictionProvider::Teacher(_) | PredictionProvider::TeacherNonBatching(_) => {
 43            TeacherPrompt::parse(example, actual_output)
 44        }
 45        PredictionProvider::Zeta2(version) => parse_zeta2_output(example, actual_output, version),
 46        _ => anyhow::bail!(
 47            "parse-output only supports Teacher and Zeta2 providers, got {:?}",
 48            provider
 49        ),
 50    }
 51}
 52
 53fn extract_zeta2_current_region(prompt: &str, version: ZetaVersion) -> Result<String> {
 54    let (current_marker, end_marker) = match version {
 55        ZetaVersion::V0112MiddleAtEnd => ("<|fim_middle|>current\n", "<|fim_middle|>updated"),
 56        ZetaVersion::V0113Ordered | ZetaVersion::V0114180EditableRegion => {
 57            ("<|fim_middle|>current\n", "<|fim_suffix|>")
 58        }
 59        ZetaVersion::V0120GitMergeMarkers => (
 60            zeta_prompt::v0120_git_merge_markers::START_MARKER,
 61            zeta_prompt::v0120_git_merge_markers::SEPARATOR,
 62        ),
 63    };
 64
 65    let start = prompt.find(current_marker).with_context(|| {
 66        format!(
 67            "missing current marker '{}' in prompt",
 68            current_marker.trim()
 69        )
 70    })? + current_marker.len();
 71
 72    let end = prompt[start..]
 73        .find(end_marker)
 74        .with_context(|| format!("missing end marker '{}' in prompt", end_marker.trim()))?
 75        + start;
 76
 77    let region = &prompt[start..end];
 78    let region = region.strip_suffix('\n').unwrap_or(region);
 79    Ok(region.replace(CURSOR_MARKER, ""))
 80}
 81
 82fn parse_zeta2_output(
 83    example: &Example,
 84    actual_output: &str,
 85    version: ZetaVersion,
 86) -> Result<(String, Option<usize>)> {
 87    let prompt = &example.prompt.as_ref().context("prompt required")?.input;
 88    let prompt_inputs = example
 89        .prompt_inputs
 90        .as_ref()
 91        .context("prompt_inputs required")?;
 92
 93    let old_text = extract_zeta2_current_region(prompt, version)?;
 94
 95    let mut new_text = actual_output.to_string();
 96    let cursor_offset = if let Some(offset) = new_text.find(CURSOR_MARKER) {
 97        new_text.replace_range(offset..offset + CURSOR_MARKER.len(), "");
 98        Some(offset)
 99    } else {
100        None
101    };
102
103    if version == ZetaVersion::V0120GitMergeMarkers {
104        if let Some(stripped) =
105            new_text.strip_suffix(zeta_prompt::v0120_git_merge_markers::END_MARKER)
106        {
107            new_text = stripped.to_string();
108        }
109    }
110
111    let mut old_text_normalized = old_text.clone();
112    if !new_text.is_empty() && !new_text.ends_with('\n') {
113        new_text.push('\n');
114    }
115    if !old_text_normalized.is_empty() && !old_text_normalized.ends_with('\n') {
116        old_text_normalized.push('\n');
117    }
118
119    let old_text_trimmed = old_text.trim_end_matches('\n');
120    let (editable_region_offset, _) = prompt_inputs
121        .content
122        .match_indices(old_text_trimmed)
123        .min_by_key(|(index, _)| index.abs_diff(prompt_inputs.cursor_offset))
124        .with_context(|| {
125            format!(
126                "could not find editable region in content.\nLooking for:\n{}\n\nIn content:\n{}",
127                old_text_trimmed, &prompt_inputs.content
128            )
129        })?;
130
131    let editable_region_start_line = prompt_inputs.content[..editable_region_offset]
132        .matches('\n')
133        .count();
134
135    // Use full context so cursor offset (relative to editable region start) aligns with diff content
136    let editable_region_lines = old_text_normalized.lines().count() as u32;
137    let diff = language::unified_diff_with_context(
138        &old_text_normalized,
139        &new_text,
140        editable_region_start_line as u32,
141        editable_region_start_line as u32,
142        editable_region_lines,
143    );
144
145    let formatted_diff = format!(
146        "--- a/{path}\n+++ b/{path}\n{diff}",
147        path = example.spec.cursor_path.to_string_lossy(),
148    );
149
150    Ok((formatted_diff, cursor_offset))
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn test_extract_zeta2_current_region_v0113() {
159        let prompt = indoc::indoc! {"
160            <|file_sep|>src/main.rs
161            <|fim_prefix|>
162            fn main() {
163            <|fim_middle|>current
164            println!(\"hello\");
165            <|fim_suffix|>
166            }
167            <|fim_middle|>updated
168        "};
169
170        let region = extract_zeta2_current_region(prompt, ZetaVersion::V0113Ordered).unwrap();
171        assert_eq!(region, "println!(\"hello\");");
172    }
173
174    #[test]
175    fn test_extract_zeta2_current_region_v0112() {
176        let prompt = indoc::indoc! {"
177            <|file_sep|>src/main.rs
178            <|fim_prefix|>
179            fn main() {
180            <|fim_suffix|>
181            }
182            <|fim_middle|>current
183            println!(\"hello\");
184            <|fim_middle|>updated
185        "};
186
187        let region = extract_zeta2_current_region(prompt, ZetaVersion::V0112MiddleAtEnd).unwrap();
188        assert_eq!(region, "println!(\"hello\");");
189    }
190
191    #[test]
192    fn test_extract_zeta2_current_region_with_cursor_marker() {
193        let prompt = indoc::indoc! {"
194            <|file_sep|>src/main.rs
195            <|fim_prefix|>
196            fn main() {
197            <|fim_middle|>current
198            print<|user_cursor|>ln!(\"hello\");
199            <|fim_suffix|>
200            }
201            <|fim_middle|>updated
202        "};
203
204        let region = extract_zeta2_current_region(prompt, ZetaVersion::V0113Ordered).unwrap();
205        assert_eq!(region, "println!(\"hello\");");
206    }
207
208    #[test]
209    fn test_extract_zeta2_current_region_v0120_git_merge_markers() {
210        let prompt = indoc::indoc! {"
211            <|file_sep|>src/main.rs
212            <|fim_prefix|>
213            fn main() {
214            <|fim_suffix|>
215            }
216            <|fim_middle|><<<<<<< CURRENT
217            println!(\"hello\");
218            =======
219        "};
220
221        let region =
222            extract_zeta2_current_region(prompt, ZetaVersion::V0120GitMergeMarkers).unwrap();
223        assert_eq!(region, "println!(\"hello\");");
224    }
225
226    #[test]
227    fn test_extract_zeta2_current_region_v0120_with_cursor_marker() {
228        let prompt = indoc::indoc! {"
229            <|file_sep|>src/main.rs
230            <|fim_prefix|>
231            fn main() {
232            <|fim_suffix|>
233            }
234            <|fim_middle|><<<<<<< CURRENT
235            print<|user_cursor|>ln!(\"hello\");
236            =======
237        "};
238
239        let region =
240            extract_zeta2_current_region(prompt, ZetaVersion::V0120GitMergeMarkers).unwrap();
241        assert_eq!(region, "println!(\"hello\");");
242    }
243}