//! Async operation types for Meerkat
//!
//! Unified abstraction for tool calls, shell commands, and delegated branches.

use crate::budget::BudgetLimits;
use crate::error::ToolError;
use crate::session::DeferredToolLoadAuthority;
use crate::types::{Message, ToolNameSet};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

/// Unique identifier for an operation
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct OperationId(pub Uuid);

/// Wait policy for async operations.
///
/// Determines whether an operation blocks the turn boundary (`Barrier`) or runs
/// independently (`Detached`).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WaitPolicy {
    /// Operation must complete before `ToolCallsResolved` can fire.
    Barrier,
    /// Operation runs independently and does not block the turn.
    Detached,
}

/// Typed async operation reference carrying an operation ID and its wait policy.
///
/// Replaces raw `OperationId` sequences in the TurnExecution machine state to
/// enable barrier-aware scheduling. Only `Barrier` ops block the turn boundary;
/// `Detached` ops are recorded but do not gate `ToolCallsResolved`.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct AsyncOpRef {
    pub operation_id: OperationId,
    pub wait_policy: WaitPolicy,
}

impl WaitPolicy {
    /// Normal tool-call operations that must complete before the turn boundary.
    pub fn barrier() -> Self {
        Self::Barrier
    }

    /// Background or mob-child operations that run independently of the turn.
    pub fn detached() -> Self {
        Self::Detached
    }
}

impl AsyncOpRef {
    /// Create a barrier op ref — blocks the turn boundary until resolved.
    pub fn barrier(operation_id: OperationId) -> Self {
        Self {
            operation_id,
            wait_policy: WaitPolicy::barrier(),
        }
    }

    /// Create a detached op ref — runs independently, does not block the turn.
    pub fn detached(operation_id: OperationId) -> Self {
        Self {
            operation_id,
            wait_policy: WaitPolicy::detached(),
        }
    }
}

/// Outcome of a tool dispatch, separating transcript data from execution metadata.
///
/// `result` is what the model sees (conversation/transcript). `async_ops` is
/// what the runtime scheduler sees (barrier/detached classification). This
/// prevents hooks, persistence, and message serialization from accidentally
/// owning barrier semantics.
/// Typed session-level effect produced by tool dispatch.
///
/// Tools that need to mutate session-owned durable state (e.g., mob authority)
/// must NOT call `SessionService` methods from inside dispatch. Instead they
/// return typed effects here, and the turn owner (agent loop) merges and commits
/// them after the parallel tool batch completes.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "effect_type", rename_all = "snake_case")]
pub enum SessionEffect {
    /// Grant management scope for a specific mob.
    GrantManageMob { mob_id: String },
    /// Record durable deferred-tool requests for subsequent boundaries.
    RequestDeferredTools {
        authorities: Vec<DeferredToolLoadAuthority>,
    },
    /// Append durable assistant blocks produced by a tool after the tool has
    /// committed any external payloads (for example generated image blobs).
    AppendAssistantBlocks {
        blocks: Vec<crate::types::AssistantBlock>,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolDispatchTerminalErrorKind {
    NotFound,
    Unavailable,
    InvalidArguments,
    ExecutionFailed,
    Timeout,
    AccessDenied,
    Other,
    CallbackPending,
}

impl From<&ToolError> for ToolDispatchTerminalErrorKind {
    fn from(error: &ToolError) -> Self {
        match error {
            ToolError::NotFound { .. } => Self::NotFound,
            ToolError::Unavailable { .. } => Self::Unavailable,
            ToolError::InvalidArguments { .. } => Self::InvalidArguments,
            ToolError::ExecutionFailed { .. } | ToolError::ExecutionFailedWithData { .. } => {
                Self::ExecutionFailed
            }
            ToolError::Timeout { .. } => Self::Timeout,
            ToolError::AccessDenied { .. } => Self::AccessDenied,
            ToolError::Other(_) => Self::Other,
            ToolError::CallbackPending { .. } => Self::CallbackPending,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolDispatchTerminalCause {
    RuntimeToolError { kind: ToolDispatchTerminalErrorKind },
}

impl ToolDispatchTerminalCause {
    #[must_use]
    pub fn runtime_tool_error(error: &ToolError) -> Self {
        Self::RuntimeToolError {
            kind: ToolDispatchTerminalErrorKind::from(error),
        }
    }

    #[must_use]
    pub fn is_runtime_tool_timeout(self) -> bool {
        matches!(
            self,
            Self::RuntimeToolError {
                kind: ToolDispatchTerminalErrorKind::Timeout
            }
        )
    }
}

#[derive(Debug, Clone)]
pub struct ToolDispatchOutcome {
    /// The tool result for the conversation transcript.
    pub result: crate::types::ToolResult,
    /// Async operations started by this dispatch, with typed wait policies.
    ///
    /// Empty for synchronous tools. Barrier ops block the turn boundary;
    /// detached ops run independently.
    pub async_ops: Vec<AsyncOpRef>,
    /// Session-level effects to be merged by the turn owner after the batch.
    ///
    /// Most tools return an empty vec. Tools that need durable session state
    /// changes (e.g., mob authority grants) emit typed effects here instead
    /// of calling `SessionService` from inside dispatch.
    pub session_effects: Vec<SessionEffect>,
    /// Runtime/tool-dispatch-authored terminal cause.
    ///
    /// Tool-authored `is_error` results intentionally leave this empty; callers
    /// must not infer canonical timeout/failure classes by parsing result text.
    terminal_cause: Option<ToolDispatchTerminalCause>,
}

/// Optional timeout policy supplied by an external tool-dispatch caller.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolDispatchTimeoutPolicy {
    /// Use the caller's default timeout value.
    Default { timeout: std::time::Duration },
    /// Do not apply a caller-specific timeout. The dispatcher may still apply
    /// its own normal execution policy.
    Disabled,
    /// Apply this finite caller-specific timeout.
    Finite { timeout: std::time::Duration },
}

impl ToolDispatchTimeoutPolicy {
    #[must_use]
    pub fn timeout(self) -> Option<std::time::Duration> {
        match self {
            Self::Default { timeout } | Self::Finite { timeout } => Some(timeout),
            Self::Disabled => None,
        }
    }

    #[must_use]
    pub fn timeout_ms(self) -> Option<u64> {
        self.timeout()
            .map(|timeout| u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX))
    }
}

impl ToolDispatchOutcome {
    /// Create an outcome with explicit async operations and session effects.
    pub fn new(
        result: crate::types::ToolResult,
        async_ops: Vec<AsyncOpRef>,
        session_effects: Vec<SessionEffect>,
    ) -> Self {
        Self {
            result,
            async_ops,
            session_effects,
            terminal_cause: None,
        }
    }

    /// Create an outcome with no async operations or session effects (synchronous tool).
    pub fn sync_result(result: crate::types::ToolResult) -> Self {
        Self::new(result, Vec::new(), Vec::new())
    }

    #[must_use]
    pub fn terminal_cause(&self) -> Option<ToolDispatchTerminalCause> {
        self.terminal_cause
    }

    #[must_use]
    pub fn is_runtime_tool_timeout(&self) -> bool {
        self.terminal_cause
            .is_some_and(ToolDispatchTerminalCause::is_runtime_tool_timeout)
    }

    pub(crate) fn clear_terminal_cause(&mut self) {
        self.terminal_cause = None;
    }
}

impl From<crate::types::ToolResult> for ToolDispatchOutcome {
    fn from(result: crate::types::ToolResult) -> Self {
        Self::sync_result(result)
    }
}

/// Convert a denied/failed tool dispatch into the canonical terminal tool
/// outcome shape used by both model-driven and external tool calls.
pub fn terminal_tool_outcome_for_error(
    tool_use_id: impl Into<String>,
    error: ToolError,
) -> ToolDispatchOutcome {
    let terminal_cause = ToolDispatchTerminalCause::runtime_tool_error(&error);
    let payload = error.to_error_payload();
    let serialized = serde_json::to_string(&payload)
        .unwrap_or_else(|_| "{\"error\":\"tool_error\",\"message\":\"tool error\"}".to_string());
    let mut outcome = ToolDispatchOutcome::sync_result(crate::types::ToolResult::new(
        tool_use_id.into(),
        serialized,
        true,
    ));
    outcome.terminal_cause = Some(terminal_cause);
    outcome
}

impl OperationId {
    /// Create a new operation ID
    pub fn new() -> Self {
        Self(crate::time_compat::new_uuid_v7())
    }
}

impl Default for OperationId {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Display for OperationId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// What kind of work the operation performs
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum WorkKind {
    /// MCP or internal tool call
    ToolCall,
    /// Shell command execution
    ShellCommand,
}

/// Shape of the operation's result
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ResultShape {
    /// Single result value
    Single,
    /// Streaming output (progress events)
    Stream,
    /// Multiple results (e.g., fork branches)
    Batch,
}

/// How much context a delegated branch receives
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(tag = "type", content = "value", rename_all = "snake_case")]
pub enum ContextStrategy {
    /// Complete conversation history (Fork default)
    #[default]
    FullHistory,
    /// Last N turns from parent
    LastTurns(u32),
    /// Compressed summary of conversation
    Summary { max_tokens: u32 },
    /// Explicit message list
    Custom { messages: Vec<Message> },
}

/// How to allocate budget when forking
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(tag = "type", content = "value", rename_all = "snake_case")]
pub enum ForkBudgetPolicy {
    /// Split remaining budget equally among branches
    #[default]
    Equal,
    /// Split proportionally based on weights
    Proportional,
    /// Fixed budget per branch
    Fixed(u64),
    /// Give all remaining budget to each branch
    Remaining,
}

/// Tool access control for delegated branches
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", content = "value", rename_all = "snake_case")]
pub enum ToolAccessPolicy {
    /// Inherit all tools from parent
    #[default]
    Inherit,
    /// Only allow specific tools
    AllowList(ToolNameSet),
    /// Block specific tools
    DenyList(ToolNameSet),
}

/// Policy for operation execution
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct OperationPolicy {
    /// Timeout for this operation
    pub timeout_ms: Option<u64>,
    /// Whether to cancel on parent cancellation
    pub cancel_on_parent_cancel: bool,
    /// Whether to include in checkpoints
    pub checkpoint_results: bool,
}

/// Complete operation specification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OperationSpec {
    pub id: OperationId,
    pub kind: WorkKind,
    pub result_shape: ResultShape,
    pub policy: OperationPolicy,
    pub budget_reservation: BudgetLimits,
    pub depth: u32,
    pub depends_on: Vec<OperationId>,
    pub context: Option<ContextStrategy>,
    pub tool_access: Option<ToolAccessPolicy>,
}

/// Result of a completed operation
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct OperationResult {
    pub id: OperationId,
    pub content: String,
    pub is_error: bool,
    pub duration_ms: u64,
    pub tokens_used: u64,
}

/// Events from operations
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum OpEvent {
    /// Operation started executing
    Started { id: OperationId, kind: WorkKind },

    /// Progress update (for streaming operations)
    Progress {
        id: OperationId,
        message: String,
        percent: Option<f32>,
    },

    /// Operation completed successfully
    Completed {
        id: OperationId,
        result: OperationResult,
    },

    /// Operation failed
    Failed { id: OperationId, error: String },

    /// Operation was cancelled
    Cancelled { id: OperationId },
}

/// Concurrency limits for operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConcurrencyLimits {
    /// Maximum delegated-branch nesting depth
    pub max_depth: u32,
    /// Maximum concurrent operations (all types)
    pub max_concurrent_ops: usize,
    /// Maximum concurrent delegated branches specifically
    pub max_concurrent_agents: usize,
    /// Maximum children per parent agent
    pub max_children_per_agent: usize,
}

impl Default for ConcurrencyLimits {
    fn default() -> Self {
        Self {
            max_depth: 3,
            max_concurrent_ops: 32,
            max_concurrent_agents: 8,
            max_children_per_agent: 5,
        }
    }
}

/// Specification for spawning a new delegated branch
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct SpawnSpec {
    /// The prompt/task for the delegated branch
    pub prompt: String,
    /// How much context the delegated branch receives
    pub context: ContextStrategy,
    /// Which tools the delegated branch can access
    pub tool_access: ToolAccessPolicy,
    /// Budget allocation for the delegated branch
    pub budget: BudgetLimits,
    /// If false, the delegated branch cannot spawn/fork further
    pub allow_spawn: bool,
    /// System prompt override (None = inherit from parent)
    pub system_prompt: Option<String>,
}

/// A branch in a fork operation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForkBranch {
    /// Identifier for this branch
    pub name: String,
    /// The prompt/task for this branch
    pub prompt: String,
    /// Tool access override (None = inherit)
    pub tool_access: Option<ToolAccessPolicy>,
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[test]
    fn barrier_constructor_produces_barrier_policy() {
        assert_eq!(WaitPolicy::barrier(), WaitPolicy::Barrier);
        let op_ref = AsyncOpRef::barrier(OperationId::new());
        assert_eq!(op_ref.wait_policy, WaitPolicy::Barrier);
    }

    #[test]
    fn detached_constructor_produces_detached_policy() {
        assert_eq!(WaitPolicy::detached(), WaitPolicy::Detached);
        let op_ref = AsyncOpRef::detached(OperationId::new());
        assert_eq!(op_ref.wait_policy, WaitPolicy::Detached);
    }

    #[test]
    fn test_operation_id_encoding() {
        let id = OperationId::new();
        let json = serde_json::to_string(&id).unwrap();

        let parsed: OperationId = serde_json::from_str(&json).unwrap();
        assert_eq!(id, parsed);
    }

    #[test]
    fn test_work_kind_serialization() {
        assert_eq!(
            serde_json::to_value(WorkKind::ToolCall).unwrap(),
            "tool_call"
        );
        assert_eq!(
            serde_json::to_value(WorkKind::ShellCommand).unwrap(),
            "shell_command"
        );
    }

    #[test]
    fn test_context_strategy_serialization() {
        let full = ContextStrategy::FullHistory;
        let json = serde_json::to_value(&full).unwrap();
        assert_eq!(json["type"], "full_history");

        let last = ContextStrategy::LastTurns(5);
        let json = serde_json::to_value(&last).unwrap();
        assert_eq!(json["type"], "last_turns");
        // Adjacently-tagged: {"type": "last_turns", "value": 5}
        assert_eq!(json["value"], 5);

        let summary = ContextStrategy::Summary { max_tokens: 1000 };
        let json = serde_json::to_value(&summary).unwrap();
        assert_eq!(json["type"], "summary");
        // Adjacently-tagged struct variant: {"type": "summary", "value": {"max_tokens": 1000}}
        assert_eq!(json["value"]["max_tokens"], 1000);

        // Roundtrip
        let parsed: ContextStrategy = serde_json::from_value(json).unwrap();
        match parsed {
            ContextStrategy::Summary { max_tokens } => assert_eq!(max_tokens, 1000),
            _ => unreachable!("Wrong variant"),
        }
    }

    #[test]
    fn test_fork_budget_policy_serialization() {
        let policies = vec![
            (ForkBudgetPolicy::Equal, "equal"),
            (ForkBudgetPolicy::Proportional, "proportional"),
            (ForkBudgetPolicy::Remaining, "remaining"),
        ];

        for (policy, expected_type) in policies {
            let json = serde_json::to_value(&policy).unwrap();
            assert_eq!(json["type"], expected_type);
        }

        let fixed = ForkBudgetPolicy::Fixed(5000);
        let json = serde_json::to_value(&fixed).unwrap();
        assert_eq!(json["type"], "fixed");
        // Adjacently-tagged: {"type": "fixed", "value": 5000}
        assert_eq!(json["value"], 5000);

        // Roundtrip
        let parsed: ForkBudgetPolicy = serde_json::from_value(json).unwrap();
        match parsed {
            ForkBudgetPolicy::Fixed(tokens) => assert_eq!(tokens, 5000),
            _ => unreachable!("Wrong variant"),
        }
    }

    #[test]
    fn test_tool_access_policy_serialization() {
        let inherit = ToolAccessPolicy::Inherit;
        let json = serde_json::to_value(&inherit).unwrap();
        assert_eq!(json["type"], "inherit");

        let allow = ToolAccessPolicy::AllowList(["read_file", "write_file"].into_iter().collect());
        let json = serde_json::to_value(&allow).unwrap();
        assert_eq!(json["type"], "allow_list");
        // Adjacently-tagged: {"type": "allow_list", "value": [...]}
        assert!(json["value"].is_array());

        let deny = ToolAccessPolicy::DenyList(["dangerous_tool"].into_iter().collect());
        let json = serde_json::to_value(&deny).unwrap();
        assert_eq!(json["type"], "deny_list");
        assert!(json["value"].is_array());

        // Roundtrip
        let parsed: ToolAccessPolicy = serde_json::from_value(json).unwrap();
        match parsed {
            ToolAccessPolicy::DenyList(tools) => {
                assert_eq!(tools.len(), 1);
                assert!(tools.contains("dangerous_tool"));
            }
            _ => unreachable!("Wrong variant"),
        }
    }

    #[test]
    fn test_op_event_serialization() {
        let events = vec![
            OpEvent::Started {
                id: OperationId::new(),
                kind: WorkKind::ToolCall,
            },
            OpEvent::Progress {
                id: OperationId::new(),
                message: "50% complete".to_string(),
                percent: Some(0.5),
            },
            OpEvent::Completed {
                id: OperationId::new(),
                result: OperationResult {
                    id: OperationId::new(),
                    content: "result".to_string(),
                    is_error: false,
                    duration_ms: 100,
                    tokens_used: 50,
                },
            },
            OpEvent::Failed {
                id: OperationId::new(),
                error: "timeout".to_string(),
            },
            OpEvent::Cancelled {
                id: OperationId::new(),
            },
        ];

        for event in events {
            let json = serde_json::to_value(&event).unwrap();
            assert!(json.get("type").is_some());

            // Roundtrip
            let _: OpEvent = serde_json::from_value(json).unwrap();
        }
    }

    #[test]
    fn test_concurrency_limits_default() {
        let limits = ConcurrencyLimits::default();
        assert_eq!(limits.max_depth, 3);
        assert_eq!(limits.max_concurrent_ops, 32);
        assert_eq!(limits.max_concurrent_agents, 8);
        assert_eq!(limits.max_children_per_agent, 5);
    }

    #[test]
    fn session_effect_grant_manage_mob_serde_round_trip() {
        let effect = SessionEffect::GrantManageMob {
            mob_id: "test-mob".into(),
        };
        let json = serde_json::to_value(&effect).unwrap();
        let parsed: SessionEffect = serde_json::from_value(json).unwrap();
        assert_eq!(effect, parsed);
    }

    #[test]
    fn tool_dispatch_outcome_with_session_effects() {
        let result = crate::types::ToolResult::new("t1".into(), "ok".into(), false);
        let outcome = ToolDispatchOutcome::new(
            result,
            vec![],
            vec![SessionEffect::GrantManageMob {
                mob_id: "mob-1".into(),
            }],
        );
        assert_eq!(outcome.session_effects.len(), 1);
        assert_eq!(outcome.terminal_cause(), None);
    }

    #[test]
    fn tool_dispatch_outcome_sync_result_has_empty_effects() {
        let result = crate::types::ToolResult::new("t1".into(), "ok".into(), false);
        let outcome = ToolDispatchOutcome::sync_result(result);
        assert!(outcome.session_effects.is_empty());
        assert_eq!(outcome.terminal_cause(), None);
    }

    #[test]
    fn terminal_tool_outcome_carries_runtime_timeout_cause() {
        let outcome = terminal_tool_outcome_for_error("t1", ToolError::timeout("slow_tool", 50));

        assert!(outcome.result.is_error);
        assert!(outcome.is_runtime_tool_timeout());
        assert_eq!(
            outcome.terminal_cause(),
            Some(ToolDispatchTerminalCause::RuntimeToolError {
                kind: ToolDispatchTerminalErrorKind::Timeout,
            })
        );
    }

    #[test]
    fn tool_authored_error_result_has_no_runtime_terminal_cause() {
        let result =
            crate::types::ToolResult::new("t1".into(), "{\"error\":\"timeout\"}".into(), true);
        let outcome = ToolDispatchOutcome::sync_result(result);

        assert!(outcome.result.is_error);
        assert!(!outcome.is_runtime_tool_timeout());
        assert_eq!(outcome.terminal_cause(), None);
    }
}
