package types

import (
	"sort"
	"time"

	"github.com/netbirdio/netbird/shared/management/http/api"
)

// UsageGranularity is the time-bucket width for the usage overview. New values
// can be added here and handled in bucketStart without touching the store.
type UsageGranularity string

const (
	UsageGranularityDay   UsageGranularity = "day"
	UsageGranularityWeek  UsageGranularity = "week"
	UsageGranularityMonth UsageGranularity = "month"
)

// ParseUsageGranularity maps the API query value to a granularity, defaulting
// to day for empty/unknown input.
func ParseUsageGranularity(s string) UsageGranularity {
	switch UsageGranularity(s) {
	case UsageGranularityWeek:
		return UsageGranularityWeek
	case UsageGranularityMonth:
		return UsageGranularityMonth
	default:
		return UsageGranularityDay
	}
}

// AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is
// the UTC start of the bucket as YYYY-MM-DD.
type AgentNetworkUsageBucket struct {
	PeriodStart          string
	InputTokens          int64
	OutputTokens         int64
	TotalTokens          int64
	CachedInputTokens    int64
	CacheCreationTokens  int64
	InputCostUSD         float64
	CachedInputCostUSD   float64
	CacheCreationCostUSD float64
	OutputCostUSD        float64
}

// TotalCostUSD is the bucket's total spend: the sum of the four per-bucket
// costs. Derived rather than accumulated separately so it cannot disagree with
// the components.
func (b *AgentNetworkUsageBucket) TotalCostUSD() float64 {
	return b.InputCostUSD + b.CachedInputCostUSD + b.CacheCreationCostUSD + b.OutputCostUSD
}

// CacheCostUSD is the bucket's prompt-cache spend: cache reads plus writes.
func (b *AgentNetworkUsageBucket) CacheCostUSD() float64 {
	return b.CachedInputCostUSD + b.CacheCreationCostUSD
}

// ToAPIResponse renders the bucket as the API representation.
func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket {
	return api.AgentNetworkUsageBucket{
		PeriodStart:          b.PeriodStart,
		InputTokens:          b.InputTokens,
		OutputTokens:         b.OutputTokens,
		TotalTokens:          b.TotalTokens,
		CachedInputTokens:    b.CachedInputTokens,
		CacheCreationTokens:  b.CacheCreationTokens,
		InputCostUsd:         b.InputCostUSD,
		CachedInputCostUsd:   b.CachedInputCostUSD,
		CacheCreationCostUsd: b.CacheCreationCostUSD,
		OutputCostUsd:        b.OutputCostUSD,
		CostUsd:              b.TotalCostUSD(),
		CacheCostUsd:         b.CacheCostUSD(),
	}
}

// bucketStart truncates t (in UTC) to the start of its bucket for the given
// granularity. Week buckets start on Monday (ISO week).
func bucketStart(t time.Time, g UsageGranularity) time.Time {
	t = t.UTC()
	switch g {
	case UsageGranularityWeek:
		// Monday-start week. time.Weekday: Sunday=0..Saturday=6.
		offset := (int(t.Weekday()) + 6) % 7
		day := time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)
		return day.AddDate(0, 0, -offset)
	case UsageGranularityMonth:
		return time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, time.UTC)
	default: // day
		return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)
	}
}

// AggregateUsageByGranularity buckets the usage rows by the requested
// granularity and returns the buckets ordered oldest-first. Aggregation is done
// in Go (rather than per-engine SQL date_trunc) so granularities stay portable
// across SQLite/Postgres/MySQL and easy to extend.
func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity) []*AgentNetworkUsageBucket {
	byPeriod := make(map[string]*AgentNetworkUsageBucket)
	for _, r := range rows {
		key := bucketStart(r.Timestamp, g).Format("2006-01-02")
		b := byPeriod[key]
		if b == nil {
			b = &AgentNetworkUsageBucket{PeriodStart: key}
			byPeriod[key] = b
		}
		b.InputTokens += r.InputTokens
		b.OutputTokens += r.OutputTokens
		b.TotalTokens += r.TotalTokens
		b.CachedInputTokens += r.CachedInputTokens
		b.CacheCreationTokens += r.CacheCreationTokens
		b.InputCostUSD += r.InputCostUSD
		b.CachedInputCostUSD += r.CachedInputCostUSD
		b.CacheCreationCostUSD += r.CacheCreationCostUSD
		b.OutputCostUSD += r.OutputCostUSD
	}

	out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))
	for _, b := range byPeriod {
		out = append(out, b)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].PeriodStart < out[j].PeriodStart })
	return out
}
