oxo-flow-cli 0.21.2

//! AI-powered error recovery for run failures.

use anyhow::Result;
use chrono::Utc;
use colored::Colorize;
use oxo_flow_ai::{knowledge::builtin, provider::AiProvider};
use std::path::{Path, PathBuf};

/// Diagnostic result from AI error analysis.
pub struct DiagnoseResult {
    pub root_cause: String,
    pub fix_action: String,
    pub modified_toml: Option<String>,
    pub safe_to_auto_apply: bool,
}

/// Parse the model's safety verdict from the mandated output line
/// `- Safe to auto-apply: <yes/no>`.
///
/// The previous substring test (`contains("safe to auto-apply")`) matched
/// every conforming reply — including `- Safe to auto-apply: no` — so the
/// auto-apply gate in `run.rs` never blocked anything. Anything other than an
/// explicit `yes`/`true` verdict (missing, unparseable, or an explicit
/// `NOT safe to auto-apply` line) means: do not rewrite the user's workflow.
fn parse_safe_to_auto_apply(response_text: &str) -> bool {
    // An explicit refusal wins wherever it appears in the reply.
    if response_text
        .to_ascii_lowercase()
        .contains("not safe to auto-apply")
    {
        return false;
    }
    for line in response_text.lines() {
        // `to_ascii_lowercase` preserves byte length, so the index stays valid
        // on the original line.
        let lower = line.to_ascii_lowercase();
        let Some(idx) = lower.find("safe to auto-apply") else {
            continue;
        };
        let verdict = line[idx + "safe to auto-apply".len()..]
            .trim_start_matches(|c: char| c == ':' || c == '*' || c == '`' || c.is_whitespace())
            .trim_matches(|c: char| c == '*' || c == '_' || c == '`' || c.is_whitespace())
            .to_ascii_lowercase();
        if verdict.starts_with("yes") || verdict.starts_with("true") {
            return true;
        }
    }
    false
}

/// Analyze a pipeline failure using AI.
pub async fn diagnose_failure(
    workflow_path: &Path,
    failed_rule: &str,
    exit_code: i32,
    stderr: &str,
    provider: &AiProvider,
) -> Result<DiagnoseResult> {
    eprintln!();
    eprintln!("{}", "AI Error Recovery".bold().green());
    eprintln!("  Failed rule: {}", failed_rule.yellow());
    eprintln!("  Exit code: {}", exit_code.to_string().red());
    eprintln!(
        "  Model: {}\n",
        provider.model().unwrap_or_else(|| "default".into())
    );

    let system = build_diagnose_prompt();
    let toml_content =
        std::fs::read_to_string(workflow_path).unwrap_or_else(|_| "(unable to read)".into());

    let user = format!(
        "## Pipeline Failure\n\n\
         Rule: **{failed_rule}**\n\
         Exit code: {exit_code}\n\n\
         ## Error Output\n```\n{stderr}\n```\n\n\
         ## Current Workflow\n```toml\n{toml_content}\n```\n\n\
         ## Task\n\
         1. Match the error against known patterns\n\
         2. Identify the root cause\n\
         3. Propose a specific fix (change exact lines in the TOML)\n\
         4. Output the corrected TOML if changes are needed (inside ```toml fences)\n\
         5. State whether the fix is safe to auto-apply",
    );

    eprintln!("{}", "  Diagnosing...".bold().cyan());

    use oxo_flow_ai::types::Message;
    let messages = vec![Message::system(&system), Message::user(&user)];
    let response = provider.chat_with_tools(&messages, &[]).await?;
    let response_text = response.content.unwrap_or_default();

    // Record token usage
    tracing::info!(
        tokens_in = response.usage.prompt_tokens,
        tokens_out = response.usage.completion_tokens,
        "AI recovery analysis completed"
    );

    // Parse the AI response
    let root_cause = extract_section(&response_text, "Root Cause");
    let fix_action = extract_section(&response_text, "Fix");
    let safe = parse_safe_to_auto_apply(&response_text);

    // Extract modified TOML
    let modified_toml = extract_toml_block(&response_text);

    eprintln!(
        "\n{}\n{}",
        "Root Cause:".bold().red(),
        root_cause
            .as_deref()
            .unwrap_or("Unknown — see full analysis below")
    );
    eprintln!(
        "{}\n{}",
        "Suggested Fix:".bold().yellow(),
        fix_action
            .as_deref()
            .unwrap_or("Manual review needed — see full analysis below")
    );
    eprintln!(
        "Safe to auto-apply: {}",
        if safe { "yes".green() } else { "no".red() }
    );
    // When section parsing failed the model answered off-format; surface the
    // raw analysis instead of a dead-end pointer (the old "see full analysis"
    // wording pointed at output that was never printed).
    if root_cause.is_none() || fix_action.is_none() {
        eprintln!(
            "\n{}\n{}",
            "Full Analysis:".bold().underline(),
            response_text
        );
    }

    Ok(DiagnoseResult {
        root_cause: root_cause.unwrap_or_else(|| "Unknown — see full analysis".into()),
        fix_action: fix_action.unwrap_or_else(|| "Manual review needed".into()),
        modified_toml,
        safe_to_auto_apply: safe,
    })
}

/// Apply a fix to the workflow file, archiving the original.
pub fn apply_fix(workflow_path: &Path, modified_toml: &str, session_id: &str) -> Result<PathBuf> {
    // Archive original
    let archive_dir = archive_dir_for(workflow_path);
    std::fs::create_dir_all(&archive_dir)?;

    let original = std::fs::read_to_string(workflow_path)?;
    let timestamp = Utc::now().format("%Y%m%d-%H%M%S");
    let backup_path = archive_dir.join(format!("{timestamp}-{session_id}-before.oxoflow"));
    std::fs::write(&backup_path, &original)?;
    eprintln!(
        "{} Original archived to {}",
        "  ✓".green(),
        backup_path.display()
    );

    // Write fix
    std::fs::write(workflow_path, modified_toml)?;
    eprintln!(
        "{} Fix applied to {}",
        "  ✓".green(),
        workflow_path.display()
    );

    Ok(backup_path)
}

fn archive_dir_for(workflow_path: &Path) -> PathBuf {
    let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
    let local = cwd.join(".oxo-flow").join("ai_archive").join(
        workflow_path
            .file_stem()
            .unwrap_or_default()
            .to_string_lossy()
            .as_ref(),
    );
    if local.parent().is_some_and(|p| p.exists()) || cwd.join(".oxo-flow").exists() {
        local
    } else {
        let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
        PathBuf::from(home)
            .join(".oxo-flow")
            .join("ai_archive")
            .join(
                workflow_path
                    .file_stem()
                    .unwrap_or_default()
                    .to_string_lossy()
                    .as_ref(),
            )
    }
}

fn build_diagnose_prompt() -> String {
    let error_patterns = builtin::format_error_patterns();

    format!(
        r#"## Role & Identity
You are a senior bioinformatics pipeline incident responder for oxo-flow. When a pipeline rule fails,
you diagnose the root cause with surgical precision and propose the minimal fix needed.
Your analysis must be specific, actionable, and safe.

## Error Pattern Reference (Common Failures)
{error_patterns}

## Diagnostic Protocol (Execute in Order)
### Step 1 — Triage
- Classify the failure: resource (OOM/timeout), data (missing/corrupt input), software (version/API change), configuration (wrong params/paths)
- Map exit code to category: 137=OOM, 139=segfault, 1=user error, 127=command not found

### Step 2 — Root Cause Analysis
- Read the error output CAREFULLY — the answer is usually in the last 5 lines
- Cross-reference against error patterns above
- Check: is the tool version compatible with the provided parameters?
- Check: are input files actually produced by the dependency rule?
- Check: does the working directory have sufficient disk space?

### Step 3 — Fix Proposal
- Identify the EXACT line(s) in the TOML that need changing
- Propose the MINIMAL change — do not restructure the workflow unless essential
- If the fix is a resource adjustment (threads/memory), specify EXACT new values
- If the fix requires a tool parameter change, verify against the tool's documentation
- NEVER propose removing QC steps, validation, or safety checks

### Step 4 — Safety Assessment
- **Safe to auto-apply**: parameter tuning (threads, memory), fixing a typo, adding a missing flag
- **NOT safe to auto-apply**: adding/removing rules, changing DAG edges, modifying shell logic, changing file paths

## Output Format
Respond in this exact structure:

```
## Root Cause
<1-2 sentences explaining WHY the failure occurred, citing specific evidence from the error output>

## Proposed Fix
<Specific action. If changing TOML parameters, list old→new values. Reference exact lines.>

## Corrected TOML
```toml
<complete corrected [workflow] section or [[rules]] block — only the changed portions>
```

## Safety Assessment
- Safe to auto-apply: <yes/no>
- Risk level: <low/medium/high>
- Rollback: <how to undo this change>

## Prevention
<1 sentence on how to prevent this failure in future pipelines>
```

⚠ ONLY propose changes that fix the reported error. Do NOT make unrelated improvements.
"#
    )
}

/// Extract TOML code block from AI response.
///
/// The diagnostic prompt allows "only the changed portions" — a `[[rules]]`
/// block is a valid fix on its own, so accept it alongside `[workflow]`.
fn extract_toml_block(response: &str) -> Option<String> {
    if let Some(start) = response.find("```toml") {
        let start = start + 7;
        if let Some(end) = response[start..].find("```") {
            let content = response[start..start + end].trim().to_string();
            if content.contains("[workflow]") || content.contains("[[rules]") {
                return Some(content);
            }
        }
    }
    None
}

/// Extract a named section from the AI response.
///
/// Tries the inline patterns first (`Marker:`, `**Marker**:`, ...), then the
/// markdown headings the diagnostic prompt itself mandates (`## Root Cause`,
/// `## Proposed Fix`). The heading form is what conforming replies actually
/// use — without it every well-formed answer fell through to the fallback
/// text and the real analysis was dropped.
fn extract_section(text: &str, marker: &str) -> Option<String> {
    let patterns = [
        format!("[{marker}]:"),
        format!("[{marker}]"),
        format!("**{marker}**:"),
        format!("{marker}:"),
    ];
    for pat in &patterns {
        if let Some(pos) = text.find(pat.as_str()) {
            let remainder = &text[pos + pat.len()..];
            let end = remainder
                .find("\n[")
                .or_else(|| remainder.find("\n**"))
                .or_else(|| remainder.find("\n```"))
                .unwrap_or(remainder.len());
            let found = remainder[..end].trim();
            if !found.is_empty() {
                return Some(found.to_string());
            }
        }
    }
    extract_heading_section(text, marker)
}

/// Capture the body under a markdown heading whose title contains `marker`
/// (case-insensitive), ending at the next heading of any level or a code fence.
fn extract_heading_section(text: &str, marker: &str) -> Option<String> {
    let marker_lower = marker.to_ascii_lowercase();
    for (idx, line) in text.lines().enumerate() {
        let trimmed = line.trim();
        if !trimmed.starts_with('#') {
            continue;
        }
        let title = trimmed.trim_start_matches('#').trim().to_ascii_lowercase();
        if !title.contains(&marker_lower) {
            continue;
        }
        let body: String = text
            .lines()
            .skip(idx + 1)
            .take_while(|l| {
                let t = l.trim();
                !(t.starts_with('#') || t.starts_with("```"))
            })
            .collect::<Vec<_>>()
            .join("\n");
        let body = body.trim();
        if !body.is_empty() {
            return Some(body.to_string());
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::{extract_section, extract_toml_block, parse_safe_to_auto_apply};

    #[test]
    fn safe_gate_requires_an_explicit_yes() {
        assert!(parse_safe_to_auto_apply("- Safe to auto-apply: yes"));
        assert!(parse_safe_to_auto_apply("- **Safe to auto-apply**: YES"));
        assert!(parse_safe_to_auto_apply("Safe to auto-apply: true"));
        // The regression: the old substring test matched this exact line.
        assert!(!parse_safe_to_auto_apply("- Safe to auto-apply: no"));
        assert!(!parse_safe_to_auto_apply("- Safe to auto-apply: <yes/no>"));
        assert!(!parse_safe_to_auto_apply(
            "- **NOT safe to auto-apply**: changing DAG edges"
        ));
        // Missing/unparseable verdicts must never authorize a rewrite.
        assert!(!parse_safe_to_auto_apply(
            "## Root Cause\nbwa ran out of memory"
        ));
        assert!(!parse_safe_to_auto_apply(""));
    }

    /// A fully-conforming reply uses the heading format the prompt mandates.
    #[test]
    fn section_extraction_matches_prompt_mandated_headings() {
        let reply = "## Root Cause\nThe rule requested 16 threads but the \
machine exposes 4; the scheduler killed it (exit 137).\n\n## Proposed Fix\n\
Change `threads = 16` to `threads = 4` in rule `align`.\n\n## Corrected TOML\n\
```toml\n[[rules]]\nname = \"align\"\nthreads = 4\n```\n\n## Safety Assessment\n\
- Safe to auto-apply: yes\n- Risk level: low\n";

        let cause = extract_section(reply, "Root Cause").expect("root cause parsed");
        assert!(cause.contains("exit 137"), "got: {cause}");
        let fix = extract_section(reply, "Fix").expect("fix parsed");
        assert!(fix.contains("threads = 4"), "got: {fix}");
    }

    #[test]
    fn section_extraction_stops_at_next_heading_or_fence() {
        let reply = "## Root Cause\nbad reference genome\n\n## Proposed Fix\npoint to hg38";
        assert_eq!(
            extract_section(reply, "Root Cause").as_deref(),
            Some("bad reference genome")
        );
        assert_eq!(
            extract_section(reply, "Fix").as_deref(),
            Some("point to hg38")
        );
    }

    #[test]
    fn section_extraction_still_supports_inline_markers() {
        assert_eq!(
            extract_section("**Root Cause**: typo in path", "Root Cause").as_deref(),
            Some("typo in path")
        );
        assert_eq!(
            extract_section("Root Cause: typo in path", "Root Cause").as_deref(),
            Some("typo in path")
        );
    }

    /// The prompt permits "only the changed portions" — a [[rules]]-only
    /// block is a legitimate fix and must not be discarded.
    #[test]
    fn toml_block_accepts_rules_only_fragments() {
        let fragment = "```toml\n[[rules]]\nname = \"align\"\nthreads = 4\n```";
        assert!(extract_toml_block(fragment).is_some());
        assert!(extract_toml_block("```toml\n[workflow]\nname = \"x\"\n```").is_some());
        assert!(extract_toml_block("```toml\nnot toml at all\n```").is_none());
        assert!(extract_toml_block("no fence here").is_none());
    }
}