use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use crank_core::{HttpMethod, IdempotencyMode, IdempotencyPolicy, Target};
use serde_json::Value;
use sha2::{Digest, Sha256};

use crate::{PreparedRequest, RuntimeError, RuntimeOperation, RuntimeRequestContext};

pub fn prepare_idempotency(
    operation: &RuntimeOperation,
    input: &Value,
    prepared_request: &mut PreparedRequest,
) -> Result<Option<String>, RuntimeError> {
    let Some(policy) = policy(operation) else {
        return Ok(None);
    };

    let key = key_from_policy(policy, input, prepared_request);
    let Some(key) = key else {
        if policy.mode == IdempotencyMode::Required {
            return Err(RuntimeError::InvalidPreparedRequest {
                field: "execution_config.idempotency".to_owned(),
                reason: "required idempotency key was not provided".to_owned(),
            });
        }
        return Ok(None);
    };

    if let Some(header_name) = policy
        .header_name
        .as_deref()
        .filter(|value| !value.is_empty())
    {
        prepared_request
            .headers
            .entry(header_name.to_owned())
            .or_insert_with(|| key.clone());
    }

    Ok(Some(key))
}

pub fn policy(operation: &RuntimeOperation) -> Option<&IdempotencyPolicy> {
    let is_mutating_rest =
        matches!(&operation.target, Target::Rest(target) if target.method != HttpMethod::Get);
    if !is_mutating_rest {
        return None;
    }

    let policy = operation.execution_config.idempotency.as_ref()?;
    if policy.mode == IdempotencyMode::Disabled || policy.ttl_ms == 0 {
        return None;
    }

    Some(policy)
}

pub fn cache_key(
    operation: &RuntimeOperation,
    idempotency_key: &str,
    request_context: Option<&RuntimeRequestContext>,
) -> Option<String> {
    if idempotency_key.is_empty() {
        return None;
    }

    let scope = request_context?.response_cache_scope()?;
    let key_hash = URL_SAFE_NO_PAD.encode(Sha256::digest(idempotency_key.as_bytes()));

    Some(format!(
        "crank:idempotency:workspace:{}:agent:{}:operation:{}:version:{}:key:{}",
        scope.workspace_key,
        scope.agent_key,
        operation.operation_id.as_str(),
        operation.operation_version,
        key_hash
    ))
}

fn key_from_policy(
    policy: &IdempotencyPolicy,
    input: &Value,
    prepared_request: &PreparedRequest,
) -> Option<String> {
    if let Some(header_name) = policy
        .header_name
        .as_deref()
        .filter(|value| !value.is_empty())
    {
        if let Some(value) = prepared_request.headers.get(header_name) {
            return Some(value.clone());
        }
    }

    let field_name = policy.input_field.as_deref()?.trim();
    if field_name.is_empty() {
        return None;
    }

    input.get(field_name).and_then(|value| match value {
        Value::String(value) if !value.trim().is_empty() => Some(value.clone()),
        Value::Number(value) => Some(value.to_string()),
        _ => None,
    })
}