package service

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"go.uber.org/zap"

	"github.com/caarlos0/httperr"
	"github.com/tidbcloud/inference-service/internal/config"
)

// UsageData represents the complete usage data structure, not concurrent-safe. Must be protected by a mutex.
type UsageData struct {
	Tenants map[string]*TenantUsage `json:"tenants"`
	Models  map[string]*ModelUsage  `json:"models"`
}

// TenantUsage tracks usage for a specific tenant
type TenantUsage struct {
	TenantID     string    `json:"tenant_id"`
	CurrentUsage int64     `json:"current_usage"` // in micro-dollars (1 dollar = 1,000,000 micro-dollars)
	QuotaLimit   int64     `json:"quota_limit"`   // in micro-dollars
	ResetDate    time.Time `json:"reset_date"`    // when quota resets
	LastUpdated  time.Time `json:"last_updated"`

	// Monthly breakdown
	MonthlyUsage map[string]int64 `json:"monthly_usage"` // key: "2024-01", value: usage in micro-dollars
}

// ModelUsage tracks usage for a specific model
type ModelUsage struct {
	ModelName    string    `json:"model_name"`
	TotalUsage   int64     `json:"total_usage"`   // in micro-dollars
	TotalTokens  int64     `json:"total_tokens"`  // total tokens processed
	RequestCount int64     `json:"request_count"` // total requests
	LastUpdated  time.Time `json:"last_updated"`

	// Daily breakdown for current month
	DailyUsage map[string]int64 `json:"daily_usage"` // key: "2024-01-15", value: usage in micro-dollars
}

// NewUsageData creates a new UsageData instance
func NewUsageData() *UsageData {
	return &UsageData{
		Tenants: make(map[string]*TenantUsage),
		Models:  make(map[string]*ModelUsage),
	}
}

// GetTenantUsage returns usage data for a tenant, creating it if it doesn't exist
// WARN: Could potentially modify the data.
func (ud *UsageData) GetTenantUsage(tenantID string) *TenantUsage {
	if usage, exists := ud.Tenants[tenantID]; exists {
		// Check if quota needs to be reset
		if ud.shouldResetQuota(usage.ResetDate) {
			usage.CurrentUsage = 0
			usage.ResetDate = ud.getNextResetDate()
		}
		return usage
	}

	// Create new tenant usage with default quota from config
	defaultQuotaMicroDollars := config.Global().Quotas.DefaultMonthlyLimit
	usage := &TenantUsage{
		TenantID:     tenantID,
		CurrentUsage: 0,
		QuotaLimit:   defaultQuotaMicroDollars,
		ResetDate:    ud.getNextResetDate(),
		LastUpdated:  time.Now(),
		MonthlyUsage: make(map[string]int64),
	}

	ud.Tenants[tenantID] = usage
	return usage
}

// GetModelUsage returns usage data for a model, creating it if it doesn't exist
// WARN: Could potentially modify the data.
func (ud *UsageData) GetModelUsage(modelName string) *ModelUsage {
	if usage, exists := ud.Models[modelName]; exists {
		return usage
	}

	// Create new model usage
	usage := &ModelUsage{
		ModelName:    modelName,
		TotalUsage:   0,
		TotalTokens:  0,
		RequestCount: 0,
		LastUpdated:  time.Now(),
		DailyUsage:   make(map[string]int64),
	}

	ud.Models[modelName] = usage
	return usage
}

// AddUsage adds usage for a tenant and model
func (ud *UsageData) AddUsage(tenantID, modelName string, costMicroDollars int64, tokens int) {
	now := time.Now()

	// Update tenant usage
	tenantUsage := ud.GetTenantUsage(tenantID)
	tenantUsage.CurrentUsage += costMicroDollars
	tenantUsage.LastUpdated = now

	// Update monthly usage
	monthKey := now.Format("2006-01")
	if tenantUsage.MonthlyUsage == nil {
		tenantUsage.MonthlyUsage = make(map[string]int64)
	}
	tenantUsage.MonthlyUsage[monthKey] += costMicroDollars

	// Update model usage
	modelUsage := ud.GetModelUsage(modelName)
	modelUsage.TotalUsage += costMicroDollars
	modelUsage.TotalTokens += int64(tokens)
	modelUsage.RequestCount++
	modelUsage.LastUpdated = now

	// Update daily usage
	dayKey := now.Format("2006-01-02")
	if modelUsage.DailyUsage == nil {
		modelUsage.DailyUsage = make(map[string]int64)
	}
	modelUsage.DailyUsage[dayKey] += costMicroDollars
}

// CanUseQuota checks if a tenant can use the specified amount
// WARN: Could potentially modify the data.
func (ud *UsageData) CanUseQuota(tenantID string, costMicroDollars int64) bool {
	tenantUsage := ud.GetTenantUsage(tenantID)
	return tenantUsage.CurrentUsage+costMicroDollars <= tenantUsage.QuotaLimit
}

// SetQuotaLimit sets the quota limit for a tenant
func (ud *UsageData) SetQuotaLimit(tenantID string, limitMicroDollars int64) {
	tenantUsage := ud.GetTenantUsage(tenantID)
	tenantUsage.QuotaLimit = limitMicroDollars
}

// shouldResetQuota checks if quota should be reset based on reset date
func (ud *UsageData) shouldResetQuota(resetDate time.Time) bool {
	return time.Now().After(resetDate)
}

// getNextResetDate calculates the next quota reset date (1st of next month)
func (ud *UsageData) getNextResetDate() time.Time {
	now := time.Now()
	year, month, _ := now.Date()

	// Next month, 1st day
	nextMonth := time.Date(year, month+1, 1, 0, 0, 0, 0, now.Location())
	return nextMonth
}

// UnmarshalJSON implements json.Unmarshaler
func (ud *UsageData) UnmarshalJSON(data []byte) error {
	type Alias UsageData
	aux := &struct {
		*Alias
	}{
		Alias: (*Alias)(ud),
	}
	if err := json.Unmarshal(data, aux); err != nil {
		return err
	}
	// Initialize maps if nil
	if ud.Tenants == nil {
		ud.Tenants = make(map[string]*TenantUsage)
	}
	if ud.Models == nil {
		ud.Models = make(map[string]*ModelUsage)
	}
	return nil
}

// UsageService manages usage tracking and quota enforcement
type UsageService struct {
	data    *UsageData
	storage StorageBackend
	dataKey string // key for main data file
	mu      sync.Mutex
	stopCh  chan struct{}

	lastSavedData []byte
}

// NewUsageService creates a new UsageService
func NewUsageService() (*UsageService, error) {
	storageConfig := config.Global().Storage

	var storage StorageBackend
	var err error

	// Check if UsageDataDir is an S3 path
	if strings.HasPrefix(storageConfig.UsageDataDir, "s3://") {
		storage, err = NewS3StorageBackend(storageConfig.UsageDataDir)
		if err != nil {
			return nil, fmt.Errorf("failed to create S3 storage backend: %w", err)
		}
	} else {
		// Create local storage directory if it doesn't exist
		if err := os.MkdirAll(storageConfig.UsageDataDir, 0755); err != nil {
			return nil, fmt.Errorf("failed to create storage directory: %w", err)
		}
		storage = NewLocalStorageBackend(storageConfig.UsageDataDir)
	}

	service := &UsageService{
		storage: storage,
		dataKey: "usage_data.json",
		stopCh:  make(chan struct{}),
	}

	// Load existing data or create new
	if err := service.loadData(); err != nil {
		zap.L().Info("Failed to load usage data, starting fresh", zap.Error(err))
		service.data = NewUsageData()
	}

	// Apply special quota limits from config
	service.applyConfiguredQuotas()

	// Start periodic save
	go service.periodicSave()

	return service, nil
}

// getDailyBackupFile returns the backup file key for a specific date
func (s *UsageService) getDailyBackupFile(date time.Time) string {
	// Create directory prefix like "usage_202504" for April 2025
	monthDir := fmt.Sprintf("usage_%s", date.Format("200601"))
	dateStr := date.Format("2006-01-02")
	return fmt.Sprintf("%s/usage_data_backup_%s.json", monthDir, dateStr)
}

// getTodayBackupFile returns today's backup file key
func (s *UsageService) getTodayBackupFile() string {
	return s.getDailyBackupFile(time.Now())
}

// CheckQuota checks if a tenant can use the specified amount
func (s *UsageService) CheckQuota(tenantID string, costMicroDollars int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.data.CanUseQuota(tenantID, costMicroDollars) {
		tenantUsage := s.data.GetTenantUsage(tenantID)
		return httperr.Errorf(http.StatusForbidden,
			"Exceeded quota limit. Current usage: $%.4f, Limit: $%.4f",
			config.MicroDollarsToDollarsForDisplay(tenantUsage.CurrentUsage),
			config.MicroDollarsToDollarsForDisplay(tenantUsage.QuotaLimit))
	}

	return nil
}

// RecordUsage records usage for a tenant and model
func (s *UsageService) RecordUsage(tenantID, modelName string, costMicroDollars int64, tokens int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.data.AddUsage(tenantID, modelName, costMicroDollars, tokens)

	zap.L().Debug("Recorded usage",
		zap.String("tenant_id", tenantID),
		zap.String("model", modelName),
		zap.Float64("cost_dollars", config.MicroDollarsToDollarsForDisplay(costMicroDollars)),
		zap.Int("tokens", tokens))
}

// GetTenantUsage returns usage information for a specific tenant
func (s *UsageService) GetTenantUsage(tenantID string) *QuotaUsageResponse {
	s.mu.Lock()
	defer s.mu.Unlock()

	tenantUsage := s.data.GetTenantUsage(tenantID)

	return &QuotaUsageResponse{
		TenantID:     tenantUsage.TenantID,
		CurrentUsage: config.MicroDollarsToDollarsForDisplay(tenantUsage.CurrentUsage),
		QuotaLimit:   config.MicroDollarsToDollarsForDisplay(tenantUsage.QuotaLimit),
		ResetDate:    tenantUsage.ResetDate,
	}
}

// GetAllTenantUsage returns usage information for all tenants
func (s *UsageService) GetAllTenantUsage() *AllQuotaUsageResponse {
	s.mu.Lock()
	defer s.mu.Unlock()

	tenants := make([]QuotaUsageResponse, 0, len(s.data.Tenants))
	for _, tenantUsage := range s.data.Tenants {
		tenants = append(tenants, QuotaUsageResponse{
			TenantID:     tenantUsage.TenantID,
			CurrentUsage: config.MicroDollarsToDollarsForDisplay(tenantUsage.CurrentUsage),
			QuotaLimit:   config.MicroDollarsToDollarsForDisplay(tenantUsage.QuotaLimit),
			ResetDate:    tenantUsage.ResetDate,
		})
	}

	return &AllQuotaUsageResponse{
		Tenants: tenants,
	}
}

// GetRawTenantUsageForMetrics returns tenant usage data in the format expected by metrics
func (s *UsageService) GetRawTenantUsageForMetrics() map[string]TenantUsageMetrics {
	s.mu.Lock()
	defer s.mu.Unlock()

	result := make(map[string]TenantUsageMetrics)
	for tenantID, usage := range s.data.Tenants {
		result[tenantID] = TenantUsageMetrics{
			TenantID:     usage.TenantID,
			CurrentUsage: usage.CurrentUsage,
			QuotaLimit:   usage.QuotaLimit,
		}
	}
	return result
}

// GetModelUsageForMetrics returns model usage data in the format expected by metrics
func (s *UsageService) GetModelUsageForMetrics() map[string]ModelUsageMetrics {
	s.mu.Lock()
	defer s.mu.Unlock()

	result := make(map[string]ModelUsageMetrics)
	for modelName, usage := range s.data.Models {
		result[modelName] = ModelUsageMetrics{
			ModelName:  usage.ModelName,
			TotalUsage: usage.TotalUsage,
		}
	}
	return result
}

// TenantUsageMetrics represents tenant usage data for metrics
type TenantUsageMetrics struct {
	TenantID     string
	CurrentUsage int64 // in micro-dollars
	QuotaLimit   int64 // in micro-dollars
}

// ModelUsageMetrics represents model usage data for metrics
type ModelUsageMetrics struct {
	ModelName  string
	TotalUsage int64 // in micro-dollars
}

// Close stops the usage service and saves data
func (s *UsageService) Close() error {
	close(s.stopCh)
	return s.saveData()
}

// loadData loads usage data from storage
func (s *UsageService) loadData() error {
	// Try to load from main data key
	err := s.loadFromKey(s.dataKey)
	if err != nil {
		return fmt.Errorf("failed to load usage data from %s: %w", s.dataKey, err)
	}
	return nil
}

// loadFromKey loads data from a specific storage key
func (s *UsageService) loadFromKey(key string) error {
	data, err := s.storage.Load(key)
	if err != nil {
		return err
	}

	s.data = NewUsageData()
	if err := json.Unmarshal(data, s.data); err != nil {
		return err
	}

	return nil
}

// saveData saves usage data to storage
func (s *UsageService) saveData() error {
	s.mu.Lock()
	data, err := json.MarshalIndent(s.data, "", "  ")
	s.mu.Unlock()
	if err != nil {
		return fmt.Errorf("failed to marshal usage data: %w", err)
	}

	if bytes.Equal(s.lastSavedData, data) {
		zap.L().Debug("Usage data unchanged, skipping save")
		return nil
	}

	todayBackupKey := s.getTodayBackupFile()
	if err := s.storage.Save(todayBackupKey, data); err != nil {
		zap.L().Warn("Failed to create daily backup",
			zap.Error(err),
			zap.String("backup_key", todayBackupKey))
	}

	if err := s.storage.Save(s.dataKey, data); err != nil {
		return fmt.Errorf("failed to save usage data: %w", err)
	}

	s.lastSavedData = data // Store last saved data to avoid unnecessary saves

	zap.L().Info("Usage data saved successfully",
		zap.String("destination", s.dataKey))
	return nil
}

// periodicSave periodically saves usage data
func (s *UsageService) periodicSave() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if err := s.saveData(); err != nil {
				zap.L().Error("Failed to save usage data", zap.Error(err))
			}
		case <-s.stopCh:
			return
		}
	}
}

// applyConfiguredQuotas applies special quota limits from configuration
func (s *UsageService) applyConfiguredQuotas() {
	for tenantID, limitMicroDollars := range config.Global().Quotas.SpecialLimits {
		s.data.SetQuotaLimit(tenantID, limitMicroDollars)
		zap.L().Info("Applied special quota limit",
			zap.String("tenant_id", tenantID),
			zap.Float64("limit_dollars", config.MicroDollarsToDollarsForDisplay(limitMicroDollars)))
	}
}