//! Read/write helpers for the `tokens` LWW-map collection.
//!
//! Token-usage entries are keyed by `agent_id` and track per-agent input/output
//! token counts and a timestamp for the most recent write.

use bft_json_crdt::json_crdt::{JsonValue, *};
use bft_json_crdt::op::ROOT_ID;
use serde_json::json;

use super::super::state::{apply_and_persist, get_crdt, rebuild_token_index};
use super::super::types::{TokenUsageCrdt, TokenUsageView};
use super::list_id_at;

/// Write or update a token-usage entry keyed by `agent_id`.
///
/// If an entry for `agent_id` already exists it is updated in place;
/// otherwise a new entry is inserted.
pub fn write_token_usage(
    agent_id: &str,
    story_id: &str,
    input_tokens: f64,
    output_tokens: f64,
    timestamp: f64,
) {
    let Some(state_mutex) = get_crdt() else {
        return;
    };
    let Ok(mut state) = state_mutex.lock() else {
        return;
    };

    if let Some(&idx) = state.token_index.get(agent_id) {
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.tokens[idx].story_id.set(story_id.to_string())
        });
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.tokens[idx].input_tokens.set(input_tokens)
        });
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.tokens[idx].output_tokens.set(output_tokens)
        });
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.tokens[idx].timestamp.set(timestamp)
        });
    } else {
        let entry: JsonValue = json!({
            "agent_id": agent_id,
            "story_id": story_id,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "timestamp": timestamp,
        })
        .into();
        apply_and_persist(&mut state, |s| s.crdt.doc.tokens.insert(ROOT_ID, entry));
        state.token_index = rebuild_token_index(&state.crdt);
    }
}

/// Read all token-usage entries.
pub fn read_all_token_usage() -> Option<Vec<TokenUsageView>> {
    let state_mutex = get_crdt()?;
    let state = state_mutex.lock().ok()?;
    let mut out = Vec::new();
    for entry in state.crdt.doc.tokens.iter() {
        if let Some(v) = extract_token_view(entry) {
            out.push(v);
        }
    }
    Some(out)
}

/// Read a single token-usage entry by `agent_id`.
pub fn read_token_usage(agent_id: &str) -> Option<TokenUsageView> {
    let state_mutex = get_crdt()?;
    let state = state_mutex.lock().ok()?;
    let &idx = state.token_index.get(agent_id)?;
    extract_token_view(&state.crdt.doc.tokens[idx])
}

/// Tombstone a token-usage entry by `agent_id`.
///
/// Returns `true` if the entry existed and a delete op was issued.
pub fn delete_token_usage(agent_id: &str) -> bool {
    let Some(state_mutex) = get_crdt() else {
        return false;
    };
    let Ok(mut state) = state_mutex.lock() else {
        return false;
    };
    let Some(&idx) = state.token_index.get(agent_id) else {
        return false;
    };
    let Some(op_id) = list_id_at(&state.crdt.doc.tokens, idx) else {
        return false;
    };
    apply_and_persist(&mut state, |s| s.crdt.doc.tokens.delete(op_id));
    state.token_index = rebuild_token_index(&state.crdt);
    true
}

/// Convert a CRDT token-usage entry into its read-only view representation.
pub(super) fn extract_token_view(entry: &TokenUsageCrdt) -> Option<TokenUsageView> {
    let agent_id = match entry.agent_id.view() {
        JsonValue::String(s) if !s.is_empty() => s,
        _ => return None,
    };
    let story_id = match entry.story_id.view() {
        JsonValue::String(s) if !s.is_empty() => Some(s),
        _ => None,
    };
    let input_tokens = match entry.input_tokens.view() {
        JsonValue::Number(n) => n,
        _ => 0.0,
    };
    let output_tokens = match entry.output_tokens.view() {
        JsonValue::Number(n) => n,
        _ => 0.0,
    };
    let timestamp = match entry.timestamp.view() {
        JsonValue::Number(n) => n,
        _ => 0.0,
    };
    Some(TokenUsageView {
        agent_id,
        story_id,
        input_tokens,
        output_tokens,
        timestamp,
    })
}
