sub2api/backend/internal/service/kiro_usage_fetcher.go

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
	kiropkg "github.com/Wei-Shaw/sub2api/internal/pkg/kiro"
	"github.com/google/uuid"
)

const (
	kiroUsageOrigin       = "AI_EDITOR"
	kiroUsageResourceType = "AGENTIC_REQUEST"

	kiroDefaultRegion = "us-east-1"
)

var resolveKiroRuntimeEndpoint = kiroRuntimeEndpoint

type kiroUsageLimitsResponse struct {
	NextDateReset        any                      `json:"nextDateReset"`
	OverageConfiguration kiroOverageConfiguration `json:"overageConfiguration"`
	SubscriptionInfo     kiroSubscriptionInfo     `json:"subscriptionInfo"`
	UsageBreakdownList   []kiroUsageBreakdown     `json:"usageBreakdownList"`
}

type kiroOverageConfiguration struct {
	OverageStatus string `json:"overageStatus"`
}

type kiroSubscriptionInfo struct {
	SubscriptionTitle string `json:"subscriptionTitle"`
	Type              string `json:"type"`
}

type kiroUsageBreakdown struct {
	Currency                     string             `json:"currency"`
	CurrentOverages              *float64           `json:"currentOverages"`
	CurrentOveragesWithPrecision *float64           `json:"currentOveragesWithPrecision"`
	CurrentUsage                 *float64           `json:"currentUsage"`
	CurrentUsageWithPrecision    *float64           `json:"currentUsageWithPrecision"`
	DisplayName                  string             `json:"displayName"`
	DisplayNamePlural            string             `json:"displayNamePlural"`
	FreeTrialInfo                *kiroFreeTrialInfo `json:"freeTrialInfo"`
	NextDateReset                any                `json:"nextDateReset"`
	OverageCharges               *float64           `json:"overageCharges"`
	ResourceType                 string             `json:"resourceType"`
	UsageLimit                   *float64           `json:"usageLimit"`
	UsageLimitWithPrecision      *float64           `json:"usageLimitWithPrecision"`
}

type kiroFreeTrialInfo struct {
	CurrentUsage              *float64 `json:"currentUsage"`
	CurrentUsageWithPrecision *float64 `json:"currentUsageWithPrecision"`
	FreeTrialExpiry           any      `json:"freeTrialExpiry"`
	FreeTrialStatus           string   `json:"freeTrialStatus"`
	UsageLimit                *float64 `json:"usageLimit"`
	UsageLimitWithPrecision   *float64 `json:"usageLimitWithPrecision"`
}

type kiroUsageHTTPError struct {
	StatusCode int
	Body       string
}

func (e *kiroUsageHTTPError) Error() string {
	if e == nil {
		return "kiro usage request failed"
	}
	if strings.TrimSpace(e.Body) == "" {
		return fmt.Sprintf("kiro usage request failed (status %d)", e.StatusCode)
	}
	return fmt.Sprintf("kiro usage request failed (status %d): %s", e.StatusCode, e.Body)
}

func (s *AccountUsageService) getKiroUsage(ctx context.Context, account *Account, source string, forceRefresh bool) (*UsageInfo, error) {
	now := time.Now()
	if account == nil {
		return &UsageInfo{
			Source:    source,
			UpdatedAt: &now,
			Error:     "account is nil",
			ErrorCode: errorCodeNetworkError,
		}, nil
	}
	if account.Platform != PlatformKiro || account.Type != AccountTypeOAuth {
		return &UsageInfo{
			Source:    source,
			UpdatedAt: &now,
		}, nil
	}

	cached, hasCached := s.getCachedKiroUsage(account.ID)
	if hasCached && (cached.ErrorCode != "" || cached.Error != "") {
		cached.Source = source
		s.attachKiroRuntimeState(ctx, account, cached)
		return cached, nil
	}
	if !forceRefresh && hasCached {
		cached.Source = source
		s.attachKiroRuntimeState(ctx, account, cached)
		return cached, nil
	}

	flightKey := fmt.Sprintf("kiro-usage:%d", account.ID)
	result, fetchErr, _ := s.cache.kiroUsageFlight.Do(flightKey, func() (any, error) {
		if !forceRefresh {
			if usage, ok := s.getCachedKiroUsage(account.ID); ok {
				return usage, nil
			}
		}
		usage, err := s.fetchAndCacheKiroUsage(ctx, account, source)
		if err != nil {
			return nil, err
		}
		return usage, nil
	})
	if fetchErr == nil {
		if usage, ok := result.(*UsageInfo); ok && usage != nil {
			usage.Source = source
			s.attachKiroRuntimeState(ctx, account, usage)
			if source == "active" {
				s.tryClearRecoverableAccountError(ctx, account)
			}
			return usage, nil
		}
	}

	degraded := buildKiroDegradedUsage(fetchErr)
	degraded.Source = source
	if hasCached {
		cached.Error = degraded.Error
		cached.ErrorCode = degraded.ErrorCode
		cached.NeedsReauth = degraded.NeedsReauth
		cached.KiroQuotaState = degraded.KiroQuotaState
		cached.KiroQuotaReason = degraded.KiroQuotaReason
		cached.KiroQuotaResetAt = degraded.KiroQuotaResetAt
		cached.Source = source
		s.attachKiroRuntimeState(ctx, account, cached)
		return cached, nil
	}
	s.storeKiroUsageSnapshot(account.ID, degraded)
	s.attachKiroRuntimeState(ctx, account, degraded)
	return degraded, nil
}

func (s *AccountUsageService) fetchAndCacheKiroUsage(ctx context.Context, account *Account, source string) (*UsageInfo, error) {
	token := strings.TrimSpace(account.GetCredential("access_token"))
	if token == "" {
		return nil, fmt.Errorf("no access token available")
	}

	region := kiroAPIRegion(account)
	profileArn := strings.TrimSpace(account.GetCredential("profile_arn"))

	resp, err := s.requestKiroUsageLimits(ctx, account, region, profileArn, token)
	if err != nil {
		return nil, err
	}

	usage := mapKiroUsageToInfo(resp)
	usage.Source = source
	s.storeKiroUsageSnapshot(account.ID, usage)
	return usage, nil
}

func (s *AccountUsageService) storeKiroUsageSnapshot(accountID int64, usage *UsageInfo) {
	if s == nil || s.cache == nil || accountID <= 0 || usage == nil {
		return
	}
	now := time.Now()
	if usage.UpdatedAt == nil {
		usage.UpdatedAt = &now
	}
	s.cache.kiroUsageCache.Store(accountID, &kiroUsageCache{
		usageInfo: cloneUsageInfo(usage),
		timestamp: now,
	})
}

func (s *AccountUsageService) getCachedKiroUsage(accountID int64) (*UsageInfo, bool) {
	if s == nil || s.cache == nil || accountID <= 0 {
		return nil, false
	}
	cached, ok := s.cache.kiroUsageCache.Load(accountID)
	if !ok {
		return nil, false
	}
	cache, ok := cached.(*kiroUsageCache)
	if !ok || cache == nil || cache.usageInfo == nil {
		return nil, false
	}
	if time.Since(cache.timestamp) >= kiroCacheTTL(cache.usageInfo) {
		return nil, false
	}
	return cloneUsageInfo(cache.usageInfo), true
}

func kiroCacheTTL(info *UsageInfo) time.Duration {
	if info == nil {
		return kiroUsageErrorTTL
	}
	if info.ErrorCode != "" || info.Error != "" {
		return kiroUsageErrorTTL
	}
	return apiCacheTTL
}

func cloneUsageInfo(info *UsageInfo) *UsageInfo {
	if info == nil {
		return nil
	}
	cloned := *info
	return &cloned
}

func (s *AccountUsageService) requestKiroUsageLimits(ctx context.Context, account *Account, region, profileArn, token string) (*kiroUsageLimitsResponse, error) {
	endpoint := resolveKiroRuntimeEndpoint(region)
	reqURL, err := url.Parse(endpoint + "/getUsageLimits")
	if err != nil {
		return nil, fmt.Errorf("build kiro usage url failed: %w", err)
	}
	q := reqURL.Query()
	q.Set("origin", kiroUsageOrigin)
	if profileArn = strings.TrimSpace(profileArn); profileArn != "" {
		q.Set("profileArn", profileArn)
	}
	q.Set("resourceType", kiroUsageResourceType)
	reqURL.RawQuery = q.Encode()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL.String(), nil)
	if err != nil {
		return nil, fmt.Errorf("create kiro usage request failed: %w", err)
	}
	s.applyKiroRuntimeHeaders(req, account, token)

	client, err := httpclient.GetClient(httpclient.Options{
		ProxyURL:           accountProxyURL(account),
		Timeout:            30 * time.Second,
		ValidateResolvedIP: true,
		AllowPrivateHosts:  isLoopbackEndpoint(endpoint),
	})
	if err != nil {
		return nil, fmt.Errorf("create kiro usage client failed: %w", err)
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("kiro usage request failed: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return nil, fmt.Errorf("read kiro usage response failed: %w", err)
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, &kiroUsageHTTPError{StatusCode: resp.StatusCode, Body: strings.TrimSpace(string(body))}
	}

	var parsed kiroUsageLimitsResponse
	if err := json.Unmarshal(body, &parsed); err != nil {
		return nil, fmt.Errorf("decode kiro usage response failed: %w", err)
	}
	return &parsed, nil
}

func (s *AccountUsageService) applyKiroRuntimeHeaders(req *http.Request, account *Account, token string) {
	if req == nil {
		return
	}
	accountKey := buildKiroAccountKey(account)
	machineID := buildKiroMachineID(account)
	req.Header.Set("Accept", "*/*")
	req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(token))
	req.Header.Set("User-Agent", kiropkg.BuildRuntimeUserAgent(accountKey, machineID))
	req.Header.Set("X-Amz-User-Agent", kiropkg.BuildRuntimeAmzUserAgent(accountKey, machineID))
	req.Header.Set("x-amzn-kiro-agent-mode", "vibe")
	req.Header.Set("x-amzn-codewhisperer-optout", "true")
	req.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
	req.Header.Set("Amz-Sdk-Invocation-Id", uuid.NewString())

	if account == nil {
		return
	}
	applyKiroConditionalHeaders(req, account)
}

func accountProxyURL(account *Account) string {
	if account == nil || account.ProxyID == nil || account.Proxy == nil {
		return ""
	}
	return account.Proxy.URL()
}

func kiroRuntimeEndpoint(region string) string {
	region = strings.TrimSpace(region)
	if region == "" {
		region = kiroDefaultRegion
	}
	switch region {
	case "us-east-1":
		return "https://q.us-east-1.amazonaws.com"
	case "eu-central-1":
		return "https://q.eu-central-1.amazonaws.com"
	case "us-gov-east-1":
		return "https://q-fips.us-gov-east-1.amazonaws.com"
	case "us-gov-west-1":
		return "https://q-fips.us-gov-west-1.amazonaws.com"
	case "us-iso-east-1":
		return "https://q.us-iso-east-1.c2s.ic.gov"
	case "us-isob-east-1":
		return "https://q.us-isob-east-1.sc2s.sgov.gov"
	case "us-isof-south-1":
		return "https://q.us-isof-south-1.csp.hci.ic.gov"
	case "us-isof-east-1":
		return "https://q.us-isof-east-1.csp.hci.ic.gov"
	default:
		if strings.HasPrefix(region, "us-gov-") {
			return "https://q-fips." + region + ".amazonaws.com"
		}
		return "https://q." + region + ".amazonaws.com"
	}
}

func isLoopbackEndpoint(raw string) bool {
	parsed, err := url.Parse(strings.TrimSpace(raw))
	if err != nil {
		return false
	}
	host := strings.TrimSpace(parsed.Hostname())
	if host == "" {
		return false
	}
	if strings.EqualFold(host, "localhost") {
		return true
	}
	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}

func mapKiroUsageToInfo(resp *kiroUsageLimitsResponse) *UsageInfo {
	now := time.Now()
	if resp == nil {
		return &UsageInfo{UpdatedAt: &now}
	}
	info := &UsageInfo{
		UpdatedAt:            &now,
		KiroSubscriptionName: strings.TrimSpace(resp.SubscriptionInfo.SubscriptionTitle),
		KiroSubscriptionType: strings.TrimSpace(resp.SubscriptionInfo.Type),
		KiroOveragesEnabled:  strings.EqualFold(strings.TrimSpace(resp.OverageConfiguration.OverageStatus), "ENABLED"),
	}

	resetAt := parseKiroTimestamp(resp.NextDateReset)
	if credit := selectKiroCreditBreakdown(resp.UsageBreakdownList); credit != nil {
		if breakdownReset := parseKiroTimestamp(credit.NextDateReset); breakdownReset != nil {
			resetAt = breakdownReset
		}
		info.KiroCredit = &KiroCreditProgress{
			CurrentUsage:   selectKiroFloat(credit.CurrentUsageWithPrecision, credit.CurrentUsage),
			UsageLimit:     selectKiroFloat(credit.UsageLimitWithPrecision, credit.UsageLimit),
			PercentageUsed: percentageOrZero(selectKiroFloat(credit.CurrentUsageWithPrecision, credit.CurrentUsage), selectKiroFloat(credit.UsageLimitWithPrecision, credit.UsageLimit)),
		}
		info.KiroOverage = &KiroOverageInfo{
			CurrentOverages: selectKiroFloat(credit.CurrentOveragesWithPrecision, credit.CurrentOverages),
			OverageCharges:  selectKiroFloat(credit.OverageCharges, nil),
			CurrencyCode:    strings.TrimSpace(credit.Currency),
			CurrencySymbol:  kiroCurrencySymbol(strings.TrimSpace(credit.Currency)),
		}
		if ft := credit.FreeTrialInfo; ft != nil && strings.EqualFold(strings.TrimSpace(ft.FreeTrialStatus), "ACTIVE") {
			expiry := parseKiroTimestamp(ft.FreeTrialExpiry)
			daysRemaining := 0
			if expiry != nil {
				daysRemaining = int(time.Until(*expiry).Hours() / 24)
				if time.Until(*expiry)%(24*time.Hour) != 0 {
					daysRemaining++
				}
				if daysRemaining < 0 {
					daysRemaining = 0
				}
			}
			current := selectKiroFloat(ft.CurrentUsageWithPrecision, ft.CurrentUsage)
			limit := selectKiroFloat(ft.UsageLimitWithPrecision, ft.UsageLimit)
			info.KiroBonus = &KiroCreditProgress{
				CurrentUsage:   current,
				UsageLimit:     limit,
				PercentageUsed: percentageOrZero(current, limit),
				DaysRemaining:  daysRemaining,
				ExpiryDate:     expiry,
			}
		}
	}
	info.KiroResetAt = resetAt
	setKiroQuotaStateFromUsage(info)
	return info
}

func selectKiroCreditBreakdown(items []kiroUsageBreakdown) *kiroUsageBreakdown {
	for i := range items {
		if strings.EqualFold(strings.TrimSpace(items[i].ResourceType), "CREDIT") {
			return &items[i]
		}
	}
	if len(items) == 0 {
		return nil
	}
	return &items[0]
}

func selectKiroFloat(preferred *float64, fallback *float64) float64 {
	switch {
	case preferred != nil:
		return *preferred
	case fallback != nil:
		return *fallback
	default:
		return 0
	}
}

func percentageOrZero(current, limit float64) float64 {
	if limit <= 0 {
		return 0
	}
	return current / limit * 100
}

func parseKiroTimestamp(raw any) *time.Time {
	if raw == nil {
		return nil
	}
	switch v := raw.(type) {
	case string:
		trimmed := strings.TrimSpace(v)
		if trimmed == "" {
			return nil
		}
		if parsed, err := time.Parse(time.RFC3339, trimmed); err == nil {
			return &parsed
		}
		if i, err := strconv.ParseInt(trimmed, 10, 64); err == nil {
			return unixishToTime(i)
		}
		if f, err := strconv.ParseFloat(trimmed, 64); err == nil {
			return unixishFloatToTime(f)
		}
	case float64:
		return unixishFloatToTime(v)
	case int64:
		return unixishToTime(v)
	case int:
		return unixishToTime(int64(v))
	case json.Number:
		if i, err := v.Int64(); err == nil {
			return unixishToTime(i)
		}
		if f, err := v.Float64(); err == nil {
			return unixishFloatToTime(f)
		}
	}
	return nil
}

func unixishFloatToTime(v float64) *time.Time {
	if v <= 0 {
		return nil
	}
	if v >= 1e12 {
		t := time.UnixMilli(int64(v))
		return &t
	}
	t := time.Unix(int64(v), 0)
	return &t
}

func unixishToTime(v int64) *time.Time {
	if v <= 0 {
		return nil
	}
	if v >= 1e12 {
		t := time.UnixMilli(v)
		return &t
	}
	t := time.Unix(v, 0)
	return &t
}

func kiroCurrencySymbol(code string) string {
	switch strings.ToUpper(strings.TrimSpace(code)) {
	case "USD":
		return "$"
	default:
		return ""
	}
}

func buildKiroDegradedUsage(err error) *UsageInfo {
	now := time.Now()
	info := &UsageInfo{
		UpdatedAt: &now,
		Error:     "usage API error",
		ErrorCode: errorCodeNetworkError,
	}
	if err == nil {
		return info
	}

	info.Error = fmt.Sprintf("usage API error: %v", err)

	classification := classifyKiroError(err)
	switch classification.Category {
	case kiroErrorAuthError:
		info.ErrorCode = errorCodeUnauthenticated
		info.NeedsReauth = true
	case kiroErrorRateLimited:
		info.ErrorCode = errorCodeRateLimited
	case kiroErrorQuotaExhausted:
		info.ErrorCode = errorCodeNetworkError
		info.KiroQuotaState = kiroQuotaStateCreditsExhausted
		info.KiroQuotaReason = classification.Message
	case kiroErrorOverageExhausted:
		info.ErrorCode = errorCodeNetworkError
		info.KiroQuotaState = kiroQuotaStateOverageExhausted
		info.KiroQuotaReason = classification.Message
	case kiroErrorSuspended, kiroErrorUsageForbidden, kiroErrorProfileError:
		info.ErrorCode = errorCodeForbidden
	default:
		info.ErrorCode = errorCodeNetworkError
	}
	return info
}

func (s *AccountUsageService) attachKiroRuntimeState(ctx context.Context, account *Account, usage *UsageInfo) {
	if s == nil || usage == nil || account == nil || account.Platform != PlatformKiro || s.kiroCooldownStore == nil {
		return
	}
	usage.KiroRuntimeState = ""
	usage.KiroRuntimeReason = ""
	usage.KiroRuntimeResetAt = nil
	state, err := s.kiroCooldownStore.GetState(ctx, buildKiroAccountKey(account))
	if err != nil || state == nil {
		return
	}
	usage.KiroRuntimeState, usage.KiroRuntimeReason, usage.KiroRuntimeResetAt = kiroRuntimeStateSnapshot(state)
}

func (s *AccountUsageService) EnrichAccountWithKiroRuntimeState(ctx context.Context, account *Account) {
	if s == nil || account == nil || account.Platform != PlatformKiro || account.Type != AccountTypeOAuth {
		return
	}
	account.KiroQuotaState = ""
	account.KiroQuotaReason = ""
	account.KiroQuotaResetAt = nil
	account.KiroRuntimeState = ""
	account.KiroRuntimeReason = ""
	account.KiroRuntimeResetAt = nil
	if usage, ok := s.getCachedKiroUsage(account.ID); ok {
		account.KiroQuotaState = usage.KiroQuotaState
		account.KiroQuotaReason = usage.KiroQuotaReason
		account.KiroQuotaResetAt = usage.KiroQuotaResetAt
	}
	if s.kiroCooldownStore == nil {
		return
	}
	state, err := s.kiroCooldownStore.GetState(ctx, buildKiroAccountKey(account))
	if err != nil || state == nil {
		return
	}
	account.KiroRuntimeState, account.KiroRuntimeReason, account.KiroRuntimeResetAt = kiroRuntimeStateSnapshot(state)
}

func setKiroQuotaStateFromUsage(info *UsageInfo) {
	if info == nil {
		return
	}
	info.KiroQuotaState = ""
	info.KiroQuotaReason = ""
	info.KiroQuotaResetAt = nil

	creditExhausted := false
	if info.KiroCredit != nil && info.KiroCredit.UsageLimit > 0 {
		creditExhausted = info.KiroCredit.CurrentUsage >= info.KiroCredit.UsageLimit
	}
	overageActive := info.KiroOverage != nil &&
		(info.KiroOverage.CurrentOverages > 0 || info.KiroOverage.OverageCharges > 0)

	switch {
	case info.KiroOveragesEnabled && (overageActive || creditExhausted):
		info.KiroQuotaState = kiroQuotaStateOverageActive
		info.KiroQuotaReason = "overages_enabled"
		info.KiroQuotaResetAt = info.KiroResetAt
	case creditExhausted:
		info.KiroQuotaState = kiroQuotaStateCreditsExhausted
		info.KiroQuotaReason = "credits_exhausted"
		info.KiroQuotaResetAt = info.KiroResetAt
	default:
		info.KiroQuotaState = kiroQuotaStateNormal
	}
}