//! Wall-clock and phase-bucketed LLM wait timing for agent runs.
//! JSON is always written to [`RUN_TIMING_JSON_FILE`]; `code`/`kpop` also print [`RUN_TIMING_SUMMARY_PREFIX`].

mod report;

use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

pub const RUN_TIMING_JSON_FILE: &str = "run_timing.json";

pub const RUN_TIMING_SUMMARY_PREFIX: &str = "TIMING: ";

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TimingPhase {
    Implement,
}

#[derive(Debug, Clone)]
pub struct RunTiming {
    wall_start: Option<Instant>,
    wall_end: Option<Instant>,
    llm_wait: Duration,
    agent_retry_backoff: Duration,
    implement: Duration,
    implement_display_name: &'static str,
    tool_calls: Duration,
}

impl Default for RunTiming {
    fn default() -> Self {
        Self {
            wall_start: None,
            wall_end: None,
            llm_wait: Duration::ZERO,
            agent_retry_backoff: Duration::ZERO,
            implement: Duration::ZERO,
            implement_display_name: "implement",
            tool_calls: Duration::ZERO,
        }
    }
}

impl RunTiming {
    pub const fn add_tool_call_wall(&mut self, d: Duration) {
        self.tool_calls = self.tool_calls.saturating_add(d);
    }
}

impl RunTiming {
    #[must_use]
    pub fn new_arc() -> Arc<Mutex<Self>> {
        Arc::new(Mutex::new(Self::default()))
    }

    pub const fn mark_wall_start(&mut self, at: Instant) {
        self.wall_start = Some(at);
    }

    pub const fn mark_wall_end(&mut self, at: Instant) {
        self.wall_end = Some(at);
    }

    pub const fn add_llm_phase(&mut self, phase: TimingPhase, d: Duration) {
        let TimingPhase::Implement = phase;
        self.llm_wait = self.llm_wait.saturating_add(d);
        self.implement = self.implement.saturating_add(d);
    }

    pub const fn add_agent_retry_backoff(&mut self, d: Duration) {
        self.agent_retry_backoff = self.agent_retry_backoff.saturating_add(d);
    }

    pub const fn set_implement_display_name(&mut self, label: &'static str) {
        self.implement_display_name = label;
    }

    pub(crate) fn wall_duration(&self) -> Option<Duration> {
        match (self.wall_start, self.wall_end) {
            (Some(a), Some(b)) => Some(b.saturating_duration_since(a)),
            _ => None,
        }
    }

    #[must_use]
    pub fn elapsed_so_far(&self) -> Duration {
        self.wall_start.map_or(Duration::ZERO, |start| {
            Instant::now().saturating_duration_since(start)
        })
    }

    /// Writes timing JSON and prints the human-readable summary line.
    ///
    /// # Errors
    ///
    /// Returns [`std::io::Error`] when writing under `run_dir` fails.
    pub fn write_json_and_print_summary(&self, run_dir: &Path) -> std::io::Result<()> {
        report::write_json_and_print_summary(self, run_dir)
    }

    /// Writes timing JSON without printing a summary line.
    ///
    /// # Errors
    ///
    /// Returns [`std::io::Error`] when writing under `run_dir` fails.
    pub fn write_json_only(&self, run_dir: &Path) -> std::io::Result<()> {
        report::write_json_only(self, run_dir)
    }
}

#[must_use]
pub fn attach_new_run_timing(
    timing_slot: &mut Option<Arc<Mutex<RunTiming>>>,
) -> Arc<Mutex<RunTiming>> {
    let timing = RunTiming::new_arc();
    *timing_slot = Some(Arc::clone(&timing));
    timing
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .mark_wall_start(Instant::now());
    timing
}

/// Anchors one wall-clock interval for a gate-kpop `code` loop (shared across iterations).
#[must_use]
pub fn attach_gate_kpop_loop_run_timing() -> Arc<Mutex<RunTiming>> {
    let mut slot = None;
    attach_new_run_timing(&mut slot)
}

pub fn record_llm(timing: Option<&Arc<Mutex<RunTiming>>>, phase: TimingPhase, elapsed: Duration) {
    let Some(t) = timing else {
        return;
    };
    let mut g = t.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
    g.add_llm_phase(phase, elapsed);
}

pub fn record_backoff(timing: Option<&Arc<Mutex<RunTiming>>>, d: Duration) {
    let Some(t) = timing else {
        return;
    };
    let mut g = t.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
    g.add_agent_retry_backoff(d);
}

fn finalize_snapshot(timing: &Arc<Mutex<RunTiming>>) -> RunTiming {
    let mut g = timing
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    if g.wall_end.is_none() {
        g.mark_wall_end(Instant::now());
    }
    g.clone()
}

/// Finalizes wall clock end time and writes JSON plus the printed summary.
///
/// # Errors
///
/// Returns [`std::io::Error`] when writing under `run_dir` fails.
pub fn finalize_and_emit_run_timing(
    run_dir: &Path,
    timing: &Arc<Mutex<RunTiming>>,
) -> std::io::Result<()> {
    finalize_snapshot(timing).write_json_and_print_summary(run_dir)
}

/// Finalizes wall clock end time and writes JSON only.
///
/// # Errors
///
/// Returns [`std::io::Error`] when writing under `run_dir` fails.
pub fn finalize_run_timing_json_only(
    run_dir: &Path,
    timing: &Arc<Mutex<RunTiming>>,
) -> std::io::Result<()> {
    finalize_snapshot(timing).write_json_only(run_dir)
}

/// Persists in-progress timing without closing the run wall clock.
///
/// # Errors
///
/// Returns [`std::io::Error`] when writing under `run_dir` fails.
pub fn persist_open_run_timing_json(
    run_dir: &Path,
    timing: &Arc<Mutex<RunTiming>>,
) -> std::io::Result<()> {
    let g = timing
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    report::write_json_only(&g, run_dir)
}

pub use report::print_summary_from_run_dir;

#[cfg(test)]
mod timing_tests;

#[cfg(test)]
mod gate_kpop_timing_regressions;

pub mod acp_post_run;
