use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};
use std::time::{Duration, Instant};

use roto::Context;

use crate::error::{Error, Result};
use crate::request::RunLimits;

#[derive(Debug)]
struct CancellationState {
    cancelled: AtomicBool,
    deadline: Option<Instant>,
}

/// Cooperative cancellation state for one runtime execution.
///
/// The engine checks this token at SDK-controlled boundaries. Embedders can
/// also pass a token to [`Agent::run_once_with_cancellation`](crate::Agent::run_once_with_cancellation)
/// when they need explicit cancellation.
#[derive(Clone, Debug)]
pub struct RunCancellation {
    state: Arc<CancellationState>,
}

impl RunCancellation {
    pub fn new(deadline: Option<Instant>) -> Self {
        Self {
            state: Arc::new(CancellationState {
                cancelled: AtomicBool::new(false),
                deadline,
            }),
        }
    }

    pub fn none() -> Self {
        Self::new(None)
    }

    pub fn with_timeout(timeout_s: u32) -> Self {
        let deadline = if timeout_s == 0 {
            None
        } else {
            Some(Instant::now() + Duration::from_secs(timeout_s as u64))
        };
        Self::new(deadline)
    }

    pub fn cancel(&self) {
        self.state.cancelled.store(true, Ordering::Release);
    }

    pub fn deadline(&self) -> Option<Instant> {
        self.state.deadline
    }

    pub fn is_cancelled(&self) -> bool {
        self.state.cancelled.load(Ordering::Acquire)
            || self
                .state
                .deadline
                .is_some_and(|deadline| Instant::now() >= deadline)
    }

    pub fn check_cancelled(&self) -> Result<()> {
        if self.state.cancelled.load(Ordering::Acquire) {
            return Err(Error::cancelled("execution cancelled")
                .with_primary_diagnostic("execution-cancelled"));
        }

        if self
            .state
            .deadline
            .is_some_and(|deadline| Instant::now() >= deadline)
        {
            return Err(
                Error::cancelled("timeout exceeded").with_primary_diagnostic("timeout-exceeded")
            );
        }

        Ok(())
    }
}

/// The execution context available to Roto scripts and registered host
/// functions during a run.
///
/// Public fields are visible to Roto scripts. Module authors should treat
/// these fields as **read-only**.
///
/// The runtime constructs a fresh context for every run.  Callers should not
/// mutate the public fields; doing so leads to undefined behaviour in the Roto
/// runtime.
#[derive(Clone)]
pub struct ExecutionContext {
    /// Maximum time the script may run, in seconds.
    ///
    /// `0` means no timeout. The SDK enforces this cooperatively at runtime
    /// checkpoints.
    pub timeout_s: u32,

    /// Maximum number of bytes the script may write to output.
    pub max_output_bytes: u64,

    /// Maximum number of bytes the script may use in temporary space.
    pub max_temp_bytes: u64,

    /// Per-run metadata as a string.
    ///
    /// This is opaque to the runtime; callers may use it for correlation IDs,
    /// trace tokens, or any other run-scoped data.
    ///
    /// An empty string means metadata was not set by the caller.
    pub metadata: Arc<str>,

    cancellation: RunCancellation,
}

// SAFETY: This manual implementation intentionally exposes only the public,
// script-visible fields to Roto. The private `cancellation` field must stay
// out of `fields()` because Roto validates every context field type at runtime.
// When adding, removing, or renaming script-visible fields on
// `ExecutionContext`, update this list in the same change.
unsafe impl Context for ExecutionContext {
    fn fields() -> Vec<roto::__internal::ContextField> {
        vec![
            roto::__internal::ContextField {
                name: "timeout_s",
                offset: std::mem::offset_of!(Self, timeout_s),
                type_name: std::any::type_name::<u32>(),
                type_id: std::any::TypeId::of::<u32>(),
                docstring: "Maximum time the script may run, in seconds.".to_string(),
            },
            roto::__internal::ContextField {
                name: "max_output_bytes",
                offset: std::mem::offset_of!(Self, max_output_bytes),
                type_name: std::any::type_name::<u64>(),
                type_id: std::any::TypeId::of::<u64>(),
                docstring: "Maximum number of bytes the script may write to output.".to_string(),
            },
            roto::__internal::ContextField {
                name: "max_temp_bytes",
                offset: std::mem::offset_of!(Self, max_temp_bytes),
                type_name: std::any::type_name::<u64>(),
                type_id: std::any::TypeId::of::<u64>(),
                docstring: "Maximum number of bytes the script may use in temporary space."
                    .to_string(),
            },
            roto::__internal::ContextField {
                name: "metadata",
                offset: std::mem::offset_of!(Self, metadata),
                type_name: std::any::type_name::<Arc<str>>(),
                type_id: std::any::TypeId::of::<Arc<str>>(),
                docstring: "Per-run metadata as a string.".to_string(),
            },
        ]
    }
}

impl ExecutionContext {
    /// Build a context from the request-level [`RunLimits`].
    pub fn from_limits(limits: &RunLimits) -> Self {
        Self::from_limits_and_cancellation(limits, RunCancellation::with_timeout(limits.timeout_s))
    }

    pub fn from_limits_and_cancellation(limits: &RunLimits, cancellation: RunCancellation) -> Self {
        Self {
            timeout_s: limits.timeout_s,
            max_output_bytes: limits.max_output_bytes as u64,
            max_temp_bytes: limits.max_temp_bytes as u64,
            metadata: limits.metadata.clone().unwrap_or_default().into(),
            cancellation,
        }
    }

    pub fn deadline(&self) -> Option<Instant> {
        self.cancellation.deadline()
    }

    pub fn is_cancelled(&self) -> bool {
        self.cancellation.is_cancelled()
    }

    pub fn check_cancelled(&self) -> Result<()> {
        self.cancellation.check_cancelled()
    }
}

#[cfg(test)]
mod tests {
    use std::thread;

    use super::*;

    #[test]
    fn no_timeout_context_is_not_cancelled() {
        let limits = RunLimits::default();
        let ctx = ExecutionContext::from_limits(&limits);
        assert!(!ctx.is_cancelled());
        assert!(ctx.check_cancelled().is_ok());
        assert!(ctx.deadline().is_none());
    }

    #[test]
    fn explicit_cancel_is_reported() {
        let cancellation = RunCancellation::none();
        let ctx = ExecutionContext::from_limits_and_cancellation(
            &RunLimits::default(),
            cancellation.clone(),
        );

        cancellation.cancel();

        let err = ctx.check_cancelled().unwrap_err();
        assert!(ctx.is_cancelled());
        assert_eq!(
            err.diagnostics()[0].code.as_deref(),
            Some("execution-cancelled")
        );
    }

    #[test]
    fn timeout_is_reported() {
        let cancellation = RunCancellation::new(Some(Instant::now() + Duration::from_millis(1)));
        let ctx =
            ExecutionContext::from_limits_and_cancellation(&RunLimits::default(), cancellation);

        thread::sleep(Duration::from_millis(5));

        let err = ctx.check_cancelled().unwrap_err();
        assert!(ctx.is_cancelled());
        assert_eq!(
            err.diagnostics()[0].code.as_deref(),
            Some("timeout-exceeded")
        );
    }
}
