enlil 0.1.1

use serde::{Deserialize, Serialize};

/// Represents the 4-layer token attribution model for a given request.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TokenAttribution {
    /// Tokens consumed by the base prompt instructions and user query.
    pub prompt_tokens: u32,
    /// Tokens consumed by tool definitions and schemas.
    pub tool_tokens: u32,
    /// Tokens consumed by retrieved context (RAG/Memory).
    pub memory_tokens: u32,
    /// Tokens generated by the model in the response.
    pub response_tokens: u32,
}

impl TokenAttribution {
    pub fn total_tokens(&self) -> u32 {
        self.prompt_tokens + self.tool_tokens + self.memory_tokens + self.response_tokens
    }
}

/// A standard LLM API Usage object (e.g., OpenAI format)
#[derive(Debug, Clone, Deserialize)]
pub struct ProviderUsage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
}

/// Canonical, provider-agnostic token usage broken out by billing category.
///
/// Different providers report prompt caching differently:
/// - **OpenAI**: `usage.prompt_tokens` *includes* cached input; the cached slice is
///   reported under `usage.prompt_tokens_details.cached_tokens`. There is no separate
///   cache-write charge.
/// - **Anthropic**: `usage.input_tokens` is the *uncached* input; cache reads and cache
///   writes are reported separately as `cache_read_input_tokens` /
///   `cache_creation_input_tokens` and are billed at different rates.
///
/// We normalize both into disjoint buckets so pricing is unambiguous.
#[derive(Debug, Clone, Default, PartialEq, Serialize)]
pub struct TokenUsage {
    /// Input tokens billed at the full input rate.
    pub uncached_input_tokens: u32,
    /// Input tokens served from the provider's prompt cache (billed at a discount).
    pub cached_read_tokens: u32,
    /// Input tokens written to the provider's prompt cache (Anthropic; billed at a premium).
    pub cache_write_tokens: u32,
    /// Generated output tokens.
    pub output_tokens: u32,
}

impl TokenUsage {
    /// All input tokens across every billing category.
    pub fn total_input(&self) -> u32 {
        self.uncached_input_tokens + self.cached_read_tokens + self.cache_write_tokens
    }

    /// Grand total of input + output tokens.
    pub fn total_tokens(&self) -> u32 {
        self.total_input() + self.output_tokens
    }

    /// Legacy `prompt_tokens`-equivalent (all input tokens) for 4-layer attribution.
    pub fn as_provider_usage(&self) -> ProviderUsage {
        ProviderUsage {
            prompt_tokens: self.total_input(),
            completion_tokens: self.output_tokens,
            total_tokens: self.total_tokens(),
        }
    }
}

/// Parses a provider `usage` JSON object into the canonical [`TokenUsage`],
/// transparently handling both OpenAI and Anthropic response shapes.
/// Returns `None` if the object carries no recognizable token counts.
pub fn parse_usage(usage: &serde_json::Value) -> Option<TokenUsage> {
    // --- OpenAI shape: prompt_tokens (incl. cached) / completion_tokens ---
    if let Some(prompt) = usage.get("prompt_tokens").and_then(|v| v.as_u64()) {
        let completion = usage
            .get("completion_tokens")
            .and_then(|v| v.as_u64())
            .unwrap_or(0);
        // cached_tokens is a subset of prompt_tokens.
        let cached = usage
            .get("prompt_tokens_details")
            .and_then(|d| d.get("cached_tokens"))
            .and_then(|v| v.as_u64())
            .unwrap_or(0)
            .min(prompt);
        return Some(TokenUsage {
            uncached_input_tokens: (prompt - cached) as u32,
            cached_read_tokens: cached as u32,
            cache_write_tokens: 0,
            output_tokens: completion as u32,
        });
    }

    // --- Anthropic shape: input_tokens (uncached) / output_tokens + cache_* ---
    if let Some(input) = usage.get("input_tokens").and_then(|v| v.as_u64()) {
        let output = usage
            .get("output_tokens")
            .and_then(|v| v.as_u64())
            .unwrap_or(0);
        let cache_read = usage
            .get("cache_read_input_tokens")
            .and_then(|v| v.as_u64())
            .unwrap_or(0);
        let cache_write = usage
            .get("cache_creation_input_tokens")
            .and_then(|v| v.as_u64())
            .unwrap_or(0);
        return Some(TokenUsage {
            uncached_input_tokens: input as u32,
            cached_read_tokens: cache_read as u32,
            cache_write_tokens: cache_write as u32,
            output_tokens: output as u32,
        });
    }

    None
}

/// Estimates tool and memory token counts from the request body.
/// Uses char_count / 4 as industry-standard token approximation.
pub fn estimate_token_layers(body: &[u8]) -> (u32, u32) {
    match serde_json::from_slice::<serde_json::Value>(body) {
        Ok(json) => estimate_token_layers_value(&json),
        Err(_) => (0, 0),
    }
}

/// Same as [`estimate_token_layers`] but from an already-parsed JSON value.
pub fn estimate_token_layers_value(json: &serde_json::Value) -> (u32, u32) {
    let tool_chars = json
        .get("tools")
        .or_else(|| json.get("functions"))
        .map(|v| v.to_string().len())
        .unwrap_or(0);

    let memory_chars = json
        .get("messages")
        .and_then(|m| m.as_array())
        .map(|msgs| {
            msgs.iter()
                .filter(|m| m.get("role").and_then(|r| r.as_str()) == Some("system"))
                .map(|m| m.get("content").map(|c| c.to_string().len()).unwrap_or(0))
                .sum::<usize>()
        })
        .unwrap_or(0);

    ((tool_chars / 4) as u32, (memory_chars / 4) as u32)
}

/// Splits raw provider usage into 4-layer attribution using estimated tool/memory tokens.
pub fn calculate_attribution(
    usage: &ProviderUsage,
    estimated_tool_tokens: u32,
    estimated_memory_tokens: u32,
) -> TokenAttribution {
    let tool = estimated_tool_tokens.min(usage.prompt_tokens);
    let memory = estimated_memory_tokens.min(usage.prompt_tokens.saturating_sub(tool));
    let prompt = usage
        .prompt_tokens
        .saturating_sub(tool)
        .saturating_sub(memory);

    TokenAttribution {
        prompt_tokens: prompt,
        tool_tokens: tool,
        memory_tokens: memory,
        response_tokens: usage.completion_tokens,
    }
}

// ---------------------------------------------------------------------------
// Model pricing — OSS.
//
// Pure functions: (model, token usage) -> cost in microdollars. These live on the
// OSS side because cost *visibility* is a developer feature (the observability
// wedge) and the local trace store records per-request cost. Stateful, per-tenant
// budget/spend accumulation stays proprietary in `finops::cost_tracker`.
// See DEVELOPMENT_PLAN.md, Step 2.
// ---------------------------------------------------------------------------

/// Cost per 1M tokens in microdollars (1 microdollar = $0.000001)
struct ModelPricing {
    input_per_1m: u64,
    output_per_1m: u64,
    /// Price for input tokens served from the provider's prompt cache.
    cached_read_per_1m: u64,
    /// Price for input tokens written to the provider's prompt cache.
    cache_write_per_1m: u64,
}

impl ModelPricing {
    /// OpenAI-style caching: cached input is billed at ~50% of the input rate and
    /// there is no separate cache-write surcharge.
    fn openai(input_per_1m: u64, output_per_1m: u64) -> Self {
        Self {
            input_per_1m,
            output_per_1m,
            cached_read_per_1m: input_per_1m / 2,
            cache_write_per_1m: input_per_1m,
        }
    }

    /// Anthropic-style caching: cache reads are billed at 0.1x the input rate and
    /// cache writes (creation) at 1.25x the input rate.
    fn anthropic(input_per_1m: u64, output_per_1m: u64) -> Self {
        Self {
            input_per_1m,
            output_per_1m,
            cached_read_per_1m: input_per_1m / 10,
            cache_write_per_1m: input_per_1m * 5 / 4,
        }
    }
}

fn get_pricing(model: &str) -> ModelPricing {
    match model {
        m if m.contains("gpt-4o-mini") => ModelPricing::openai(150_000, 600_000),
        m if m.contains("gpt-4o") => ModelPricing::openai(2_500_000, 10_000_000),
        m if m.contains("gpt-4-turbo") => ModelPricing::openai(10_000_000, 30_000_000),
        m if m.contains("gpt-4") => ModelPricing::openai(30_000_000, 60_000_000),
        m if m.contains("gpt-3.5") => ModelPricing::openai(500_000, 1_500_000),
        m if m.contains("claude-3-5-sonnet") => ModelPricing::anthropic(3_000_000, 15_000_000),
        m if m.contains("claude-3-opus") => ModelPricing::anthropic(15_000_000, 75_000_000),
        m if m.contains("claude-3-haiku") => ModelPricing::anthropic(250_000, 1_250_000),
        m if m.contains("claude") => ModelPricing::anthropic(3_000_000, 15_000_000),
        _ => ModelPricing::openai(2_500_000, 10_000_000), // default to gpt-4o
    }
}

/// Calculate cost in microdollars from a simple prompt/completion split.
/// Backwards-compatible helper: treats all prompt tokens as full-price (uncached) input.
pub fn calculate_cost(model: &str, prompt_tokens: u32, completion_tokens: u32) -> u64 {
    calculate_cost_detailed(
        model,
        &TokenUsage {
            uncached_input_tokens: prompt_tokens,
            output_tokens: completion_tokens,
            ..Default::default()
        },
    )
}

/// Calculate the *actual* cost in microdollars, pricing each billing category
/// (uncached input, cached read, cache write, output) at its own rate.
pub fn calculate_cost_detailed(model: &str, usage: &TokenUsage) -> u64 {
    let p = get_pricing(model);
    (usage.uncached_input_tokens as u64 * p.input_per_1m
        + usage.cached_read_tokens as u64 * p.cached_read_per_1m
        + usage.cache_write_tokens as u64 * p.cache_write_per_1m
        + usage.output_tokens as u64 * p.output_per_1m)
        / 1_000_000
}

/// Calculate the *naive* cost in microdollars — every input token billed at the
/// full input rate, as a simplistic "total tokens × list price" counter would.
/// The gap between this and [`calculate_cost_detailed`] is the prompt-cache saving.
pub fn calculate_cost_naive(model: &str, usage: &TokenUsage) -> u64 {
    let p = get_pricing(model);
    (usage.total_input() as u64 * p.input_per_1m + usage.output_tokens as u64 * p.output_per_1m)
        / 1_000_000
}

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

    #[test]
    fn test_estimate_token_layers_with_tools_and_system() {
        let body = serde_json::json!({
            "model": "gpt-4o",
            "messages": [
                {"role": "system", "content": "You are a helpful assistant with access to company docs."},
                {"role": "user", "content": "What is our refund policy?"}
            ],
            "tools": [
                {"type": "function", "function": {"name": "search_docs", "parameters": {"type": "object"}}}
            ]
        });
        let bytes = serde_json::to_vec(&body).unwrap();
        let (tool, memory) = estimate_token_layers(&bytes);
        assert!(tool > 0, "tool tokens should be > 0");
        assert!(memory > 0, "memory tokens should be > 0");
    }

    #[test]
    fn test_estimate_no_tools_no_memory() {
        let body = serde_json::json!({
            "model": "gpt-4o",
            "messages": [{"role": "user", "content": "Hello"}]
        });
        let bytes = serde_json::to_vec(&body).unwrap();
        let (tool, memory) = estimate_token_layers(&bytes);
        assert_eq!(tool, 0);
        assert_eq!(memory, 0);
    }

    #[test]
    fn test_calculate_attribution_with_layers() {
        let usage = ProviderUsage {
            prompt_tokens: 100,
            completion_tokens: 50,
            total_tokens: 150,
        };
        let attr = calculate_attribution(&usage, 20, 30);
        assert_eq!(attr.tool_tokens, 20);
        assert_eq!(attr.memory_tokens, 30);
        assert_eq!(attr.prompt_tokens, 50);
        assert_eq!(attr.response_tokens, 50);
        assert_eq!(attr.total_tokens(), 150);
    }

    #[test]
    fn test_parse_usage_openai_with_cache() {
        let usage = serde_json::json!({
            "prompt_tokens": 1000,
            "completion_tokens": 100,
            "total_tokens": 1100,
            "prompt_tokens_details": { "cached_tokens": 200 }
        });
        let u = parse_usage(&usage).unwrap();
        assert_eq!(u.uncached_input_tokens, 800);
        assert_eq!(u.cached_read_tokens, 200);
        assert_eq!(u.cache_write_tokens, 0);
        assert_eq!(u.output_tokens, 100);
        assert_eq!(u.total_input(), 1000);
        assert_eq!(u.total_tokens(), 1100);
    }

    #[test]
    fn test_parse_usage_openai_no_cache() {
        let usage = serde_json::json!({ "prompt_tokens": 500, "completion_tokens": 50, "total_tokens": 550 });
        let u = parse_usage(&usage).unwrap();
        assert_eq!(u.uncached_input_tokens, 500);
        assert_eq!(u.cached_read_tokens, 0);
        assert_eq!(u.output_tokens, 50);
    }

    #[test]
    fn test_parse_usage_anthropic() {
        let usage = serde_json::json!({
            "input_tokens": 1000,
            "output_tokens": 200,
            "cache_read_input_tokens": 500,
            "cache_creation_input_tokens": 400
        });
        let u = parse_usage(&usage).unwrap();
        assert_eq!(u.uncached_input_tokens, 1000);
        assert_eq!(u.cached_read_tokens, 500);
        assert_eq!(u.cache_write_tokens, 400);
        assert_eq!(u.output_tokens, 200);
        assert_eq!(u.total_input(), 1900);
    }

    #[test]
    fn test_parse_usage_empty() {
        assert!(parse_usage(&serde_json::json!({})).is_none());
    }
}