package bakeoff

import (
	"bufio"
	"encoding/json"
	"errors"
	"io"
)

// Tokens is what one agent session cost.
type Tokens struct {
	Input, Output, CacheRead, CacheWrite int64
	Turns                                int
	Seconds                              float64
	CostUSD                              float64
}

// Add sums two sessions.
func (t Tokens) Add(o Tokens) Tokens {
	return Tokens{t.Input + o.Input, t.Output + o.Output, t.CacheRead + o.CacheRead, t.CacheWrite + o.CacheWrite,
		t.Turns + o.Turns, t.Seconds + o.Seconds, t.CostUSD + o.CostUSD}
}

func scanLines(r io.Reader, f func([]byte)) error {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 1<<20), 1<<26)
	for sc.Scan() {
		f(sc.Bytes())
	}
	return sc.Err()
}

// ParseClaudeLog reads `claude -p --output-format stream-json` output. A
// session can emit several result events (it resumes itself when
// background work finishes): their per-model usage and cost are
// cumulative for the whole session, so the last one counts, while turns
// and duration are per result, so they are summed. Per-model usage, which
// includes any subagents, is preferred to the top-level usage.
func ParseClaudeLog(r io.Reader) (Tokens, error) {
	type usage struct {
		Input      int64 `json:"input_tokens"`
		Output     int64 `json:"output_tokens"`
		CacheRead  int64 `json:"cache_read_input_tokens"`
		CacheWrite int64 `json:"cache_creation_input_tokens"`
	}
	type modelUsage struct {
		Input      int64 `json:"inputTokens"`
		Output     int64 `json:"outputTokens"`
		CacheRead  int64 `json:"cacheReadInputTokens"`
		CacheWrite int64 `json:"cacheCreationInputTokens"`
	}
	var last *Tokens
	turns, seconds := 0, 0.0
	err := scanLines(r, func(b []byte) {
		var e struct {
			Type       string                `json:"type"`
			Usage      usage                 `json:"usage"`
			ModelUsage map[string]modelUsage `json:"modelUsage"`
			Turns      int                   `json:"num_turns"`
			DurationMs float64               `json:"duration_ms"`
			Cost       float64               `json:"total_cost_usd"`
		}
		if json.Unmarshal(b, &e) != nil || e.Type != "result" {
			return
		}
		t := Tokens{Input: e.Usage.Input, Output: e.Usage.Output, CacheRead: e.Usage.CacheRead, CacheWrite: e.Usage.CacheWrite,
			Turns: e.Turns, Seconds: e.DurationMs / 1000, CostUSD: e.Cost}
		if len(e.ModelUsage) > 0 {
			t.Input, t.Output, t.CacheRead, t.CacheWrite = 0, 0, 0, 0
			for _, m := range e.ModelUsage {
				t.Input += m.Input
				t.Output += m.Output
				t.CacheRead += m.CacheRead
				t.CacheWrite += m.CacheWrite
			}
		}
		turns += e.Turns
		seconds += e.DurationMs / 1000
		last = &t
	})
	if err != nil {
		return Tokens{}, err
	}
	if last == nil {
		return Tokens{}, errors.New("no result event in claude log")
	}
	last.Turns, last.Seconds = turns, seconds
	return *last, nil
}

// ParseCodexLog sums the usage on every turn.completed event of
// `codex exec --json` output. Codex's input count includes cached input.
func ParseCodexLog(r io.Reader) (Tokens, error) {
	var t Tokens
	err := scanLines(r, func(b []byte) {
		var e struct {
			Type  string `json:"type"`
			Usage struct {
				Input      int64 `json:"input_tokens"`
				Cached     int64 `json:"cached_input_tokens"`
				CacheWrite int64 `json:"cache_write_input_tokens"`
				Output     int64 `json:"output_tokens"`
			} `json:"usage"`
		}
		if json.Unmarshal(b, &e) != nil || e.Type != "turn.completed" {
			return
		}
		t.Input += e.Usage.Input
		t.CacheRead += e.Usage.Cached
		t.CacheWrite += e.Usage.CacheWrite
		t.Output += e.Usage.Output
		t.Turns++
	})
	if err == nil && t.Turns == 0 {
		err = errors.New("no turn.completed event in codex log")
	}
	return t, err
}

// Since turns a resumed session's cumulative totals into that phase's own:
// claude reports usage and cost for the whole session on resume, but turns
// and duration for the resumed run only.
func (t Tokens) Since(prev Tokens) Tokens {
	return Tokens{t.Input - prev.Input, t.Output - prev.Output, t.CacheRead - prev.CacheRead, t.CacheWrite - prev.CacheWrite,
		t.Turns, t.Seconds, t.CostUSD - prev.CostUSD}
}