// usage.mbt — upstream `src/usage.ts` mirror: per-model usage totals and the
// `_lody/session/usage_update` notification payload.

///|
/// `ModelUsage` — token/cost totals for one model bucket.
pub(all) struct ModelUsage {
  input_tokens : Int
  output_tokens : Int
  cache_read_input_tokens : Int
  cache_creation_input_tokens : Int?
  reasoning_output_tokens : Int?
  web_search_requests : Int?
  cost_usd : Double?
  context_window : Int64?
} derive(Eq, Debug)

///|
pub fn ModelUsage::to_json(self : ModelUsage) -> Json {
  let pairs : Array[(String, Json)] = [
    ("inputTokens", Json::number(self.input_tokens.to_double())),
    ("outputTokens", Json::number(self.output_tokens.to_double())),
    (
      "cacheReadInputTokens",
      Json::number(self.cache_read_input_tokens.to_double()),
    ),
  ]
  match self.cache_creation_input_tokens {
    Some(v) =>
      pairs.push(("cacheCreationInputTokens", Json::number(v.to_double())))
    None => ()
  }
  match self.reasoning_output_tokens {
    Some(v) =>
      pairs.push(("reasoningOutputTokens", Json::number(v.to_double())))
    None => ()
  }
  match self.web_search_requests {
    Some(v) => pairs.push(("webSearchRequests", Json::number(v.to_double())))
    None => ()
  }
  match self.cost_usd {
    Some(v) => pairs.push(("costUSD", Json::number(v)))
    None => ()
  }
  match self.context_window {
    Some(v) => pairs.push(("contextWindow", Json::number(v.to_double())))
    None => ()
  }
  Json::object(Map::from_array(pairs))
}

///|
pub fn ModelUsage::from_json(raw : Json) -> Result[ModelUsage, String] {
  try {
    let ctx = "lody.usage.model"
    let fields = object_fields(raw, ctx)
    Ok(ModelUsage::{
      input_tokens: req_int(fields, ctx, "inputTokens"),
      output_tokens: req_int(fields, ctx, "outputTokens"),
      cache_read_input_tokens: req_int(fields, ctx, "cacheReadInputTokens"),
      cache_creation_input_tokens: opt_int(
        fields, ctx, "cacheCreationInputTokens",
      ),
      reasoning_output_tokens: opt_int(fields, ctx, "reasoningOutputTokens"),
      web_search_requests: opt_int(fields, ctx, "webSearchRequests"),
      cost_usd: opt_double(fields, ctx, "costUSD"),
      context_window: opt_int64(fields, ctx, "contextWindow"),
    })
  } catch {
    DecodeError::Msg(message) => Err(message)
  }
}

///|
/// `SessionUsageUpdate` — `_lody/session/usage_update` notification payload:
/// whole-session usage plus optional per-model buckets.
pub(all) struct SessionUsageUpdate {
  session_id : String
  usage : ModelUsage
  model_usage : Map[String, ModelUsage]?
} derive(Eq, Debug)

///|
pub fn SessionUsageUpdate::to_json(self : SessionUsageUpdate) -> Json {
  let pairs : Array[(String, Json)] = [
    ("sessionId", Json::string(self.session_id)),
    ("usage", self.usage.to_json()),
  ]
  match self.model_usage {
    Some(buckets) => {
      let bucket_pairs : Array[(String, Json)] = []
      for model_id, usage in buckets {
        bucket_pairs.push((model_id, usage.to_json()))
      }
      pairs.push(("modelUsage", Json::object(Map::from_array(bucket_pairs))))
    }
    None => ()
  }
  Json::object(Map::from_array(pairs))
}

///|
pub fn SessionUsageUpdate::from_json(
  raw : Json,
) -> Result[SessionUsageUpdate, String] {
  try {
    let ctx = "lody.usage.update"
    let fields = object_fields(raw, ctx)
    let usage = req_sub(fields, ctx, "usage", ModelUsage::from_json)
    let model_usage : Map[String, ModelUsage]? = match
      opt_object(fields, ctx, "modelUsage") {
      None => None
      Some(buckets) => {
        let parsed : Array[(String, ModelUsage)] = []
        for model_id, bucket in buckets {
          match ModelUsage::from_json(bucket) {
            Ok(usage) => parsed.push((model_id, usage))
            Err(message) => raise DecodeError::Msg(message)
          }
        }
        Some(Map::from_array(parsed))
      }
    }
    Ok(SessionUsageUpdate::{
      session_id: req_string(fields, ctx, "sessionId"),
      usage,
      model_usage,
    })
  } catch {
    DecodeError::Msg(message) => Err(message)
  }
}