package metrics

import (
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
	"unicode/utf8"
)

const (
	tokenTrackingAuditFilename      = "tokentracking.audit"
	tokenTrackingAuditSchemaVersion = 1
	tokenTrackingAuditMaxBytes      = 10 * 1024 * 1024
	tokenTrackingAuditMaxBackups    = 3

	tokenTrackingAuditMaxLineBytes       = 512 * 1024
	tokenTrackingAuditMaxRawCommandBytes = 64 * 1024
	tokenTrackingAuditMaxEventBytes      = 64
	tokenTrackingAuditMaxRuntimeBytes    = 64
	tokenTrackingAuditMaxToolNameBytes   = 120
	tokenTrackingAuditMaxSourceKindBytes = 64
	tokenTrackingAuditMaxSourceNameBytes = 120
	tokenTrackingAuditMaxDecisionBytes   = 64
	tokenTrackingAuditMaxCandidates      = 64
	tokenTrackingAuditMaxAllocations     = 65
)

// TokenTrackingAuditSource is a bounded source identity attached to an audit
// entry. It deliberately contains no tool arguments or runtime identifiers.
type TokenTrackingAuditSource struct {
	SourceKind     string `json:"kind,omitempty"`
	SourceName     string `json:"name,omitempty"`
	InputEstimate  uint64 `json:"inputEstimate,omitempty"`
	OutputEstimate uint64 `json:"outputEstimate,omitempty"`
}

// TokenTrackingAuditAllocation records one persisted split without retaining
// any source content.
type TokenTrackingAuditAllocation struct {
	SourceKind string `json:"kind,omitempty"`
	SourceName string `json:"name,omitempty"`
	Input      uint64 `json:"input,omitempty"`
	Cached     uint64 `json:"cached,omitempty"`
	Output     uint64 `json:"output,omitempty"`
}

// TokenTrackingAuditEntry is one local diagnostic record explaining a
// token-source classification or association decision.
type TokenTrackingAuditEntry struct {
	SchemaVersion       int                            `json:"schemaVersion"`
	Timestamp           time.Time                      `json:"timestamp"`
	Event               string                         `json:"event,omitempty"`
	Runtime             string                         `json:"runtime,omitempty"`
	ToolName            string                         `json:"toolName,omitempty"`
	RawCommand          string                         `json:"rawCommand,omitempty"`
	RawCommandTruncated bool                           `json:"rawCommandTruncated,omitempty"`
	SourceKind          string                         `json:"sourceKind,omitempty"`
	SourceName          string                         `json:"sourceName,omitempty"`
	Decision            string                         `json:"decision,omitempty"`
	Candidates          []TokenTrackingAuditSource     `json:"candidates,omitempty"`
	Allocations         []TokenTrackingAuditAllocation `json:"allocations,omitempty"`
}

// TokenTrackingAuditRecorder accepts diagnostic token-source audit entries.
type TokenTrackingAuditRecorder interface {
	RecordTokenTrackingAudit(TokenTrackingAuditEntry)
}

type tokenTrackingAuditWriter struct {
	mu             sync.Mutex
	file           *os.File
	path           string
	maxBytes       int64
	writeErrorOnce sync.Once
}

func newTokenTrackingAuditWriter() (*tokenTrackingAuditWriter, error) {
	path := filepath.Join(filepath.Dir(tokenTrackingPath), tokenTrackingAuditFilename)
	return openTokenTrackingAuditWriter(path, tokenTrackingAuditMaxBytes)
}

func openTokenTrackingAuditWriter(path string, maxBytes int64) (*tokenTrackingAuditWriter, error) {
	if maxBytes <= 0 {
		return nil, fmt.Errorf("open token tracking audit: invalid maximum size")
	}
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return nil, fmt.Errorf("create token tracking audit directory: %w", err)
	}
	file, err := openTokenTrackingAuditFile(path)
	if err != nil {
		return nil, err
	}
	return &tokenTrackingAuditWriter{
		file:     file,
		path:     path,
		maxBytes: maxBytes,
	}, nil
}

func openTokenTrackingAuditFile(path string) (*os.File, error) {
	file, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		return nil, fmt.Errorf("open token tracking audit: %w", err)
	}
	if err := file.Chmod(0600); err != nil {
		_ = file.Close()
		return nil, fmt.Errorf("set token tracking audit permissions: %w", err)
	}
	return file, nil
}

// RecordTokenTrackingAudit appends one normalized JSONL record. Persistence
// failures are diagnostic-only and never expose the entry payload to slog.
func (writer *tokenTrackingAuditWriter) RecordTokenTrackingAudit(entry TokenTrackingAuditEntry) {
	if writer == nil {
		return
	}
	line, err := marshalTokenTrackingAuditEntry(entry)
	if err != nil {
		writer.warnWriteError(err)
		return
	}

	writer.mu.Lock()
	defer writer.mu.Unlock()
	if writer.file == nil {
		return
	}
	info, err := writer.file.Stat()
	if err != nil {
		writer.warnWriteError(fmt.Errorf("stat active file: %w", err))
		return
	}
	if info.Size() > 0 && info.Size()+int64(len(line)) > writer.maxBytes {
		if err := writer.rotateLocked(); err != nil {
			writer.warnWriteError(err)
			return
		}
	}
	if _, err := writer.file.Write(line); err != nil {
		writer.warnWriteError(fmt.Errorf("append record: %w", err))
	}
}

func marshalTokenTrackingAuditEntry(entry TokenTrackingAuditEntry) ([]byte, error) {
	entry = normalizeTokenTrackingAuditEntry(entry)
	line, err := json.Marshal(entry)
	if err != nil {
		return nil, fmt.Errorf("encode record: %w", err)
	}
	line = append(line, '\n')
	if len(line) > tokenTrackingAuditMaxLineBytes {
		return nil, fmt.Errorf("encoded record exceeds %d-byte limit", tokenTrackingAuditMaxLineBytes)
	}
	return line, nil
}

func normalizeTokenTrackingAuditEntry(entry TokenTrackingAuditEntry) TokenTrackingAuditEntry {
	entry.SchemaVersion = tokenTrackingAuditSchemaVersion
	if entry.Timestamp.IsZero() {
		entry.Timestamp = time.Now().UTC()
	} else {
		entry.Timestamp = entry.Timestamp.UTC()
	}
	entry.Event, _ = boundedTokenTrackingAuditString(entry.Event, tokenTrackingAuditMaxEventBytes)
	entry.Runtime, _ = boundedTokenTrackingAuditString(entry.Runtime, tokenTrackingAuditMaxRuntimeBytes)
	entry.ToolName, _ = boundedTokenTrackingAuditString(entry.ToolName, tokenTrackingAuditMaxToolNameBytes)
	var truncated bool
	entry.RawCommand, truncated = boundedTokenTrackingAuditString(
		entry.RawCommand,
		tokenTrackingAuditMaxRawCommandBytes,
	)
	entry.RawCommandTruncated = entry.RawCommandTruncated || truncated
	entry.SourceKind, _ = boundedTokenTrackingAuditString(
		entry.SourceKind,
		tokenTrackingAuditMaxSourceKindBytes,
	)
	entry.SourceName, _ = boundedTokenTrackingAuditString(
		entry.SourceName,
		tokenTrackingAuditMaxSourceNameBytes,
	)
	entry.Decision, _ = boundedTokenTrackingAuditString(
		entry.Decision,
		tokenTrackingAuditMaxDecisionBytes,
	)
	candidateCount := min(len(entry.Candidates), tokenTrackingAuditMaxCandidates)
	candidates := make([]TokenTrackingAuditSource, candidateCount)
	for i := range candidateCount {
		candidates[i] = entry.Candidates[i]
	}
	entry.Candidates = candidates
	for i := range candidates {
		candidates[i].SourceKind, _ = boundedTokenTrackingAuditString(
			candidates[i].SourceKind,
			tokenTrackingAuditMaxSourceKindBytes,
		)
		candidates[i].SourceName, _ = boundedTokenTrackingAuditString(
			candidates[i].SourceName,
			tokenTrackingAuditMaxSourceNameBytes,
		)
	}
	allocationCount := min(len(entry.Allocations), tokenTrackingAuditMaxAllocations)
	allocations := make([]TokenTrackingAuditAllocation, allocationCount)
	copy(allocations, entry.Allocations[:allocationCount])
	entry.Allocations = allocations
	for i := range allocations {
		allocations[i].SourceKind, _ = boundedTokenTrackingAuditString(
			allocations[i].SourceKind,
			tokenTrackingAuditMaxSourceKindBytes,
		)
		allocations[i].SourceName, _ = boundedTokenTrackingAuditString(
			allocations[i].SourceName,
			tokenTrackingAuditMaxSourceNameBytes,
		)
	}
	return entry
}

func boundedTokenTrackingAuditString(value string, maxBytes int) (string, bool) {
	changed := !utf8.ValidString(value)
	value = strings.ToValidUTF8(value, "\uFFFD")
	if len(value) <= maxBytes {
		return value, changed
	}
	end := maxBytes
	for end > 0 && !utf8.RuneStart(value[end]) {
		end--
	}
	return value[:end], true
}

func (writer *tokenTrackingAuditWriter) rotateLocked() error {
	if err := writer.file.Close(); err != nil {
		writer.file = nil
		return fmt.Errorf("close active file for rotation: %w", err)
	}
	writer.file = nil

	var rotationErr error
	if err := os.Remove(fmt.Sprintf("%s.%d", writer.path, tokenTrackingAuditMaxBackups)); err != nil &&
		!errors.Is(err, os.ErrNotExist) {
		rotationErr = fmt.Errorf("remove oldest backup: %w", err)
	}
	for i := tokenTrackingAuditMaxBackups - 1; rotationErr == nil && i >= 1; i-- {
		if err := os.Rename(
			fmt.Sprintf("%s.%d", writer.path, i),
			fmt.Sprintf("%s.%d", writer.path, i+1),
		); err != nil &&
			!errors.Is(err, os.ErrNotExist) {
			rotationErr = fmt.Errorf("shift backup %d: %w", i, err)
		}
	}
	if rotationErr == nil {
		if err := os.Rename(writer.path, fmt.Sprintf("%s.%d", writer.path, 1)); err != nil {
			rotationErr = fmt.Errorf("rotate active file: %w", err)
		}
	}

	file, openErr := openTokenTrackingAuditFile(writer.path)
	if openErr == nil {
		writer.file = file
	}
	if rotationErr != nil && openErr != nil {
		return errors.Join(rotationErr, openErr)
	}
	if rotationErr != nil {
		return rotationErr
	}
	return openErr
}

func (writer *tokenTrackingAuditWriter) warnWriteError(err error) {
	writer.writeErrorOnce.Do(func() {
		slog.Warn("failed to write token tracking audit", "path", writer.path, "error", err)
	})
}

func (writer *tokenTrackingAuditWriter) Close() error {
	if writer == nil {
		return nil
	}
	writer.mu.Lock()
	defer writer.mu.Unlock()
	if writer.file == nil {
		return nil
	}
	err := writer.file.Close()
	writer.file = nil
	return err
}