use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};

use tokio_util::sync::CancellationToken;

#[derive(Debug, Clone)]
pub struct AgentBudget {
    pub remaining_runs: usize,
    pub remaining_depth: u16,
}

/// A signal this run emitted, preserved so the terminal completion payload
/// can reference the signal IDs for cross-checking.
#[derive(Debug, Clone)]
pub struct EmittedSignal {
    pub signal_id: String,
}

#[derive(Debug, Clone)]
pub struct AgentExecutionContext {
    pub root_session_id: String,
    pub root_turn_id: Option<String>,
    pub run_id: String,
    /// Execution attempt identifier owning conditional state transitions in
    /// Storage.  Equal to `run_id` for the first attempt.
    pub execution_id: String,
    pub parent_run_id: Option<String>,
    pub caller_agent_id: String,
    pub current_agent_id: String,
    /// Agent IDs already present in this execution chain, including current.
    pub ancestry: Vec<String>,
    pub depth: u16,
    pub plan_item_id: Option<String>,
    pub cancellation: CancellationToken,
    pub budget: AgentBudget,
    /// Shared admission counter across the whole delegation tree, including
    /// the root run that owns this context.  Durable run accounting replaces
    /// it once the Coordinator persists runs; until then it is the only
    /// tree-wide enforcement of `max_runs_per_tree`.
    pub tree_runs: Arc<AtomicUsize>,
    /// Durable emit_signal contract; `None` means this run has no signal
    /// capability and must not see the `emit_signal` tool.
    pub signal_contract: Option<Arc<crate::agent::definition::SignalContract>>,
    /// Signals accepted by this run so far, in emit order.
    pub emitted_signals: Arc<Mutex<Vec<EmittedSignal>>>,
}

impl AgentExecutionContext {
    pub fn child(
        parent: &Arc<Self>,
        run_id: String,
        target: String,
        plan_item_id: Option<String>,
        cancellation: CancellationToken,
    ) -> Self {
        let mut ancestry = parent.ancestry.clone();
        ancestry.push(target.clone());
        Self {
            root_session_id: parent.root_session_id.clone(),
            root_turn_id: parent.root_turn_id.clone(),
            run_id: run_id.clone(),
            execution_id: run_id,
            parent_run_id: Some(parent.run_id.clone()),
            caller_agent_id: parent.current_agent_id.clone(),
            current_agent_id: target,
            ancestry,
            depth: parent.depth.saturating_add(1),
            plan_item_id,
            cancellation,
            budget: AgentBudget {
                remaining_runs: parent.budget.remaining_runs.saturating_sub(1),
                remaining_depth: parent.budget.remaining_depth.saturating_sub(1),
            },
            tree_runs: parent.tree_runs.clone(),
            signal_contract: parent.signal_contract.clone(),
            emitted_signals: Arc::new(Mutex::new(Vec::new())),
        }
    }

    /// Reserve one tree-run slot for a child that has already passed the
    /// caller's budget checks.  Returns the total number of runs in the tree
    /// after reservation, or `None` (with the counter restored) when the
    /// reservation would exceed `max_runs_per_tree`.
    pub fn reserve_tree_run(&self, max_runs_per_tree: usize) -> Option<usize> {
        let previous = self.tree_runs.fetch_add(1, Ordering::SeqCst);
        let total = previous.saturating_add(1);
        if total > max_runs_per_tree {
            self.tree_runs.fetch_sub(1, Ordering::SeqCst);
            return None;
        }
        Some(total)
    }

    /// Remaining tree capacity for batch admission checks.
    pub fn remaining_tree_runs(&self, max_runs_per_tree: usize) -> usize {
        max_runs_per_tree.saturating_sub(self.tree_runs.load(Ordering::SeqCst))
    }
}

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

    fn root_context() -> AgentExecutionContext {
        AgentExecutionContext {
            root_session_id: "cli:test:dialog".to_string(),
            root_turn_id: None,
            run_id: "run-root".to_string(),
            execution_id: "run-root".to_string(),
            parent_run_id: None,
            caller_agent_id: "ROOT".to_string(),
            current_agent_id: "researcher".to_string(),
            ancestry: vec!["researcher".to_string()],
            depth: 1,
            plan_item_id: None,
            cancellation: CancellationToken::new(),
            budget: AgentBudget {
                remaining_runs: 15,
                remaining_depth: 3,
            },
            tree_runs: Arc::new(AtomicUsize::new(1)),
            signal_contract: None,
            emitted_signals: Arc::new(Mutex::new(Vec::new())),
        }
    }

    #[test]
    fn tree_run_reservation_is_shared_and_bounded() {
        let parent = Arc::new(root_context());
        assert_eq!(parent.remaining_tree_runs(3), 2);

        assert_eq!(parent.reserve_tree_run(3), Some(2));
        assert_eq!(parent.reserve_tree_run(3), Some(3));
        assert_eq!(parent.reserve_tree_run(3), None);
        assert_eq!(parent.tree_runs.load(Ordering::SeqCst), 3);
        assert_eq!(parent.remaining_tree_runs(3), 0);

        let child = AgentExecutionContext::child(
            &parent,
            "run-child".to_string(),
            "reviewer".to_string(),
            None,
            CancellationToken::new(),
        );
        assert!(Arc::ptr_eq(&parent.tree_runs, &child.tree_runs));
        assert_eq!(child.remaining_tree_runs(3), 0);
    }
}
