use std::time::{Duration, Instant};

use crank_core::{ExecutionFailure, ExecutionOrigin, InvocationSource, WorkspaceId};
use crank_runtime::{
    ExecutionAuthorization, ResolvedAuth, RuntimeExecutionRequest, RuntimeExecutor,
    RuntimeOperation, RuntimeRequestContext,
};
use serde_json::Value;

#[allow(dead_code)]
pub trait RuntimeExecutorTestExt {
    async fn execute(
        &self,
        operation: &RuntimeOperation,
        input: &Value,
    ) -> Result<Value, ExecutionFailure>;

    async fn execute_with_context(
        &self,
        operation: &RuntimeOperation,
        input: &Value,
        context: Option<&RuntimeRequestContext>,
    ) -> Result<Value, ExecutionFailure>;

    async fn execute_with_auth_and_context(
        &self,
        operation: &RuntimeOperation,
        input: &Value,
        resolved_auth: Option<&ResolvedAuth>,
        context: Option<&RuntimeRequestContext>,
    ) -> Result<Value, ExecutionFailure>;
}

impl RuntimeExecutorTestExt for RuntimeExecutor {
    async fn execute(
        &self,
        operation: &RuntimeOperation,
        input: &Value,
    ) -> Result<Value, ExecutionFailure> {
        self.execute_with_auth_and_context(operation, input, None, None)
            .await
    }

    async fn execute_with_context(
        &self,
        operation: &RuntimeOperation,
        input: &Value,
        context: Option<&RuntimeRequestContext>,
    ) -> Result<Value, ExecutionFailure> {
        self.execute_with_auth_and_context(operation, input, None, context)
            .await
    }

    async fn execute_with_auth_and_context(
        &self,
        operation: &RuntimeOperation,
        input: &Value,
        resolved_auth: Option<&ResolvedAuth>,
        context: Option<&RuntimeRequestContext>,
    ) -> Result<Value, ExecutionFailure> {
        let mut context = context.cloned().unwrap_or_else(|| {
            RuntimeRequestContext::from_correlation(&crank_core::CorrelationContext::generate())
        });
        if context.metering_context().is_none() {
            context = context.with_metering_context(
                WorkspaceId::new("ws_test"),
                None,
                InvocationSource::AdminTestRun,
            );
        }
        let metering = context.metering_context().expect("test metering context");
        let origin = if metering.agent_id.is_some() {
            ExecutionOrigin::AgentSnapshot
        } else {
            ExecutionOrigin::AdminDraft
        };
        let request = RuntimeExecutionRequest::try_new(
            &metering.workspace_id,
            origin,
            metering.agent_id.as_ref(),
            operation,
            input,
            ExecutionAuthorization::Authorized,
            resolved_auth,
            &context,
            Instant::now() + Duration::from_secs(30),
        )
        .map_err(|_| {
            crank_core::ExecutionFailure::new(
                crank_core::ExecutionErrorCode::RuntimeInternal,
                crank_core::CorrelationContext::new(
                    context.request_id.clone(),
                    context.trace_context.clone(),
                ),
            )
        })?;
        self.execute_outcome(request)
            .await
            .map(|success| success.output)
    }
}