//! Harness middleware ported from the AHE evolved harness:
//! - round & token reminders injected before each model call
//! - execution risk hints appended to shell tool results
//! - context compaction trigger logic

use regex::Regex;
use std::collections::BTreeSet;
use std::sync::OnceLock;

macro_rules! re {
    ($name:ident, $pat:expr) => {
        fn $name() -> &'static Regex {
            static RE: OnceLock<Regex> = OnceLock::new();
            RE.get_or_init(|| Regex::new($pat).expect("valid middleware regex"))
        }
    };
}

re!(localhost_re, r"(?i)https?://(127\.0\.0\.1|localhost|0\.0\.0\.0)(:\d+)?/");
re!(validation_desc_re, r"(?i)\b(validate|validation|verify|verification|acceptance|final|evaluator(-style)?|check|smoke test|layout)\b");
re!(help_or_syntax_re, r"(?i)\b(--help|-h\b|py_compile|compileall)\b");
re!(existence_only_re, r"(?i)^\s*(pwd|ls|stat|file|which|type|readlink|test\s+-[efsd]|wc\s+-[clmw]|head|tail|sed\s+-n|grep|find)\b");
re!(missing_dep_re, r"(?i)(command not found|No module named|ModuleNotFoundError|not installed|not found)");
re!(env_probe_re, r"(?i)\b(command\s+-v|which\b|type\b|python\d*\s+-c|python\d*\s+--version|pip\b|uv\b|apt-get\b|brew\b)");
re!(success_signal_re, r"(?i)(VALIDATION_OK|acceptance check passed|all (diff )?checks passed|exact_match\s*[:=]\s*True|diff_exit=0)");
re!(benchmark_re, r"(?i)\b(benchmark|speedup|median_s|elapsed_ms|latency|throughput|runtime)\b");
re!(reference_re, r"(?i)\b(golden|baseline|reference|threshold|ratio|target)\b");

pub const LONG_STEP_MS: u128 = 90_000;
pub const BUDGET_WARN_TOTAL_MS: u128 = 420_000;

/// Per-session middleware state (risk-hint counters, emitted-note dedup).
#[derive(Debug, Default)]
pub struct HintState {
    pub long_steps: u32,
    pub long_ms: u128,
    pub timeouts: u32,
    pub same_error_repeats: u32,
    pub last_error_signature: String,
    pub dependency_probe_failures: u32,
    emitted: BTreeSet<&'static str>,
}

impl HintState {
    fn emit_once(&mut self, key: &'static str, note: String, out: &mut Vec<String>) {
        if self.emitted.insert(key) {
            out.push(note);
        }
    }

    /// Inspect a finished shell command; returns up to 2 "Execution note:" hints
    /// to append to the tool result.
    pub fn after_shell(
        &mut self,
        command: &str,
        description: &str,
        output: &str,
        exit_code: i32,
        duration_ms: u128,
        timed_out: bool,
    ) -> Vec<String> {
        let mut notes: Vec<String> = Vec::new();

        if duration_ms >= LONG_STEP_MS || timed_out {
            self.long_steps += 1;
            self.long_ms += duration_ms;
        }
        if timed_out {
            self.timeouts += 1;
        }

        // Repeated identical error family.
        let signature = error_signature(output, exit_code, timed_out);
        if !signature.is_empty() {
            if signature == self.last_error_signature {
                self.same_error_repeats += 1;
            } else {
                self.last_error_signature = signature;
                self.same_error_repeats = 1;
            }
        } else {
            self.last_error_signature.clear();
            self.same_error_repeats = 0;
        }
        if self.same_error_repeats >= 2 {
            self.emit_once("retry", "The same error family just repeated. Stop retrying the same command pattern; inspect the contract/path/tool signature once, then switch strategy.".into(), &mut notes);
        }

        // Missing dependency probe loop.
        let combined = format!("{description}\n{command}");
        if (env_probe_re().is_match(&combined)
            || existence_only_re().is_match(command)
            || help_or_syntax_re().is_match(command))
            && missing_dep_re().is_match(output)
            && exit_code != 0
        {
            self.dependency_probe_failures += 1;
        } else {
            self.dependency_probe_failures = 0;
        }
        if self.dependency_probe_failures >= 2 {
            self.emit_once("missing_dependency", "A key runtime/dependency probe failed again. Stop re-checking the environment; move to the best direct implementation or the cheapest bounded fallback.".into(), &mut notes);
        }

        // Time budget.
        if self.long_steps >= 2 && (self.timeouts >= 1 || self.long_ms >= BUDGET_WARN_TOTAL_MS) {
            let secs = (self.long_ms / 1000) as u64;
            self.emit_once("budget", format!("Budget warning: about {secs}s spent across long foreground runs/timeouts. Stop starting new heavy searches; use background jobs with short polls or finalize the best verified candidate now."), &mut notes);
        }

        if duration_ms >= LONG_STEP_MS && !timed_out {
            self.emit_once("long_step", format!("This step consumed about {}s. Reassess now: keep the cheapest viable path, and if a candidate already meets the contract, save it to the required target path before more exploration.", duration_ms / 1000), &mut notes);
        }

        // Localhost-only validation.
        if localhost_re().is_match(command) && validation_desc_re().is_match(&combined) {
            self.emit_once("localhost", "A localhost/127.0.0.1 check only proves local reachability. Mirror the evaluator's exact public host/path/port and final deployed state before publishing.".into(), &mut notes);
        }

        // Shallow validation (existence/help only).
        if validation_desc_re().is_match(description)
            && (existence_only_re().is_match(command) || help_or_syntax_re().is_match(command))
        {
            self.emit_once("shallow", "This was only help/syntax/file-existence validation. Run the real entry point with evaluator-style args and confirm the exact output path/public endpoint from disk before treating it as proof.".into(), &mut notes);
        }

        // Acceptance-style success: freeze.
        if exit_code == 0 && success_signal_re().is_match(output) {
            self.emit_once("success_freeze", "This looks like acceptance-style success. Do not edit the deliverable again unless you have new failing evidence; if you do, rerun the exact final check from the canonical entry point.".into(), &mut notes);
        }

        // Benchmarks without a reference comparator.
        let all = format!("{combined}\n{output}");
        if benchmark_re().is_match(&all) && !reference_re().is_match(&all) {
            self.emit_once("benchmark_reference", "A speedup against an arbitrary comparator is not enough. Benchmark against the task's golden/reference baseline under the same setup and decide by the required median/threshold.".into(), &mut notes);
        }

        notes.truncate(2);
        notes
    }
}

fn error_signature(output: &str, exit_code: i32, timed_out: bool) -> String {
    if exit_code == 0 && !timed_out {
        return String::new();
    }
    for line in output.lines() {
        let t = line.trim();
        if t.starts_with("Error")
            || t.starts_with("Timeout")
            || t.starts_with("Traceback")
            || t.contains("AssertionError")
            || t.contains("FileNotFoundError")
            || t.contains("ModuleNotFoundError")
            || t.starts_with("error[")
            || t.starts_with("error:")
        {
            return t.chars().take(180).collect();
        }
    }
    format!("exit:{exit_code}")
}

/// Round/token reminder injected before each model call (after the first
/// assistant turn exists).
pub fn round_token_reminder(
    iteration: u32,
    max_iterations: u32,
    prompt_tokens: u32,
    max_context_tokens: u32,
) -> String {
    let remaining_iters = max_iterations.saturating_sub(iteration);
    let iter_hint = if remaining_iters <= 1 {
        format!("WARNING: iteration {iteration}/{max_iterations} — only {remaining_iters} left. Provide a conclusive response now; avoid further tool calls unless absolutely critical.")
    } else if remaining_iters <= 3 {
        format!("Iteration {iteration}/{max_iterations} — {remaining_iters} remaining. Work towards a conclusion.")
    } else {
        format!("Iteration {iteration}/{max_iterations}.")
    };

    let remaining_tokens = max_context_tokens.saturating_sub(prompt_tokens);
    // Warn proportionally (<25% of the window left) so small local models
    // (16-40K) don't trip a fixed absolute threshold on every single turn while
    // 1M-context models still get an early heads-up.
    let token_hint = if max_context_tokens > 0 && remaining_tokens < max_context_tokens / 4 {
        format!("WARNING: context usage {prompt_tokens}/{max_context_tokens} — only {remaining_tokens} tokens left. Be brief; finalize soon.")
    } else {
        format!("Context: {prompt_tokens}/{max_context_tokens} tokens used.")
    };

    format!("[harness] {iter_hint} {token_hint}")
}

/// Returns true when history should be compacted (>=75% of context used).
pub fn should_compact(prompt_tokens: u32, max_context_tokens: u32) -> bool {
    max_context_tokens > 0 && prompt_tokens as u64 * 100 >= max_context_tokens as u64 * 75
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn repeated_error_hint_fires_once() {
        let mut s = HintState::default();
        let n1 = s.after_shell("python x.py", "", "ModuleNotFoundError: No module named 'x'", 1, 100, false);
        assert!(n1.is_empty());
        let n2 = s.after_shell("python x.py", "", "ModuleNotFoundError: No module named 'x'", 1, 100, false);
        assert!(n2.iter().any(|n| n.contains("same error family")));
        let n3 = s.after_shell("python x.py", "", "ModuleNotFoundError: No module named 'x'", 1, 100, false);
        assert!(!n3.iter().any(|n| n.contains("same error family")));
    }

    #[test]
    fn shallow_validation_hint() {
        let mut s = HintState::default();
        let notes = s.after_shell("ls -la out/", "final validation of layout", "result.txt", 0, 50, false);
        assert!(notes.iter().any(|n| n.contains("file-existence validation")));
    }

    #[test]
    fn compaction_threshold() {
        assert!(!should_compact(100, 1000));
        assert!(should_compact(750, 1000));
        assert!(should_compact(900, 1000));
    }

    #[test]
    fn compaction_tracks_per_model_window() {
        // Same prompt size: safe on a 1M model, must compact after a mid-session
        // downgrade to a 200K model (the per-model context bug this fixes).
        let prompt = 300_000;
        assert!(!should_compact(prompt, 1_048_576));
        assert!(should_compact(prompt, 200_000));
        // A tiny local window (16K) compacts very early.
        assert!(should_compact(13_000, 16_384));
    }

    #[test]
    fn token_warning_is_proportional() {
        // 16K local model at 13K used (<25% left) warns...
        assert!(round_token_reminder(2, 150, 13_000, 16_384).contains("WARNING"));
        // ...but at 4K used (lots of room) it does not.
        assert!(!round_token_reminder(2, 150, 4_000, 16_384).contains("WARNING"));
        // 1M model with 300K used still has >25% free → no token warning.
        assert!(!round_token_reminder(2, 150, 300_000, 1_048_576).contains("WARNING"));
    }

    #[test]
    fn reminder_warns_near_end() {
        let r = round_token_reminder(149, 150, 1000, 100_000);
        assert!(r.contains("WARNING"));
        let r2 = round_token_reminder(3, 150, 1000, 100_000);
        assert!(!r2.contains("WARNING"));
    }
}
