package model

import (
	"errors"
	"fmt"
	"time"

	"github.com/QuantumNous/new-api/common"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// RelayAttempt is the durable financial fact and log outbox for one upstream
// generation. It contains numeric usage and pricing, never a prompt or key.
type RelayAttempt struct {
	ID                    int64  `json:"id"`
	AttemptKey            string `json:"attempt_key" gorm:"type:varchar(96);uniqueIndex"`
	RequestID             string `json:"request_id" gorm:"type:varchar(64);index"`
	Number                int    `json:"number"`
	UserID                int    `json:"-" gorm:"index"`
	TokenID               int    `json:"-"`
	ChannelID             int    `json:"channel_id"`
	GroupName             string `json:"group"`
	RequestedModel        string `json:"requested_model"`
	ReturnedModel         string `json:"returned_model"`
	GuardResult           string `json:"guard_result"`
	Funding               string `json:"billing_source"`
	SubscriptionID        int    `json:"-"`
	SubscriptionResetTime int64  `json:"-"`
	Reserved              int    `json:"-"`
	Quota                 int    `json:"quota"`
	ObservedQuota         int    `json:"-"`
	EstimateQuota         int    `json:"-"`
	Billable              bool   `json:"-"`
	UsageSource           string `json:"usage_source"`
	State                 string `json:"state" gorm:"type:varchar(24);index"`
	UsageJSON             string `json:"-" gorm:"type:text"`
	PriceJSON             string `json:"-" gorm:"type:text"`
	LogJSON               string `json:"-" gorm:"type:text"`
	LogDelivered          bool   `json:"-" gorm:"index"`
	CreatedAt             int64  `json:"created_at" gorm:"autoCreateTime:false"`
	UpdatedAt             int64  `json:"-" gorm:"autoUpdateTime:false"`
	Deadline              int64  `json:"-" gorm:"index"`
}

// SQLite cannot upgrade a read transaction while another attempt is writing.
// Retry the entire rolled-back transaction, never an individual debit statement.
func relayAttemptTransaction(fn func(*gorm.DB) error) error {
	for retry := 0; ; retry++ {
		err := DB.Transaction(fn)
		var coded interface{ Code() int }
		busy := errors.As(err, &coded) && (coded.Code()&255 == 5 || coded.Code()&255 == 6)
		if !busy || retry >= 24 {
			return err
		}
		delay := time.Duration(retry+1) * 5 * time.Millisecond
		if delay > 100*time.Millisecond {
			delay = 100 * time.Millisecond
		}
		time.Sleep(delay)
	}
}

var ErrAttemptQuota = errors.New("insufficient quota for another upstream attempt")

func ReserveRelayAttempt(attempt *RelayAttempt, preference string) error {
	if attempt.Reserved < 0 || attempt.Reserved > common.MaxWalletQuota {
		return ErrAttemptQuota
	}
	attempt.CreatedAt = time.Now().Unix()
	attempt.UpdatedAt = attempt.CreatedAt
	attempt.State = "reserved"
	var tokenKey string
	original := *attempt
	err := relayAttemptTransaction(func(tx *gorm.DB) error {
		*attempt = original
		var existing RelayAttempt
		if err := tx.Where("attempt_key = ?", attempt.AttemptKey).First(&existing).Error; err == nil {
			*attempt = existing
			return nil
		} else if !errors.Is(err, gorm.ErrRecordNotFound) {
			return err
		}
		var user User
		if err := lockForUpdate(tx).First(&user, attempt.UserID).Error; err != nil {
			return err
		}
		var token Token
		if err := lockForUpdate(tx).First(&token, attempt.TokenID).Error; err != nil {
			return err
		}
		tokenKey = token.Key
		if token.UserId != attempt.UserID || token.Status != common.TokenStatusEnabled || (!token.UnlimitedQuota && token.RemainQuota < attempt.Reserved) {
			return ErrAttemptQuota
		}
		if err := invalidateTokenCacheForMutation(token.Key); err != nil {
			return err
		}
		amount := int64(attempt.Reserved)
		wallet := func() bool { return user.Quota >= attempt.Reserved }
		trySub := func() (bool, bool, error) {
			var subs []UserSubscription
			if err := lockForUpdate(tx).Where("user_id = ? AND status = ? AND end_time > ?", attempt.UserID, "active", attempt.CreatedAt).Order("end_time asc, id asc").Find(&subs).Error; err != nil {
				return false, false, err
			}
			allow := len(subs) == 0
			for _, sub := range subs {
				allow = allow || sub.AllowWalletOverflow
				plan, err := getSubscriptionPlanByIdTx(tx, sub.PlanId)
				if err != nil {
					return false, allow, err
				}
				if err := maybeResetUserSubscriptionWithPlanTx(tx, &sub, plan, attempt.CreatedAt); err != nil {
					return false, allow, err
				}
				if sub.AmountTotal > 0 && sub.AmountTotal-sub.AmountUsed < amount {
					continue
				}
				if err := tx.Model(&UserSubscription{}).Where("id = ?", sub.Id).Update("amount_used", gorm.Expr("amount_used + ?", amount)).Error; err != nil {
					return false, allow, err
				}
				attempt.Funding = "subscription"
				attempt.SubscriptionID = sub.Id
				attempt.SubscriptionResetTime = sub.LastResetTime
				return true, allow, nil
			}
			return false, allow, nil
		}
		useWallet := false
		switch common.NormalizeBillingPreference(preference) {
		case "wallet_only":
			useWallet = wallet()
		case "wallet_first":
			useWallet = wallet()
			if !useWallet {
				ok, _, err := trySub()
				if err != nil {
					return err
				}
				if !ok {
					return ErrAttemptQuota
				}
			}
		case "subscription_only":
			ok, _, err := trySub()
			if err != nil {
				return err
			}
			if !ok {
				return ErrAttemptQuota
			}
		default:
			ok, allow, err := trySub()
			if err != nil {
				return err
			}
			if !ok {
				useWallet = allow && wallet()
			}
		}
		if useWallet {
			attempt.Funding = "wallet"
			if err := tx.Model(&User{}).Where("id = ?", user.Id).Update("quota", gorm.Expr("quota - ?", attempt.Reserved)).Error; err != nil {
				return err
			}
		} else if attempt.Funding != "subscription" {
			return ErrAttemptQuota
		}
		if err := tx.Model(&Token{}).Where("id = ?", token.Id).Updates(map[string]any{"remain_quota": gorm.Expr("remain_quota - ?", attempt.Reserved), "used_quota": gorm.Expr("used_quota + ?", attempt.Reserved), "accessed_time": attempt.CreatedAt}).Error; err != nil {
			return err
		}
		return tx.Create(attempt).Error
	})
	if err == nil {
		_ = invalidateUserCache(attempt.UserID)
		_ = invalidateTokenCacheForMutation(tokenKey)
	}
	return err
}

func CheckpointRelayAttempt(id int64, quota int, usage, logJSON, returned, guard, source string, billable bool, deadline int64) error {
	return DB.Model(&RelayAttempt{}).Where("id = ? AND state <> ?", id, "settled").Updates(map[string]any{"observed_quota": quota, "usage_json": usage, "log_json": logJSON, "returned_model": returned, "guard_result": guard, "usage_source": source, "billable": billable, "state": "collecting", "updated_at": time.Now().Unix(), "deadline": deadline}).Error
}

// Settlement and all balances/statistics commit together. Retrying after any
// process crash cannot repeat a debit, including the delta-zero case.
func SettleRelayAttempt(id int64, quota int, log *Log, details ...RelayAttempt) error {
	if quota < 0 || quota > common.MaxWalletQuota {
		return fmt.Errorf("invalid attempt charge")
	}
	var attempt RelayAttempt
	var tokenKey string
	err := relayAttemptTransaction(func(tx *gorm.DB) error {
		attempt = RelayAttempt{}
		if err := lockForUpdate(tx).First(&attempt, id).Error; err != nil {
			return err
		}
		if attempt.State == "settled" {
			return nil
		}
		delta := quota - attempt.Reserved
		var user User
		if err := lockForUpdate(tx).Unscoped().First(&user, attempt.UserID).Error; err != nil {
			return err
		}
		var token Token
		if err := lockForUpdate(tx).Unscoped().First(&token, attempt.TokenID).Error; err != nil {
			return err
		}
		tokenKey = token.Key
		if err := invalidateTokenCacheForMutation(token.Key); err != nil {
			return err
		}
		if attempt.Funding == "subscription" {
			if err := tx.Model(&UserSubscription{}).Where("id = ? AND last_reset_time = ?", attempt.SubscriptionID, attempt.SubscriptionResetTime).Update("amount_used", gorm.Expr("CASE WHEN amount_used + ? < 0 THEN 0 ELSE amount_used + ? END", delta, delta)).Error; err != nil {
				return err
			}
		} else if err := tx.Unscoped().Model(&User{}).Where("id = ?", attempt.UserID).Update("quota", gorm.Expr("quota - ?", delta)).Error; err != nil {
			return err
		}
		if err := tx.Unscoped().Model(&Token{}).Where("id = ?", attempt.TokenID).Updates(map[string]any{"remain_quota": gorm.Expr("remain_quota - ?", delta), "used_quota": gorm.Expr("used_quota + ?", delta)}).Error; err != nil {
			return err
		}
		if quota > 0 || attempt.Billable || (len(details) > 0 && details[0].Billable) {
			if err := tx.Unscoped().Model(&User{}).Where("id = ?", attempt.UserID).Updates(map[string]any{"used_quota": gorm.Expr("used_quota + ?", quota), "request_count": gorm.Expr("request_count + 1")}).Error; err != nil {
				return err
			}
			if err := tx.Model(&Channel{}).Where("id = ?", attempt.ChannelID).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error; err != nil {
				return err
			}
		}
		if len(details) > 0 {
			d := details[0]
			attempt.UsageJSON = d.UsageJSON
			attempt.ReturnedModel = d.ReturnedModel
			attempt.GuardResult = d.GuardResult
			attempt.UsageSource = d.UsageSource
			attempt.Billable = d.Billable
			attempt.ObservedQuota = quota
		}
		attempt.Quota = quota
		attempt.State = "settled"
		attempt.UpdatedAt = time.Now().Unix()
		if log != nil {
			log.Quota = quota
			log.AttemptKey = &attempt.AttemptKey
			if common.DataExportEnabled {
				if err := tx.Create(&QuotaData{AttemptKey: &attempt.AttemptKey, UserID: attempt.UserID, Username: log.Username, ModelName: log.ModelName, CreatedAt: log.CreatedAt - (log.CreatedAt % 3600), UseGroup: attempt.GroupName, TokenID: attempt.TokenID, ChannelID: attempt.ChannelID, NodeName: common.NodeName, Count: 1, Quota: quota, TokenUsed: log.PromptTokens + log.CompletionTokens}).Error; err != nil {
					return err
				}
			}
			body, err := common.Marshal(log)
			if err != nil {
				return err
			}
			attempt.LogJSON = string(body)
		}
		return tx.Save(&attempt).Error
	})
	if err == nil {
		_ = invalidateUserCache(attempt.UserID)
		_ = invalidateTokenCacheForMutation(tokenKey)
	}
	return err
}

func DeliverRelayAttemptLogs() error {
	var attempts []RelayAttempt
	if err := DB.Where("state = ? AND log_delivered = ? AND log_json <> ?", "settled", false, "").Limit(100).Find(&attempts).Error; err != nil {
		return err
	}
	for _, a := range attempts {
		var log Log
		if err := common.UnmarshalJsonStr(a.LogJSON, &log); err != nil {
			return err
		}
		log.Id = 0
		log.AttemptKey = &a.AttemptKey
		if err := LOG_DB.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "attempt_key"}}, DoNothing: true}).Create(&log).Error; err != nil {
			return err
		}
		if err := DB.Model(&RelayAttempt{}).Where("id = ?", a.ID).Update("log_delivered", true).Error; err != nil {
			return err
		}
	}
	return nil
}