package usage

import (
	"time"

	"github.com/shopspring/decimal"
	"go.uber.org/zap"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"openteam/server/internal/store"
)

type Record struct {
	RequestID           string
	UserID              int64
	KeyID               int64
	ChannelID           int64
	ModelID             int64
	ModelName           string
	InputTokens         int64
	OutputTokens        int64
	CacheReadTokens     int64
	CacheCreationTokens int64
	InputPrice          decimal.Decimal
	OutputPrice         decimal.Decimal
	CacheReadPrice      decimal.Decimal
	Cost                decimal.Decimal
	LatencyMs           int
	Status              string
	ErrorCode           string
}

type Service struct {
	db  *gorm.DB
	log *zap.Logger
}

func NewService(db *gorm.DB, log *zap.Logger) *Service {
	return &Service{db: db, log: log}
}

// Record inserts a request-level usage log and upserts the daily aggregate.
func (s *Service) Record(r Record) error {
	date := time.Now().Format("2006-01-02")
	err := s.db.Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(&store.UsageLog{
			RequestID:           r.RequestID,
			UserID:              r.UserID,
			KeyID:               r.KeyID,
			ChannelID:           r.ChannelID,
			ModelID:             r.ModelID,
			ModelName:           r.ModelName,
			InputTokens:         r.InputTokens,
			OutputTokens:        r.OutputTokens,
			CacheReadTokens:     r.CacheReadTokens,
			CacheCreationTokens: r.CacheCreationTokens,
			InputPrice:          r.InputPrice,
			OutputPrice:         r.OutputPrice,
			CacheReadPrice:      r.CacheReadPrice,
			Cost:                r.Cost,
			LatencyMs:           r.LatencyMs,
			Status:              r.Status,
			ErrorCode:           r.ErrorCode,
			CreatedAt:           time.Now(),
		}).Error; err != nil {
			return err
		}

		var daily store.UsageDaily
		err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
			Where("user_id = ? AND model_id = ? AND date = ?", r.UserID, r.ModelID, date).
			FirstOrCreate(&daily, store.UsageDaily{
				UserID: r.UserID, ModelID: r.ModelID, Date: date,
			}).Error
		if err != nil {
			return err
		}
		daily.Requests++
		daily.InputTokens += r.InputTokens
		daily.OutputTokens += r.OutputTokens
		daily.CacheReadTokens += r.CacheReadTokens
		daily.CacheCreationTokens += r.CacheCreationTokens
		daily.Cost = daily.Cost.Add(r.Cost)
		return tx.Save(&daily).Error
	})
	if err != nil {
		s.log.Warn("record usage failed", zap.Error(err))
	}
	return err
}