package extensionflow

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/extension/egress"
	"github.com/tidbcloud/AgentTeam/internal/extension/spec"
)

// OAuthCredential is the closed decrypted credential shape supplied to an
// operator-owned OAuth refresher. Token endpoint/origin is deliberately absent:
// connection JSON cannot turn refresh into an arbitrary egress channel.
type OAuthCredential struct {
	AccessToken  string
	RefreshToken string
	Scopes       []string
	ExpiresAt    time.Time
}

// OAuthRefreshRequest is safe execution metadata plus the current provider
// credential. It is process-local only and must never be logged or audited.
type OAuthRefreshRequest struct {
	DefinitionDigest string
	AuthProfileID    string
	RequiredScopes   []string
	Credential       OAuthCredential
}

// OAuthRefresher is an operator-wired, provider-specific refresh collaborator.
// It owns the fixed provider token endpoint and must not derive it from stored
// credential JSON or Agent input.
type OAuthRefresher interface {
	RefreshExtensionOAuth(ctx context.Context, request OAuthRefreshRequest) (OAuthCredential, error)
}

type storedOAuthCredential struct {
	AccessToken       string   `json:"accessToken"`
	RefreshToken      string   `json:"refreshToken,omitempty"`
	Scopes            []string `json:"scopes"`
	ExpiresAt         string   `json:"expiresAt,omitempty"`
	LegacyAccessToken string   `json:"access_token,omitempty"`
	LegacyRefresh     string   `json:"refresh_token,omitempty"`
	LegacyScope       string   `json:"scope,omitempty"`
	LegacyExpiresAt   string   `json:"expires_at,omitempty"`
}

type storedAPIKeyCredential struct {
	Value string `json:"value"`
}

func resolveProviderCredential(
	profile *spec.AuthProfile,
	requirement *spec.ActionAuthRequirement,
	plaintext []byte,
	now time.Time,
) (egress.Credential, OAuthCredential, bool, error) {
	if requirement == nil {
		if profile != nil {
			return egress.Credential{}, OAuthCredential{}, false,
				errors.New("no-auth action unexpectedly resolved an auth profile")
		}
		return egress.Credential{}, OAuthCredential{}, false, nil
	}
	if profile == nil || profile.ID != requirement.Profile {
		return egress.Credential{}, OAuthCredential{}, false,
			errors.New("action auth profile is inconsistent")
	}
	switch profile.Type {
	case egress.CredentialOAuth2Bearer:
		stored := storedOAuthCredential{}
		if err := decodeStrictCredentialJSON(plaintext, &stored); err != nil {
			return egress.Credential{}, OAuthCredential{}, false,
				fmt.Errorf("decode OAuth credential: %w", err)
		}
		credential, err := validateOAuthCredential(stored)
		if err != nil {
			return egress.Credential{}, OAuthCredential{}, false, err
		}
		if err := requireOAuthScopes(credential.Scopes, requirement.RequiredScopes); err != nil {
			return egress.Credential{}, OAuthCredential{}, false, err
		}
		expired := !credential.ExpiresAt.IsZero() &&
			!credential.ExpiresAt.After(now.UTC().Add(2*time.Minute))
		if credential.AccessToken == "" && !expired {
			return egress.Credential{}, OAuthCredential{}, false,
				errors.New("OAuth access token is required")
		}
		return egress.Credential{
			Type:  egress.CredentialOAuth2Bearer,
			Value: credential.AccessToken,
		}, credential, expired, nil
	case egress.CredentialAPIKey:
		stored := storedAPIKeyCredential{}
		if err := decodeStrictCredentialJSON(plaintext, &stored); err != nil {
			return egress.Credential{}, OAuthCredential{}, false,
				fmt.Errorf("decode API key credential: %w", err)
		}
		if stored.Value == "" || strings.ContainsAny(stored.Value, "\r\n") {
			return egress.Credential{}, OAuthCredential{}, false,
				errors.New("API key value is invalid")
		}
		return egress.Credential{
			Type:  egress.CredentialAPIKey,
			In:    profile.In,
			Name:  profile.Name,
			Value: stored.Value,
		}, OAuthCredential{}, false, nil
	default:
		return egress.Credential{}, OAuthCredential{}, false,
			fmt.Errorf("unsupported credential profile type %q", profile.Type)
	}
}

func normalizeStoredOAuthCredential(stored storedOAuthCredential) storedOAuthCredential {
	if stored.AccessToken == "" {
		stored.AccessToken = stored.LegacyAccessToken
	}
	if stored.RefreshToken == "" {
		stored.RefreshToken = stored.LegacyRefresh
	}
	if stored.ExpiresAt == "" {
		stored.ExpiresAt = stored.LegacyExpiresAt
	}
	if stored.Scopes == nil && stored.LegacyScope != "" {
		stored.Scopes = strings.Fields(stored.LegacyScope)
	}
	return stored
}

func validateOAuthCredential(stored storedOAuthCredential) (OAuthCredential, error) {
	if stored.AccessToken != "" && stored.LegacyAccessToken != "" ||
		stored.RefreshToken != "" && stored.LegacyRefresh != "" ||
		stored.ExpiresAt != "" && stored.LegacyExpiresAt != "" ||
		stored.Scopes != nil && stored.LegacyScope != "" {
		return OAuthCredential{}, errors.New("OAuth credential mixes canonical and legacy fields")
	}
	stored = normalizeStoredOAuthCredential(stored)
	if strings.ContainsAny(stored.AccessToken, "\r\n") ||
		strings.ContainsAny(stored.RefreshToken, "\r\n") {
		return OAuthCredential{}, errors.New("OAuth token is invalid")
	}
	scopes, err := validatedSortedScopes(stored.Scopes)
	if err != nil {
		return OAuthCredential{}, err
	}
	var expiresAt time.Time
	if stored.ExpiresAt != "" {
		expiresAt, err = time.Parse(time.RFC3339, stored.ExpiresAt)
		if err != nil {
			return OAuthCredential{}, errors.New("OAuth expiresAt must be RFC3339")
		}
		expiresAt = expiresAt.UTC()
	}
	return OAuthCredential{
		AccessToken:  stored.AccessToken,
		RefreshToken: stored.RefreshToken,
		Scopes:       scopes,
		ExpiresAt:    expiresAt,
	}, nil
}

func validateRefreshedOAuthCredential(
	credential OAuthCredential,
	requiredScopes []string,
	now time.Time,
) (OAuthCredential, error) {
	if credential.AccessToken == "" ||
		strings.ContainsAny(credential.AccessToken, "\r\n") ||
		strings.ContainsAny(credential.RefreshToken, "\r\n") {
		return OAuthCredential{}, errors.New("refreshed OAuth token is invalid")
	}
	scopes, err := validatedSortedScopes(credential.Scopes)
	if err != nil {
		return OAuthCredential{}, err
	}
	credential.Scopes = scopes
	credential.ExpiresAt = credential.ExpiresAt.UTC()
	if !credential.ExpiresAt.IsZero() &&
		!credential.ExpiresAt.After(now.UTC().Add(2*time.Minute)) {
		return OAuthCredential{}, errors.New("refreshed OAuth token expires inside the refresh skew")
	}
	if err := requireOAuthScopes(scopes, requiredScopes); err != nil {
		return OAuthCredential{}, err
	}
	return credential, nil
}

func encodeStoredOAuthCredential(credential OAuthCredential) ([]byte, error) {
	stored := storedOAuthCredential{
		AccessToken:  credential.AccessToken,
		RefreshToken: credential.RefreshToken,
		Scopes:       append([]string(nil), credential.Scopes...),
	}
	if !credential.ExpiresAt.IsZero() {
		stored.ExpiresAt = credential.ExpiresAt.UTC().Format(time.RFC3339)
	}
	return json.Marshal(stored)
}

func decodeStrictCredentialJSON(raw []byte, target any) error {
	if len(bytes.TrimSpace(raw)) == 0 {
		return errors.New("credential JSON is empty")
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(target); err != nil {
		return err
	}
	var trailing json.RawMessage
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		if err == nil {
			return errors.New("credential JSON must contain one document")
		}
		return errors.New("credential JSON has trailing data")
	}
	return nil
}

func validatedSortedScopes(scopes []string) ([]string, error) {
	out := append([]string(nil), scopes...)
	sort.Strings(out)
	for i, scope := range out {
		if scope == "" || strings.TrimSpace(scope) != scope {
			return nil, errors.New("OAuth scopes must be non-empty exact strings")
		}
		if i > 0 && out[i-1] == scope {
			return nil, fmt.Errorf("OAuth scope %q is duplicated", scope)
		}
	}
	return out, nil
}

func requireOAuthScopes(granted, required []string) error {
	grantedSet := make(map[string]struct{}, len(granted))
	for _, scope := range granted {
		grantedSet[scope] = struct{}{}
	}
	for _, scope := range required {
		if _, ok := grantedSet[scope]; !ok {
			return fmt.Errorf("OAuth credential is missing required scope %q", scope)
		}
	}
	return nil
}