// Package metrics writes a JSONL event log: one JSON object per line, one line
// per thing worth counting later. It is wired in before there is anything to
// measure on purpose — retrofitting instrumentation is how projects end up with
// no numbers.
package metrics

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/johannSo/codey/internal/statedir"
)

// DefaultPath is where Codey logs unless told otherwise.
const DefaultPath = statedir.Dir + "/events.jsonl"

// Event is one line of the log. Extra is flattened into the same JSON object so
// the file stays greppable with jq without nested lookups.
type Event struct {
	Time    time.Time `json:"ts"`
	Session string    `json:"session,omitempty"`
	Type    string    `json:"type"`

	// Tool dispatch.
	Tool     string `json:"tool,omitempty"`
	Duration int64  `json:"dur_ms,omitempty"`
	OK       *bool  `json:"ok,omitempty"`
	Err      string `json:"err,omitempty"`

	// Turn accounting.
	Turn       int    `json:"turn,omitempty"`
	StopReason string `json:"stop_reason,omitempty"`

	// Token accounting, mirroring the Messages API usage block.
	InputTokens      int `json:"input_tokens,omitempty"`
	OutputTokens     int `json:"output_tokens,omitempty"`
	CacheReadTokens  int `json:"cache_read_input_tokens,omitempty"`
	CacheWriteTokens int `json:"cache_creation_input_tokens,omitempty"`

	Extra map[string]any `json:"-"`
}

// MarshalJSON flattens Extra alongside the declared fields. Keys in Extra never
// clobber a declared field.
func (e Event) MarshalJSON() ([]byte, error) {
	type alias Event
	b, err := json.Marshal(alias(e))
	if err != nil {
		return nil, err
	}
	if len(e.Extra) == 0 {
		return b, nil
	}
	var m map[string]json.RawMessage
	if err := json.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	for k, v := range e.Extra {
		if _, taken := m[k]; taken {
			continue
		}
		ev, err := json.Marshal(v)
		if err != nil {
			return nil, fmt.Errorf("metrics: extra field %q: %w", k, err)
		}
		m[k] = ev
	}
	return json.Marshal(m)
}

// sink is the writer a logger and everything derived from it share. Derived
// loggers exist so a subagent can stamp its id on every line without opening a
// second file, and sharing the lock is what keeps their lines whole.
type sink struct {
	mu     sync.Mutex
	w      io.Writer
	closer io.Closer
}

// Logger appends events to a writer. The zero value and a nil *Logger are both
// usable and simply drop events, so callers never need a nil check.
type Logger struct {
	sink    *sink
	session string
	// agent is the subagent this logger belongs to, empty for the session's
	// own. It is stamped into every event so one log file can be split by
	// agent afterwards.
	agent string
}

// For returns a logger that stamps every event with an agent id and writes to
// the same place this one does. Closing the original closes both; a derived
// logger owns no file of its own.
func (l *Logger) For(agent string) *Logger {
	if l == nil {
		return nil
	}
	return &Logger{sink: l.sink, session: l.session, agent: agent}
}

// Agent reports which subagent this logger stamps, empty for the session's own.
func (l *Logger) Agent() string {
	if l == nil {
		return ""
	}
	return l.agent
}

// Open appends to path, creating parent directories as needed.
func Open(path, session string) (*Logger, error) {
	if dir := filepath.Dir(path); dir != "." {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return nil, fmt.Errorf("metrics: %w", err)
		}
	}
	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
	if err != nil {
		return nil, fmt.Errorf("metrics: %w", err)
	}
	return &Logger{sink: &sink{w: f, closer: f}, session: session}, nil
}

// NewWriter logs to an arbitrary writer. Useful in tests and in the eval runner,
// which keeps a separate log per run.
func NewWriter(w io.Writer, session string) *Logger {
	return &Logger{sink: &sink{w: w}, session: session}
}

// Session returns the session id stamped onto every event.
func (l *Logger) Session() string {
	if l == nil {
		return ""
	}
	return l.session
}

// Log writes one event. Errors are dropped: a broken metrics file must never
// take down a run.
func (l *Logger) Log(ev Event) {
	if l == nil || l.sink == nil || l.sink.w == nil {
		return
	}
	if ev.Time.IsZero() {
		ev.Time = time.Now().UTC()
	}
	if ev.Session == "" {
		ev.Session = l.session
	}
	if l.agent != "" {
		// A copy, because the caller's map is theirs: stamping it in place
		// would mutate a map they may reuse or share with another goroutine.
		extra := make(map[string]any, len(ev.Extra)+1)
		for k, v := range ev.Extra {
			extra[k] = v
		}
		if _, taken := extra["agent"]; !taken {
			extra["agent"] = l.agent
		}
		ev.Extra = extra
	}
	b, err := json.Marshal(ev)
	if err != nil {
		return
	}
	l.sink.mu.Lock()
	defer l.sink.mu.Unlock()
	l.sink.w.Write(append(b, '\n'))
}

// ToolCall records one dispatch. err nil means success.
func (l *Logger) ToolCall(name string, d time.Duration, err error, extra map[string]any) {
	ok := err == nil
	ev := Event{
		Type:     "tool_call",
		Tool:     name,
		Duration: d.Milliseconds(),
		OK:       &ok,
		Extra:    extra,
	}
	if err != nil {
		ev.Err = err.Error()
	}
	l.Log(ev)
}

// Close closes the underlying file, if there is one.
func (l *Logger) Close() error {
	if l == nil || l.sink == nil || l.sink.closer == nil {
		return nil
	}
	return l.sink.closer.Close()
}