sub2api/backend/internal/service/kimi_oauth_service.go

package service

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

	"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
	"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
)

const (
	kimiAuthBaseURL       = "https://auth.kimi.com"
	kimiDeviceAuthURL     = kimiAuthBaseURL + "/api/oauth/device_authorization"
	kimiTokenURL          = kimiAuthBaseURL + "/api/oauth/token"
	kimiClientID          = "17e5f671-d194-4dfb-9706-5516cb48c098"
	kimiDefaultScope      = "kimi-code"
	KimiTokenSafetyWindow = 300 // 5 minutes
	KimiTokenMinTTL       = 30  // 30 seconds
)

// KimiDeviceAuthResponse represents the response from the device authorization endpoint.
type KimiDeviceAuthResponse struct {
	DeviceCode            string `json:"device_code"`
	UserCode              string `json:"user_code"`
	VerificationURI       string `json:"verification_uri"`
	VerificationURIComplete string `json:"verification_uri_complete"`
	ExpiresIn             int64  `json:"expires_in"`
	Interval              int64  `json:"interval"`
}

// KimiTokenResponse represents the response from the token endpoint.
type KimiTokenResponse struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	ExpiresIn    int64  `json:"expires_in"`
	Scope        string `json:"scope"`
	TokenType    string `json:"token_type"`
}

// KimiTokenInfo holds token information for an account.
type KimiTokenInfo struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	ExpiresIn    int64  `json:"expires_in"`
	ExpiresAt    int64  `json:"expires_at"`
	Scope        string `json:"scope,omitempty"`
	TokenType    string `json:"token_type,omitempty"`
}

// KimiOAuthService handles Kimi OAuth device flow and token refresh.
type KimiOAuthService struct {
	sessions sync.Map // device_code -> *kimiOAuthSession
}

type kimiOAuthSession struct {
	DeviceCode              string
	UserCode                string
	VerificationURI         string
	VerificationURIComplete string
	ExpiresAt               time.Time
	Interval                time.Duration
	ProxyURL                string
}

// NewKimiOAuthService creates a new KimiOAuthService.
func NewKimiOAuthService() *KimiOAuthService {
	return &KimiOAuthService{}
}

// InitiateDeviceAuth starts the OAuth 2.0 Device Authorization Grant flow.
func (s *KimiOAuthService) InitiateDeviceAuth(ctx context.Context, proxyID *int64, proxyRepo ProxyRepository) (*KimiDeviceAuthResponse, error) {
	data := url.Values{}
	data.Set("client_id", kimiClientID)
	data.Set("scope", kimiDefaultScope)

	var proxyURL string
	if proxyID != nil && proxyRepo != nil {
		if proxy, err := proxyRepo.GetByID(ctx, *proxyID); err == nil && proxy != nil {
			proxyURL = proxy.URL()
		}
	}

	req, err := http.NewRequestWithContext(ctx, "POST", kimiDeviceAuthURL, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create device auth request: %w", err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")

	resp, err := s.doRequest(req, proxyURL)
	if err != nil {
		return nil, fmt.Errorf("device auth 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("failed to read device auth response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("device auth failed: status=%d body=%s", resp.StatusCode, string(body))
	}

	var result KimiDeviceAuthResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, fmt.Errorf("failed to parse device auth response: %w", err)
	}

	interval := time.Duration(result.Interval) * time.Second
	if interval <= 0 {
		interval = 5 * time.Second
	}

	session := &kimiOAuthSession{
		DeviceCode:              result.DeviceCode,
		UserCode:                result.UserCode,
		VerificationURI:         result.VerificationURI,
		VerificationURIComplete: result.VerificationURIComplete,
		ExpiresAt:               time.Now().Add(time.Duration(result.ExpiresIn) * time.Second),
		Interval:                interval,
		ProxyURL:                proxyURL,
	}
	s.sessions.Store(result.DeviceCode, session)

	logger.LegacyPrintf("service.kimi_oauth", "[KimiOAuth] Device auth initiated: user_code=%s", result.UserCode)
	return &result, nil
}

// PollToken polls the token endpoint for a device_code.
func (s *KimiOAuthService) PollToken(ctx context.Context, deviceCode string) (*KimiTokenInfo, error) {
	sessionVal, ok := s.sessions.Load(deviceCode)
	if !ok {
		return nil, fmt.Errorf("session not found or expired")
	}
	session := sessionVal.(*kimiOAuthSession)

	if time.Now().After(session.ExpiresAt) {
		s.sessions.Delete(deviceCode)
		return nil, fmt.Errorf("device auth session expired")
	}

	data := url.Values{}
	data.Set("client_id", kimiClientID)
	data.Set("grant_type", "urn:ietf:params:oauth:grant-type:device_code")
	data.Set("device_code", deviceCode)

	req, err := http.NewRequestWithContext(ctx, "POST", kimiTokenURL, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create token request: %w", err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")

	resp, err := s.doRequest(req, session.ProxyURL)
	if err != nil {
		return nil, fmt.Errorf("token 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("failed to read token response: %w", err)
	}

	if resp.StatusCode == http.StatusBadRequest || resp.StatusCode == http.StatusForbidden {
		// Check for authorization_pending
		var errResp struct {
			Error string `json:"error"`
		}
		if json.Unmarshal(body, &errResp) == nil && errResp.Error == "authorization_pending" {
			return nil, fmt.Errorf("authorization_pending")
		}
		if errResp.Error == "expired_token" || errResp.Error == "invalid_grant" {
			s.sessions.Delete(deviceCode)
		}
		return nil, fmt.Errorf("token poll failed: %s", errResp.Error)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("token poll failed: status=%d body=%s", resp.StatusCode, string(body))
	}

	var tokenResp KimiTokenResponse
	if err := json.Unmarshal(body, &tokenResp); err != nil {
		return nil, fmt.Errorf("failed to parse token response: %w", err)
	}

	s.sessions.Delete(deviceCode)

	expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - KimiTokenSafetyWindow
	minExpiresAt := time.Now().Unix() + KimiTokenMinTTL
	if expiresAt < minExpiresAt {
		expiresAt = minExpiresAt
	}

	logger.LegacyPrintf("service.kimi_oauth", "[KimiOAuth] Token obtained successfully")
	return &KimiTokenInfo{
		AccessToken:  tokenResp.AccessToken,
		RefreshToken: tokenResp.RefreshToken,
		ExpiresIn:    tokenResp.ExpiresIn,
		ExpiresAt:    expiresAt,
		Scope:        tokenResp.Scope,
		TokenType:    tokenResp.TokenType,
	}, nil
}

// RefreshToken refreshes the access token using a refresh token.
func (s *KimiOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*KimiTokenInfo, error) {
	if strings.TrimSpace(refreshToken) == "" {
		return nil, fmt.Errorf("refresh_token is empty")
	}

	data := url.Values{}
	data.Set("client_id", kimiClientID)
	data.Set("grant_type", "refresh_token")
	data.Set("refresh_token", refreshToken)

	req, err := http.NewRequestWithContext(ctx, "POST", kimiTokenURL, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create refresh request: %w", err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")

	resp, err := s.doRequest(req, proxyURL)
	if err != nil {
		return nil, fmt.Errorf("refresh 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("failed to read refresh response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("refresh failed: status=%d body=%s", resp.StatusCode, string(body))
	}

	var tokenResp KimiTokenResponse
	if err := json.Unmarshal(body, &tokenResp); err != nil {
		return nil, fmt.Errorf("failed to parse refresh response: %w", err)
	}

	expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - KimiTokenSafetyWindow
	minExpiresAt := time.Now().Unix() + KimiTokenMinTTL
	if expiresAt < minExpiresAt {
		expiresAt = minExpiresAt
	}

	return &KimiTokenInfo{
		AccessToken:  tokenResp.AccessToken,
		RefreshToken: tokenResp.RefreshToken,
		ExpiresIn:    tokenResp.ExpiresIn,
		ExpiresAt:    expiresAt,
		Scope:        tokenResp.Scope,
		TokenType:    tokenResp.TokenType,
	}, nil
}

// BuildAccountCredentials builds the credentials map for a Kimi account.
func (s *KimiOAuthService) BuildAccountCredentials(tokenInfo *KimiTokenInfo) map[string]any {
	creds := map[string]any{
		"access_token": tokenInfo.AccessToken,
		"expires_at":   strconv.FormatInt(tokenInfo.ExpiresAt, 10),
	}
	if tokenInfo.RefreshToken != "" {
		creds["refresh_token"] = tokenInfo.RefreshToken
	}
	if tokenInfo.TokenType != "" {
		creds["token_type"] = tokenInfo.TokenType
	}
	if tokenInfo.Scope != "" {
		creds["scope"] = tokenInfo.Scope
	}
	return creds
}

// Stop cleans up the service.
func (s *KimiOAuthService) Stop() {
	s.sessions.Range(func(key, value any) bool {
		s.sessions.Delete(key)
		return true
	})
}

func (s *KimiOAuthService) doRequest(req *http.Request, proxyURL string) (*http.Response, error) {
	opts := httpclient.Options{
		ProxyURL: proxyURL,
		Timeout:  30 * time.Second,
	}
	client, err := httpclient.GetClient(opts)
	if err != nil {
		return nil, err
	}
	return client.Do(req)
}