ccd-cli 1.0.0-beta.4

use std::env;
use std::fs::{self, File};
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use serde::Deserialize;
use serde_json::Value;

use super::string_key;
use crate::telemetry::host::{self as host_telemetry, HostContextSnapshot};

pub(crate) fn current() -> Result<Option<HostContextSnapshot>> {
    let thread_id = match env::var("CODEX_THREAD_ID") {
        Ok(value) if !value.trim().is_empty() => value,
        _ => return Ok(None),
    };

    let codex_home = codex_home()?;
    let session_path = match find_session_log(&codex_home.join("sessions"), &thread_id)? {
        Some(path) => path,
        None => return Ok(None),
    };

    let observed_at_epoch_s = match host_telemetry::file_mtime_epoch_s(&session_path)? {
        Some(epoch_s) => epoch_s,
        None => return Ok(None),
    };

    latest_token_count(&session_path, observed_at_epoch_s)
}

fn codex_home() -> Result<PathBuf> {
    if let Some(home) = env::var_os("CODEX_HOME") {
        return Ok(PathBuf::from(home));
    }

    match env::var_os("HOME") {
        Some(home) => Ok(PathBuf::from(home).join(".codex")),
        None => anyhow::bail!("HOME environment variable is not set"),
    }
}

fn find_session_log(root: &Path, thread_id: &str) -> Result<Option<PathBuf>> {
    if !root.is_dir() {
        return Ok(None);
    }

    let mut stack = vec![root.to_path_buf()];
    while let Some(dir) = stack.pop() {
        for entry in
            fs::read_dir(&dir).with_context(|| format!("failed to read {}", dir.display()))?
        {
            let entry =
                entry.with_context(|| format!("failed to read entry in {}", dir.display()))?;
            let path = entry.path();
            if path.is_dir() {
                stack.push(path);
                continue;
            }

            let matches = path
                .file_name()
                .and_then(|name| name.to_str())
                .map(|name| name.contains(thread_id) && name.ends_with(".jsonl"))
                .unwrap_or(false);
            if matches {
                return Ok(Some(path));
            }
        }
    }

    Ok(None)
}

fn latest_token_count(
    path: &Path,
    observed_at_epoch_s: u64,
) -> Result<Option<HostContextSnapshot>> {
    let file = File::open(path).with_context(|| format!("failed to open {}", path.display()))?;
    let reader = BufReader::new(file);
    let mut latest = None;
    let mut latest_model = host_telemetry::env_string(&["CCD_CODEX_MODEL", "CODEX_MODEL"])
        .or_else(|| host_telemetry::env_string(&["CCD_HOST_MODEL"]));

    for line in reader.lines() {
        let line = line.with_context(|| format!("failed to read {}", path.display()))?;
        let value: Value = match serde_json::from_str(&line) {
            Ok(value) => value,
            Err(_) => continue,
        };
        if let Some(model_name) = string_key(
            &value,
            &[
                "model",
                "model_name",
                "modelName",
                "model_slug",
                "modelSlug",
            ],
        ) {
            latest_model = Some(model_name);
        }
        let entry: SessionEntry = match serde_json::from_value(value) {
            Ok(entry) => entry,
            Err(_) => continue,
        };

        if entry.entry_type != "event_msg" {
            continue;
        }

        let payload = match entry.payload {
            Some(payload) => payload,
            None => continue,
        };

        if payload.payload_type != "token_count" {
            continue;
        }

        // Precedence: native payload -> CCD_CONTEXT_WINDOW_TOKENS
        // (general fallback for the rare case where the Codex runtime
        // emits token_count events without a populated
        // model_context_window field). See ccd#529.
        let native_window = match payload.info.model_context_window {
            0 => None,
            value => Some(value),
        };
        let context_window = native_window.or_else(host_telemetry::general_context_window);

        let total_tokens = payload
            .info
            .last_token_usage
            .as_ref()
            .map(|usage| usage.total_tokens)
            .unwrap_or(payload.info.total_token_usage.total_tokens);
        latest = Some(HostContextSnapshot {
            host: "codex",
            observed_at_epoch_s,
            model_name: latest_model.clone(),
            context_used_pct: host_telemetry::compute_pct(total_tokens, context_window),
            total_tokens: Some(total_tokens),
            model_context_window: context_window,
            compacted: None,
            cost_usage: host_telemetry::HostCostUsage {
                blended_total_tokens: Some(payload.info.total_token_usage.total_tokens),
                ..host_telemetry::HostCostUsage::default()
            },
        });
    }

    Ok(latest)
}

#[derive(Deserialize)]
struct SessionEntry {
    #[serde(rename = "type")]
    entry_type: String,
    payload: Option<TokenCountPayload>,
}

#[derive(Deserialize)]
struct TokenCountPayload {
    #[serde(rename = "type")]
    payload_type: String,
    info: TokenCountInfo,
}

#[derive(Deserialize)]
struct TokenCountInfo {
    total_token_usage: TotalTokenUsage,
    last_token_usage: Option<TotalTokenUsage>,
    model_context_window: u64,
}

#[derive(Deserialize)]
struct TotalTokenUsage {
    total_tokens: u64,
}