package api

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

	"github.com/muyue/muyue/internal/config"
)

type consumptionEntry struct {
	Date    string `json:"date"`
	Tokens  int    `json:"tokens"`
	Requests int   `json:"requests"`
}

type providerConsumption struct {
	Name     string             `json:"name"`
	Daily    []consumptionEntry `json:"daily"`
	Total    int                `json:"total_tokens"`
	Requests int                `json:"total_requests"`
}

type consumptionStore struct {
	mu        sync.Mutex
	providers map[string]*providerConsumption
}

func newConsumptionStore() *consumptionStore {
	cs := &consumptionStore{
		providers: make(map[string]*providerConsumption),
	}
	cs.load()
	cs.prune()
	return cs
}

func (cs *consumptionStore) Record(providerName string, tokens int) {
	if tokens <= 0 || providerName == "" {
		return
	}
	cs.mu.Lock()
	defer cs.mu.Unlock()

	today := time.Now().UTC().Format("2006-01-02")

	p, ok := cs.providers[providerName]
	if !ok {
		p = &providerConsumption{Name: providerName}
		cs.providers[providerName] = p
	}

	p.Total += tokens
	p.Requests++

	if len(p.Daily) > 0 && p.Daily[len(p.Daily)-1].Date == today {
		p.Daily[len(p.Daily)-1].Tokens += tokens
		p.Daily[len(p.Daily)-1].Requests++
	} else {
		p.Daily = append(p.Daily, consumptionEntry{
			Date:     today,
			Tokens:   tokens,
			Requests: 1,
		})
	}

	cs.save()
}

func (cs *consumptionStore) GetAll() map[string]*providerConsumption {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	result := make(map[string]*providerConsumption)
	for k, v := range cs.providers {
		pc := *v
		daily := make([]consumptionEntry, len(v.Daily))
		copy(daily, v.Daily)
		pc.Daily = daily
		result[k] = &pc
	}
	return result
}

func (cs *consumptionStore) prune() {
	cutoff := time.Now().UTC().AddDate(0, 0, -7).Format("2006-01-02")
	for _, p := range cs.providers {
		filtered := make([]consumptionEntry, 0)
		for _, d := range p.Daily {
			if d.Date >= cutoff {
				filtered = append(filtered, d)
			}
		}
		p.Daily = filtered
	}
}

func (cs *consumptionStore) filePath() string {
	dir, err := config.ConfigDir()
	if err != nil {
		return ""
	}
	return filepath.Join(dir, "consumption.json")
}

func (cs *consumptionStore) load() {
	fp := cs.filePath()
	if fp == "" {
		return
	}
	data, err := os.ReadFile(fp)
	if err != nil {
		return
	}
	json.Unmarshal(data, &cs.providers)
}

func (cs *consumptionStore) save() {
	fp := cs.filePath()
	if fp == "" {
		return
	}
	data, _ := json.Marshal(cs.providers)
	os.WriteFile(fp, data, 0644)
}