package ai

import (
	"context"

	"github.com/youruser/base/internal/svc"
	"github.com/youruser/base/internal/types"
	"github.com/youruser/base/model"

	"github.com/zeromicro/go-zero/core/logx"
)

type AiUsageStatsLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewAiUsageStatsLogic(ctx context.Context, svcCtx *svc.ServiceContext) *AiUsageStatsLogic {
	return &AiUsageStatsLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
	}
}

func (l *AiUsageStatsLogic) AiUsageStats(req *types.AIUsageStatsRequest) (resp *types.AIUsageStatsResponse, err error) {
	days := req.Days
	if days <= 0 {
		days = 30
	}

	// Total stats
	totalStats, err := model.AIUsageRecordTotalStats(l.ctx, l.svcCtx.DB, days)
	if err != nil {
		return nil, err
	}

	// Model stats
	modelResults, err := model.AIUsageRecordModelStats(l.ctx, l.svcCtx.DB, days)
	if err != nil {
		return nil, err
	}

	modelStats := make([]types.AIModelStatItem, len(modelResults))
	for i, m := range modelResults {
		modelStats[i] = types.AIModelStatItem{
			ModelId:      m.ModelId,
			Calls:        m.Calls,
			InputTokens:  m.InputTokens,
			OutputTokens: m.OutputTokens,
			TotalCost:    m.TotalCost,
		}
	}

	// Daily stats
	dailyResults, err := model.AIUsageRecordDailyStats(l.ctx, l.svcCtx.DB, days)
	if err != nil {
		return nil, err
	}

	dailyStats := make([]types.AIDailyStatItem, len(dailyResults))
	for i, d := range dailyResults {
		dailyStats[i] = types.AIDailyStatItem{
			Date:        d.Date,
			Calls:       d.Calls,
			TotalTokens: d.TotalTokens,
			TotalCost:   d.TotalCost,
		}
	}

	return &types.AIUsageStatsResponse{
		TotalCalls:  totalStats.TotalCalls,
		TotalTokens: totalStats.TotalTokens,
		TotalCost:   totalStats.TotalCost,
		TotalUsers:  totalStats.TotalUsers,
		ModelStats:  modelStats,
		DailyStats:  dailyStats,
	}, nil
}