package model

import (
	"net/http"
	"opencatd-open/store"
	"time"

	"github.com/gin-gonic/gin"
)

type Usage struct {
	ID               int64     `gorm:"column:id;primaryKey;autoIncrement"`
	UserID           int64     `gorm:"column:user_id;index:idx_user_id"`
	TokenID          int64     `gorm:"column:token_id;index:idx_token_id"`
	Capability       string    `gorm:"column:capability;index:idx_usage_capability;comment:模型能力"`
	Date             time.Time `gorm:"column:date;autoCreateTime;index:idx_date"`
	Model            string    `gorm:"column:model"`
	Stream           bool      `gorm:"column:stream"`
	PromptTokens     float64   `gorm:"column:prompt_tokens"`
	CompletionTokens float64   `gorm:"column:completion_tokens"`
	TotalTokens      float64   `gorm:"column:total_tokens"`
	Cost             string    `gorm:"column:cost"`
}

func (Usage) TableName() string {
	return "usages"
}

type DailyUsage struct {
	ID               int64     `gorm:"column:id;primaryKey;autoIncrement"`
	UserID           int64     `gorm:"column:user_id;uniqueIndex:idx_daily_unique,priority:1"`
	TokenID          int64     `gorm:"column:token_id;index:idx_daily_token_id"`
	Capability       string    `gorm:"column:capability;uniqueIndex:idx_daily_unique,priority:2;comment:模型能力"`
	Date             time.Time `gorm:"column:date;autoCreateTime;uniqueIndex:idx_daily_unique,priority:3"`
	Model            string    `gorm:"column:model"`
	Stream           bool      `gorm:"column:stream"`
	PromptTokens     float64   `gorm:"column:prompt_tokens"`
	CompletionTokens float64   `gorm:"column:completion_tokens"`
	TotalTokens      float64   `gorm:"column:total_tokens"`
	Cost             string    `gorm:"column:cost"`
}

func (DailyUsage) TableName() string {
	return "daily_usages"
}

func HandleUsage(c *gin.Context) {
	fromStr := c.Query("from")
	toStr := c.Query("to")
	getMonthStartAndEnd := func() (start, end string) {
		loc, _ := time.LoadLocation("Local")
		now := time.Now().In(loc)

		year, month, _ := now.Date()

		startOfMonth := time.Date(year, month, 1, 0, 0, 0, 0, loc)
		endOfMonth := startOfMonth.AddDate(0, 1, 0)

		start = startOfMonth.Format("2006-01-02")
		end = endOfMonth.Format("2006-01-02")
		return
	}
	if fromStr == "" || toStr == "" {
		fromStr, toStr = getMonthStartAndEnd()
	}

	usage, err := store.QueryUsage(fromStr, toStr)
	if err != nil {
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}

	c.JSON(200, usage)
}