package codex

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"io"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const (
	oauthClientID    = "app_EMoamEEZ73f0CkXaXp7hrann"
	oauthScope       = "openid profile email"
	oauthTokenURL    = "https://auth.openai.com/oauth/token"
	refreshInterval  = 30 * time.Minute
	expiryBuffer     = 30 * time.Minute
	refreshThreshold = 28 * 24 * time.Hour // 28 days
)

// TokenManager manages ChatGPT OAuth tokens with automatic refresh.
type TokenManager struct {
	mu           sync.RWMutex
	accessToken  string
	refreshToken string
	accountID    string
	idToken      string
	lastRefresh  time.Time
	lastError    string
	rawAuth      map[string]interface{}
	onUpdate     func(map[string]interface{}) error

	httpClient *http.Client
	initialCtx context.Context
	stopCh     chan struct{}
	stopped    bool
}

// TokenManagerOption customizes TokenManager behavior.
type TokenManagerOption func(*TokenManager)

// WithAuthPersistence stores refreshed tokens back into the provided auth.json
// shape and invokes onUpdate with a deep copy after each successful refresh.
func WithAuthPersistence(raw map[string]interface{}, onUpdate func(map[string]interface{}) error) TokenManagerOption {
	return func(tm *TokenManager) {
		tm.rawAuth = cloneMap(raw)
		tm.onUpdate = onUpdate
		if lastRefresh := parseAuthLastRefresh(tm.rawAuth["last_refresh"]); !lastRefresh.IsZero() {
			tm.lastRefresh = lastRefresh
		}
	}
}

func withTokenManagerHTTPClient(httpClient *http.Client) TokenManagerOption {
	return func(tm *TokenManager) {
		if httpClient != nil {
			tm.httpClient = httpClient
		}
	}
}

func withTokenManagerInitialContext(ctx context.Context) TokenManagerOption {
	return func(tm *TokenManager) {
		if ctx != nil {
			tm.initialCtx = ctx
		}
	}
}

// NewTokenManager creates a TokenManager from the initial auth token.
// The authToken can be a JSON string or a plain JWT access token.
func NewTokenManager(authToken string, opts ...TokenManagerOption) (*TokenManager, error) {
	accessToken, refreshToken, idToken, accountID := parseChatGPTAuthToken(authToken)
	if accessToken == "" {
		return nil, fmt.Errorf("no access token found in auth token")
	}

	tm := &TokenManager{
		accessToken:  accessToken,
		refreshToken: refreshToken,
		idToken:      idToken,
		accountID:    accountID,
		httpClient:   &http.Client{Timeout: 30 * time.Second},
		initialCtx:   context.Background(),
		stopCh:       make(chan struct{}),
	}
	for _, opt := range opts {
		opt(tm)
	}

	// Trigger an immediate refresh if Codex auth.json semantics say it is due:
	// expired/near-expiry access token, missing last_refresh, or stale last_refresh.
	if tm.shouldRefreshLocked() {
		ctx, cancel := context.WithTimeout(tm.initialCtx, 30*time.Second)
		defer cancel()
		if err := tm.refreshTokenLocked(ctx); err != nil {
			tm.lastError = err.Error()
			slog.Warn("Initial token refresh failed", "error", err)
		}
	}

	return tm, nil
}

// GetAccessToken returns the current access token and account ID.
func (tm *TokenManager) GetAccessToken() (string, string) {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	return tm.accessToken, tm.accountID
}

func (tm *TokenManager) hasRefreshToken() bool {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	return tm.refreshToken != ""
}

func (tm *TokenManager) tokenSnapshot() (accessToken, refreshToken, accountID string, lastRefresh time.Time, lastError string) {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	return tm.accessToken, tm.refreshToken, tm.accountID, tm.lastRefresh, tm.lastError
}

// StartBackgroundRefresh starts a goroutine that periodically refreshes the token.
func (tm *TokenManager) StartBackgroundRefresh(ctx context.Context) {
	logging.Go(nil, func() {
		ticker := time.NewTicker(refreshInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-tm.stopCh:
				return
			case <-ticker.C:
				if err := tm.doRefresh(ctx); err != nil {
					slog.Warn("Background token refresh failed", "error", err)
				}
			}
		}
	})
}

// Stop signals the background refresh goroutine to stop.
func (tm *TokenManager) Stop() {
	tm.mu.Lock()
	defer tm.mu.Unlock()
	if !tm.stopped {
		close(tm.stopCh)
		tm.stopped = true
	}
}

// ForceRefresh forces an immediate token refresh regardless of expiry.
func (tm *TokenManager) ForceRefresh(ctx context.Context) error {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	if tm.refreshToken == "" {
		err := fmt.Errorf("no refresh token available")
		tm.lastError = err.Error()
		return err
	}

	err := tm.refreshTokenLocked(ctx)
	if err != nil {
		tm.lastError = err.Error()
	} else {
		tm.lastError = ""
	}
	return err
}

func (tm *TokenManager) doRefresh(ctx context.Context) error {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	if !tm.shouldRefreshLocked() {
		return nil
	}

	err := tm.refreshTokenLocked(ctx)
	if err != nil {
		tm.lastError = err.Error()
	} else {
		tm.lastError = ""
	}
	return err
}

func (tm *TokenManager) shouldRefreshLocked() bool {
	if tm.refreshToken == "" {
		return false
	}
	if isAccessTokenExpired(tm.accessToken, expiryBuffer) {
		return true
	}
	if tm.lastRefresh.IsZero() {
		return true
	}
	return time.Since(tm.lastRefresh) >= refreshThreshold
}

func (tm *TokenManager) refreshTokenLocked(ctx context.Context) error {
	payload := map[string]string{
		"client_id":     oauthClientID,
		"grant_type":    "refresh_token",
		"refresh_token": tm.refreshToken,
		"scope":         oauthScope,
	}

	body, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("failed to marshal refresh payload: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", oauthTokenURL, strings.NewReader(string(body)))
	if err != nil {
		return fmt.Errorf("failed to create refresh request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := tm.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("refresh request failed: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read refresh response: %w", err)
	}

	if resp.StatusCode >= 400 {
		return fmt.Errorf("token refresh failed (HTTP %d): %s", resp.StatusCode, string(respBody))
	}

	var result struct {
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
		IDToken      string `json:"id_token"`
	}
	if err := json.Unmarshal(respBody, &result); err != nil {
		return fmt.Errorf("failed to parse refresh response: %w", err)
	}

	if result.AccessToken != "" {
		tm.accessToken = result.AccessToken
	}
	if result.RefreshToken != "" {
		tm.refreshToken = result.RefreshToken
	}
	if result.IDToken != "" {
		tm.idToken = result.IDToken
	}

	// Re-extract account_id from the new access token
	if _, _, _, accountID := parseChatGPTAuthToken(tm.accessToken); accountID != "" {
		tm.accountID = accountID
	}

	tm.lastRefresh = time.Now()
	tm.persistAuthLocked()

	slog.Info("Refreshed Codex OAuth token")
	return nil
}

func (tm *TokenManager) persistAuthLocked() {
	if tm.onUpdate == nil {
		return
	}
	if tm.rawAuth == nil {
		tm.rawAuth = map[string]interface{}{}
	}
	tokens, _ := tm.rawAuth["tokens"].(map[string]interface{})
	if tokens == nil {
		tokens = map[string]interface{}{}
		tm.rawAuth["tokens"] = tokens
	}
	tokens["access_token"] = tm.accessToken
	if tm.refreshToken != "" {
		tokens["refresh_token"] = tm.refreshToken
	}
	if tm.idToken != "" {
		tokens["id_token"] = tm.idToken
	}
	if tm.accountID != "" {
		tokens["account_id"] = tm.accountID
	}
	tm.rawAuth["last_refresh"] = time.Now().UTC().Format(time.RFC3339)

	if err := tm.onUpdate(cloneMap(tm.rawAuth)); err != nil {
		slog.Warn("Failed to persist refreshed Codex auth.json", "error", err)
	}
}

// parseChatGPTAuthToken parses a ChatGPT auth token in either JSON or JWT format.
// Returns (accessToken, refreshToken, idToken, accountID).
func parseChatGPTAuthToken(authToken string) (string, string, string, string) {
	if authToken == "" {
		return "", "", "", ""
	}

	// Strategy 1: JSON format with tokens object
	var authData struct {
		Tokens struct {
			AccessToken  string `json:"access_token"`
			RefreshToken string `json:"refresh_token"`
			IDToken      string `json:"id_token"`
			AccountID    string `json:"account_id"`
		} `json:"tokens"`
	}
	if err := json.Unmarshal([]byte(authToken), &authData); err == nil {
		if authData.Tokens.AccessToken != "" {
			return authData.Tokens.AccessToken, authData.Tokens.RefreshToken, authData.Tokens.IDToken, authData.Tokens.AccountID
		}
	}

	// Strategy 2: Plain JWT — extract account_id from payload
	accessToken, accountID := parseJWTAccountID(authToken)
	return accessToken, "", "", accountID
}

// parseJWTAccountID extracts the account ID from a JWT access token.
func parseJWTAccountID(token string) (string, string) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return token, ""
	}

	payload := parts[1]
	// Base64url padding
	if m := len(payload) % 4; m != 0 {
		payload += strings.Repeat("=", 4-m)
	}

	decoded, err := base64.URLEncoding.DecodeString(payload)
	if err != nil {
		// Try std encoding as fallback
		decoded, err = base64.StdEncoding.DecodeString(payload)
		if err != nil {
			return token, ""
		}
	}

	var claims map[string]interface{}
	if err := json.Unmarshal(decoded, &claims); err != nil {
		return token, ""
	}

	// Extract account_id from https://api.openai.com/auth.chatgpt_account_id
	if auth, ok := claims["https://api.openai.com/auth"].(map[string]interface{}); ok {
		if accountID, ok := auth["chatgpt_account_id"].(string); ok {
			return token, accountID
		}
	}

	return token, ""
}

// isAccessTokenExpired checks if a JWT access token is expired or will expire within buffer.
func isAccessTokenExpired(token string, buffer time.Duration) bool {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return false
	}

	payload := parts[1]
	if m := len(payload) % 4; m != 0 {
		payload += strings.Repeat("=", 4-m)
	}

	decoded, err := base64.URLEncoding.DecodeString(payload)
	if err != nil {
		decoded, err = base64.StdEncoding.DecodeString(payload)
		if err != nil {
			return false
		}
	}

	var claims map[string]interface{}
	if err := json.Unmarshal(decoded, &claims); err != nil {
		return false
	}

	exp, ok := claims["exp"].(float64)
	if !ok {
		return false
	}

	return time.Now().Unix() >= int64(exp)-int64(buffer.Seconds())
}

type authJSONCredentials struct {
	AccessToken  string
	RefreshToken string
	IDToken      string
	AccountID    string
	APIKey       string
	IsAPIKey     bool
	Raw          map[string]interface{}
}

// loadAuthJSON reads an auth.json file and extracts either ChatGPT OAuth tokens
// or a top-level OPENAI_API_KEY/openai_api_key credential.
func loadAuthJSON(path string) (*authJSONCredentials, error) {
	data, err := os.ReadFile(expandPath(path))
	if err != nil {
		return nil, fmt.Errorf("failed to read auth.json: %w", err)
	}
	return parseAuthJSON(data)
}

func loadAuthJSONContent(content string) (*authJSONCredentials, error) {
	return parseAuthJSON([]byte(content))
}

func parseAuthJSON(data []byte) (*authJSONCredentials, error) {
	var raw map[string]interface{}
	if err := json.Unmarshal(data, &raw); err != nil {
		return nil, fmt.Errorf("failed to parse auth.json: %w", err)
	}

	// Check for top-level OPENAI_API_KEY / openai_api_key (api_key mode)
	for _, key := range []string{"OPENAI_API_KEY", "openai_api_key"} {
		if v, ok := raw[key].(string); ok && strings.TrimSpace(v) != "" {
			return &authJSONCredentials{
				APIKey:   strings.TrimSpace(v),
				IsAPIKey: true,
				Raw:      cloneMap(raw),
			}, nil
		}
	}

	// Standard tokens structure
	tokens, ok := raw["tokens"].(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("auth.json: no tokens object found")
	}

	accessToken, _ := tokens["access_token"].(string)
	refreshToken, _ := tokens["refresh_token"].(string)
	idToken, _ := tokens["id_token"].(string)
	accountID, _ := tokens["account_id"].(string)

	return &authJSONCredentials{
		AccessToken:  accessToken,
		RefreshToken: refreshToken,
		IDToken:      idToken,
		AccountID:    accountID,
		Raw:          cloneMap(raw),
	}, nil
}

func expandPath(path string) string {
	if strings.HasPrefix(path, "~/") {
		home, _ := os.UserHomeDir()
		return filepath.Join(home, path[2:])
	}
	return path
}

func cloneMap(src map[string]interface{}) map[string]interface{} {
	if src == nil {
		return nil
	}
	data, err := json.Marshal(src)
	if err != nil {
		dst := make(map[string]interface{}, len(src))
		for k, v := range src {
			dst[k] = v
		}
		return dst
	}
	var dst map[string]interface{}
	if err := json.Unmarshal(data, &dst); err != nil {
		return nil
	}
	return dst
}

func parseAuthLastRefresh(value interface{}) time.Time {
	s, ok := value.(string)
	if !ok {
		return time.Time{}
	}
	s = strings.TrimSpace(s)
	if s == "" {
		return time.Time{}
	}
	for _, layout := range []string{time.RFC3339Nano, time.RFC3339, "2006-01-02 15:04:05"} {
		if t, err := time.Parse(layout, s); err == nil {
			return t
		}
	}
	return time.Time{}
}