package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"time"
)

type usage struct {
	InputTokens               int `json:"input_tokens"`
	OutputTokens              int `json:"output_tokens"`
	CacheCreationInputTokens  int `json:"cache_creation_input_tokens"`
	CacheReadInputTokens      int `json:"cache_read_input_tokens"`
}

type message struct {
	Usage usage `json:"usage"`
}

type entry struct {
	Message message `json:"message"`
}

type record struct {
	Timestamp  string `json:"timestamp"`
	Date       string `json:"date"`
	User       string `json:"user"`
	Input      int    `json:"input"`
	Output     int    `json:"output"`
	CacheWrite int    `json:"cache_write"`
	CacheRead  int    `json:"cache_read"`
	Actual     int    `json:"actual"`
}

func main() {
	home, _ := os.UserHomeDir()
	user := os.Getenv("USER")
	if user == "" {
		user = filepath.Base(home)
	}

	var total record
	pattern := filepath.Join(home, ".claude", "projects", "**", "*.jsonl")
	matches, _ := filepath.Glob(pattern)

	// filepath.Glob doesn't support **, so walk manually
	projectsDir := filepath.Join(home, ".claude", "projects")
	filepath.Walk(projectsDir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() || filepath.Ext(path) != ".jsonl" {
			return nil
		}
		_ = matches
		f, err := os.Open(path)
		if err != nil {
			return nil
		}
		defer f.Close()

		scanner := bufio.NewScanner(f)
		scanner.Buffer(make([]byte, 1024*1024), 1024*1024)
		for scanner.Scan() {
			var e entry
			if json.Unmarshal(scanner.Bytes(), &e) == nil {
				u := e.Message.Usage
				total.Input += u.InputTokens
				total.Output += u.OutputTokens
				total.CacheWrite += u.CacheCreationInputTokens
				total.CacheRead += u.CacheReadInputTokens
			}
		}
		return nil
	})

	now := time.Now().UTC()
	total.Timestamp = now.Format(time.RFC3339)
	total.Date = now.Format("2006-01-02")
	total.User = user
	total.Actual = total.Input + total.Output

	backup := filepath.Join(home, ".claude", "token-stats-backup.jsonl")
	f, err := os.OpenFile(backup, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		os.Exit(1)
	}
	defer f.Close()

	line, _ := json.Marshal(total)
	fmt.Fprintf(f, "%s\n", line)
}
