package codex

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/apr/types"
)

type setupOptions struct {
	onPersist       func(*CodexConfig) error
	httpClient      *http.Client
	tokenHTTPClient *http.Client
	lifecycle       context.Context
}

// SetupOption customizes Codex provider setup.
type SetupOption func(*setupOptions)

// WithConfigPersistence registers a callback invoked after refreshed auth data
// is copied into the setup config. The callback should persist the root config
// when durable inline auth_json_content updates are desired.
func WithConfigPersistence(onPersist func(*CodexConfig) error) SetupOption {
	return func(opts *setupOptions) {
		opts.onPersist = onPersist
	}
}

// WithHTTPClient sets the client used for Codex upstream API calls.
func WithHTTPClient(httpClient *http.Client) SetupOption {
	return func(opts *setupOptions) {
		opts.httpClient = httpClient
	}
}

// WithTokenHTTPClient sets the client used for Codex OAuth token refresh calls.
func WithTokenHTTPClient(httpClient *http.Client) SetupOption {
	return func(opts *setupOptions) {
		opts.tokenHTTPClient = httpClient
	}
}

// WithLifecycleContext sets the context that owns background provider work.
func WithLifecycleContext(ctx context.Context) SetupOption {
	return func(opts *setupOptions) {
		if ctx != nil {
			opts.lifecycle = ctx
		}
	}
}

// Setup initializes the Codex provider from configuration.
func Setup(cfg *CodexConfig, tr types.HTTPProxy, ctx context.Context, opts ...SetupOption) (*Codex, error) {
	options := setupOptions{}
	for _, opt := range opts {
		opt(&options)
	}
	return setupWithOptions(cfg, tr, ctx, options)
}

func setupWithOptions(cfg *CodexConfig, tr types.HTTPProxy, ctx context.Context, options setupOptions) (*Codex, error) {
	normalizeCodexConfig(cfg)
	lifecycle := options.lifecycle
	if lifecycle == nil {
		lifecycle = ctx
	}

	mode := cfg.Mode
	baseURL := cfg.BaseURL
	if baseURL == "" {
		baseURL = "https://chatgpt.com"
	}

	// Auto-detect mode if not explicitly set
	if mode == "" {
		if cfg.AuthToken != "" || cfg.AuthJSONContent != "" || cfg.AuthJSONPath != "" {
			mode = "chatgpt"
		} else if cfg.APIKey != "" {
			mode = "api_key"
		} else if strings.EqualFold(cfg.Mode, "passthrough") {
			mode = "passthrough"
		} else {
			return nil, fmt.Errorf("codex: no authentication configured (need auth_token, auth_json_content, auth_json_path, or api_key)")
		}
	}

	transport := &http.Transport{
		MaxIdleConns:        100,
		MaxIdleConnsPerHost: 100,
		IdleConnTimeout:     90 * time.Second,
	}
	if tr != nil {
		proxyTransport := tr.Transport()
		if proxyTransport != nil && proxyTransport.Proxy != nil {
			transport.Proxy = proxyTransport.Proxy
		}
	}
	httpClient := options.httpClient
	if httpClient == nil {
		httpClient = &http.Client{
			Transport: transport,
			Timeout:   10 * time.Minute,
		}
	}

	c := &Codex{
		mode:                mode,
		baseURL:             correctBaseURL(baseURL),
		backendName:         cfg.Name,
		backendID:           cfg.ID,
		active:              isActive(cfg.Active),
		activeSet:           true,
		weight:              cfg.Weight,
		allowedAPIKeys:      stringSet(cfg.AllowedAPIKeys),
		requiresOAuthBeta:   codexRequiresOAuthBeta(cfg),
		customPrompt:        cfg.CustomPrompt,
		customPromptModels:  append([]string(nil), cfg.CustomPromptModels...),
		httpClient:          httpClient,
		tokenHTTPClient:     options.tokenHTTPClient,
		lifecycleContext:    lifecycle,
		adaptiveWeights:     cfg.AdaptiveWeights,
		weightDiffThreshold: codexWeightDiffThreshold(cfg),
		authSWRR:            make(map[*Codex]float64),
		fallbackSWRR:        make(map[*Codex]float64),
		authEffective:       make(map[*Codex]float64),
		fallbackEffective:   make(map[*Codex]float64),
	}
	c.lastError.Store("")

	switch mode {
	case "chatgpt":
		authToken := strings.TrimSpace(cfg.AuthToken)
		var rawAuth map[string]interface{}

		// Prefer inline auth_json_content over auth_json_path, matching llmproxy.
		if authToken == "" {
			creds, err := loadConfiguredAuthJSON(cfg)
			if err != nil {
				return nil, err
			}
			if creds != nil {
				if creds.IsAPIKey {
					c.mode = "api_key"
					c.apiKey = creds.APIKey
					cfg.APIKey = creds.APIKey
					break
				}
				rawAuth = creds.Raw
				authToken, err = authJSONToTokenManagerInput(creds)
				if err != nil {
					return nil, err
				}
			}
		} else if parsed, err := loadAuthJSONContent(authToken); err == nil {
			if parsed.IsAPIKey {
				c.mode = "api_key"
				c.apiKey = parsed.APIKey
				cfg.APIKey = parsed.APIKey
				break
			}
			rawAuth = parsed.Raw
		}

		if authToken == "" {
			return nil, fmt.Errorf("codex: no auth_token, auth_json_content, or auth_json_path configured for chatgpt mode")
		}

		opts := []TokenManagerOption{}
		if options.tokenHTTPClient != nil {
			opts = append(opts, withTokenManagerHTTPClient(options.tokenHTTPClient))
		}
		opts = append(opts, withTokenManagerInitialContext(ctx))
		if rawAuth != nil {
			opts = append(opts, WithAuthPersistence(rawAuth, persistAuthJSON(cfg, options.onPersist)))
		}
		tm, err := NewTokenManager(authToken, opts...)
		if err != nil {
			return nil, fmt.Errorf("codex: failed to initialize token manager: %w", err)
		}
		c.tokenManager = tm
		if accessToken, _ := tm.GetAccessToken(); accessToken != "" {
			cfg.AuthToken = accessToken
		}
		tm.StartBackgroundRefresh(lifecycle)

	case "api_key":
		if cfg.APIKey == "" {
			return nil, fmt.Errorf("codex: no api_key configured for api_key mode")
		}
		c.apiKey = cfg.APIKey
	case "passthrough":
		// No configured auth is injected.

	default:
		return nil, fmt.Errorf("codex: unknown mode %q (expected \"chatgpt\", \"api_key\", or \"passthrough\")", mode)
	}

	// Static model list
	modelProvider := types.NewStaticModelProvider(providerName, codexModels)
	modelCache := types.NewModelCache(modelProvider)
	if err := modelCache.Initialize(); err != nil {
		return nil, fmt.Errorf("codex: failed to initialize model cache: %w", err)
	}
	c.modelCache = modelCache

	for i := range cfg.AuthFiles {
		authBackendCfg := codexConfigForAuthFile(cfg, cfg.AuthFiles[i])
		authIndex := i
		childOptions := options
		childOptions.onPersist = func(updated *CodexConfig) error {
			copyCodexAuthToAuthFile(&cfg.AuthFiles[authIndex], updated)
			if options.onPersist != nil {
				return options.onPersist(cfg)
			}
			return nil
		}
		authBackend, err := setupWithOptions(&authBackendCfg, tr, ctx, childOptions)
		if err != nil {
			return nil, fmt.Errorf("codex: failed to initialize auth file %d: %w", i, err)
		}
		c.authBackends = append(c.authBackends, authBackend)
	}

	for i := range cfg.Fallbacks {
		fallbackCfg := cfg.Fallbacks[i]
		fallbackCfg.Fallbacks = nil
		fallbackIndex := i
		childOptions := options
		childOptions.onPersist = func(updated *CodexConfig) error {
			copyCodexRuntimeAuth(&cfg.Fallbacks[fallbackIndex], updated)
			if options.onPersist != nil {
				return options.onPersist(cfg)
			}
			return nil
		}
		fallback, err := setupWithOptions(&fallbackCfg, tr, ctx, childOptions)
		if err != nil {
			return nil, fmt.Errorf("codex: failed to initialize fallback %d: %w", i, err)
		}
		c.fallbacks = append(c.fallbacks, fallback)
	}

	slog.Info("Registered Codex provider",
		"mode", c.mode,
		"base_url", c.baseURL,
		"auth_backends", len(c.authBackends)+1,
		"fallbacks", len(c.fallbacks),
	)

	return c, nil
}

// ReloadAuth reinitializes auth backends and fallbacks from an updated config.
// It stops old token managers, rebuilds auth state, and atomically swaps backends.
func (c *Codex) ReloadAuth(cfg *CodexConfig, tr types.HTTPProxy, ctx context.Context) error {
	// Stop old token managers
	if c.tokenManager != nil {
		c.tokenManager.Stop()
	}
	for _, backend := range c.authBackends {
		if backend.tokenManager != nil {
			backend.tokenManager.Stop()
		}
	}
	for _, fallback := range c.fallbacks {
		if fallback.tokenManager != nil {
			fallback.tokenManager.Stop()
		}
	}

	// Rebuild primary auth by re-running setup logic
	normalized := *cfg
	normalizeCodexConfig(&normalized)

	mode := normalized.Mode
	if mode == "" {
		if normalized.AuthToken != "" || normalized.AuthJSONContent != "" || normalized.AuthJSONPath != "" {
			mode = "chatgpt"
		} else if normalized.APIKey != "" {
			mode = "api_key"
		}
	}

	baseURL := normalized.BaseURL
	if baseURL == "" {
		baseURL = "https://chatgpt.com"
	}
	c.baseURL = correctBaseURL(baseURL)
	c.backendName = normalized.Name
	c.backendID = normalized.ID
	c.active = isActive(normalized.Active)
	c.activeSet = true
	c.weight = normalized.Weight
	c.allowedAPIKeys = stringSet(normalized.AllowedAPIKeys)
	c.requiresOAuthBeta = codexRequiresOAuthBeta(&normalized)
	c.customPrompt = normalized.CustomPrompt
	c.customPromptModels = append([]string(nil), normalized.CustomPromptModels...)
	c.adaptiveWeights = normalized.AdaptiveWeights
	c.weightDiffThreshold = codexWeightDiffThreshold(&normalized)

	// Update primary auth fields
	switch mode {
	case "chatgpt":
		authToken := strings.TrimSpace(normalized.AuthToken)
		var rawAuth map[string]interface{}

		if authToken == "" {
			creds, err := loadConfiguredAuthJSON(&normalized)
			if err != nil {
				return err
			}
			if creds != nil {
				if creds.IsAPIKey {
					c.mode = "api_key"
					c.apiKey = creds.APIKey
					c.tokenManager = nil
					goto rebuildBackends
				}
				rawAuth = creds.Raw
				authToken, err = authJSONToTokenManagerInput(creds)
				if err != nil {
					return err
				}
			}
		} else if parsed, err := loadAuthJSONContent(authToken); err == nil {
			if parsed.IsAPIKey {
				c.mode = "api_key"
				c.apiKey = parsed.APIKey
				c.tokenManager = nil
				goto rebuildBackends
			}
			rawAuth = parsed.Raw
		}

		if authToken == "" {
			return fmt.Errorf("codex: no auth token available for reload")
		}

		opts := []TokenManagerOption{}
		if c.tokenHTTPClient != nil {
			opts = append(opts, withTokenManagerHTTPClient(c.tokenHTTPClient))
		}
		opts = append(opts, withTokenManagerInitialContext(ctx))
		if rawAuth != nil {
			opts = append(opts, WithAuthPersistence(rawAuth, persistAuthJSON(cfg, nil)))
		}
		tm, err := NewTokenManager(authToken, opts...)
		if err != nil {
			return fmt.Errorf("codex: failed to reinitialize token manager: %w", err)
		}
		c.tokenManager = tm
		c.mode = "chatgpt"
		c.apiKey = ""
		if accessToken, _ := tm.GetAccessToken(); accessToken != "" {
			cfg.AuthToken = accessToken
		}
		lifecycle := c.lifecycleContext
		if lifecycle == nil {
			lifecycle = ctx
		}
		tm.StartBackgroundRefresh(lifecycle)

	case "api_key":
		c.mode = "api_key"
		c.tokenManager = nil
		c.apiKey = normalized.APIKey
	case "passthrough":
		c.mode = "passthrough"
		c.tokenManager = nil
		c.apiKey = ""
	default:
		return fmt.Errorf("codex: no authentication configured (need auth_token, auth_json_content, auth_json_path, api_key, or mode: passthrough)")
	}

rebuildBackends:
	// Rebuild auth backends
	newAuthBackends := make([]*Codex, 0, len(normalized.AuthFiles))
	for i := range normalized.AuthFiles {
		authBackendCfg := codexConfigForAuthFile(&normalized, normalized.AuthFiles[i])
		authBackend, err := setupWithOptions(&authBackendCfg, tr, ctx, setupOptions{
			httpClient:      c.httpClient,
			tokenHTTPClient: c.tokenHTTPClient,
			lifecycle:       c.lifecycleContext,
		})
		if err != nil {
			slog.Warn("Codex reload: failed to initialize auth file", "index", i, "error", err)
			continue
		}
		newAuthBackends = append(newAuthBackends, authBackend)
	}

	// Rebuild fallbacks
	newFallbacks := make([]*Codex, 0, len(normalized.Fallbacks))
	for i := range normalized.Fallbacks {
		fallbackCfg := normalized.Fallbacks[i]
		fallbackCfg.Fallbacks = nil
		fallback, err := setupWithOptions(&fallbackCfg, tr, ctx, setupOptions{
			httpClient:      c.httpClient,
			tokenHTTPClient: c.tokenHTTPClient,
			lifecycle:       c.lifecycleContext,
		})
		if err != nil {
			slog.Warn("Codex reload: failed to initialize fallback", "index", i, "error", err)
			continue
		}
		newFallbacks = append(newFallbacks, fallback)
	}

	// Atomically swap
	c.selectorMu.Lock()
	c.authBackends = newAuthBackends
	c.fallbacks = newFallbacks
	c.authSWRR = make(map[*Codex]float64)
	c.fallbackSWRR = make(map[*Codex]float64)
	c.authEffective = make(map[*Codex]float64)
	c.fallbackEffective = make(map[*Codex]float64)
	c.selectorMu.Unlock()

	slog.Info("Codex provider auth reloaded",
		"mode", c.mode,
		"auth_backends", len(c.authBackends)+1,
		"fallbacks", len(c.fallbacks),
	)
	return nil
}

func normalizeCodexConfig(cfg *CodexConfig) {
	authFiles := collectAuthFiles(cfg)
	cfg.AuthJSONContents = nil
	cfg.AuthJSONPaths = nil
	cfg.AuthFiles = nil

	if !hasPrimaryAuth(cfg) && len(authFiles) > 0 {
		applyAuthFile(cfg, authFiles[0])
		authFiles = authFiles[1:]
	}
	cfg.AuthFiles = authFiles
}

func collectAuthFiles(cfg *CodexConfig) []CodexAuthConfig {
	authFiles := make([]CodexAuthConfig, 0, len(cfg.AuthFiles)+len(cfg.AuthJSONContents)+len(cfg.AuthJSONPaths))
	authFiles = append(authFiles, cfg.AuthFiles...)
	for _, content := range cfg.AuthJSONContents {
		if strings.TrimSpace(content) == "" {
			continue
		}
		authFiles = append(authFiles, CodexAuthConfig{AuthJSONContent: content})
	}
	for _, path := range cfg.AuthJSONPaths {
		if strings.TrimSpace(path) == "" {
			continue
		}
		authFiles = append(authFiles, CodexAuthConfig{AuthJSONPath: path})
	}
	return authFiles
}

func hasPrimaryAuth(cfg *CodexConfig) bool {
	return strings.TrimSpace(cfg.AuthToken) != "" ||
		strings.TrimSpace(cfg.AuthJSONContent) != "" ||
		strings.TrimSpace(cfg.AuthJSONPath) != "" ||
		strings.TrimSpace(cfg.APIKey) != ""
}

func codexConfigForAuthFile(base *CodexConfig, authFile CodexAuthConfig) CodexConfig {
	cfg := *base
	cfg.ID = 0
	cfg.Name = ""
	cfg.AuthToken = ""
	cfg.AuthJSONContent = ""
	cfg.AuthJSONPath = ""
	cfg.APIKey = ""
	cfg.Mode = ""
	cfg.AuthJSONContents = nil
	cfg.AuthJSONPaths = nil
	cfg.AuthFiles = nil
	cfg.Fallbacks = nil
	applyAuthFile(&cfg, authFile)
	return cfg
}

func applyAuthFile(cfg *CodexConfig, authFile CodexAuthConfig) {
	cfg.ID = authFile.ID
	cfg.Name = authFile.Name
	cfg.Active = authFile.Active
	cfg.Weight = authFile.Weight
	cfg.AllowedAPIKeys = append([]string(nil), authFile.AllowedAPIKeys...)
	cfg.AuthToken = authFile.AuthToken
	cfg.AuthJSONContent = authFile.AuthJSONContent
	cfg.AuthJSONPath = authFile.AuthJSONPath
	cfg.APIKey = authFile.APIKey
	cfg.Mode = authFile.Mode
	if authFile.RequiresOAuthBeta || authFile.RequiresOAuthBetaLegacy {
		cfg.RequiresOAuthBeta = true
	}
}

func codexRequiresOAuthBeta(cfg *CodexConfig) bool {
	if cfg == nil {
		return false
	}
	return cfg.RequiresOAuthBeta || cfg.RequiresOAuthBetaLegacy
}

func isActive(active *bool) bool {
	return active == nil || *active
}

func codexWeightDiffThreshold(cfg *CodexConfig) float64 {
	if cfg.WeightDiffThreshold > 0 {
		return cfg.WeightDiffThreshold
	}
	return 10.0
}

func stringSet(values []string) map[string]struct{} {
	if len(values) == 0 {
		return nil
	}
	set := make(map[string]struct{}, len(values))
	for _, value := range values {
		value = strings.TrimSpace(value)
		if value != "" {
			set[value] = struct{}{}
		}
	}
	return set
}

func copyCodexAuthToAuthFile(dst *CodexAuthConfig, src *CodexConfig) {
	dst.AuthToken = src.AuthToken
	dst.AuthJSONContent = src.AuthJSONContent
	dst.AuthJSONPath = src.AuthJSONPath
	dst.APIKey = src.APIKey
	dst.Mode = src.Mode
}

func copyCodexRuntimeAuth(dst *CodexConfig, src *CodexConfig) {
	dst.AuthToken = src.AuthToken
	dst.AuthJSONContent = src.AuthJSONContent
	dst.AuthJSONPath = src.AuthJSONPath
	dst.APIKey = src.APIKey
	dst.Mode = src.Mode
	dst.AuthFiles = src.AuthFiles
	dst.AuthJSONContents = src.AuthJSONContents
	dst.AuthJSONPaths = src.AuthJSONPaths
}

func loadConfiguredAuthJSON(cfg *CodexConfig) (*authJSONCredentials, error) {
	if strings.TrimSpace(cfg.AuthJSONContent) != "" {
		creds, err := loadAuthJSONContent(cfg.AuthJSONContent)
		if err != nil {
			return nil, fmt.Errorf("codex: failed to parse auth_json_content: %w", err)
		}
		return creds, nil
	}
	if strings.TrimSpace(cfg.AuthJSONPath) != "" {
		creds, err := loadAuthJSON(cfg.AuthJSONPath)
		if err != nil {
			return nil, fmt.Errorf("codex: failed to load auth.json: %w", err)
		}
		return creds, nil
	}
	return nil, nil
}

func authJSONToTokenManagerInput(creds *authJSONCredentials) (string, error) {
	wrapper := map[string]interface{}{
		"tokens": map[string]string{
			"access_token":  creds.AccessToken,
			"refresh_token": creds.RefreshToken,
			"id_token":      creds.IDToken,
			"account_id":    creds.AccountID,
		},
	}
	data, err := json.Marshal(wrapper)
	if err != nil {
		return "", fmt.Errorf("codex: failed to marshal auth data: %w", err)
	}
	return string(data), nil
}

func persistAuthJSON(cfg *CodexConfig, onPersist func(*CodexConfig) error) func(map[string]interface{}) error {
	return func(raw map[string]interface{}) error {
		data, err := json.MarshalIndent(raw, "", "  ")
		if err != nil {
			return fmt.Errorf("marshal auth json: %w", err)
		}
		serialized := string(data)
		cfg.AuthJSONContent = serialized

		if tokens, ok := raw["tokens"].(map[string]interface{}); ok {
			if accessToken, ok := tokens["access_token"].(string); ok {
				cfg.AuthToken = accessToken
			}
		}
		if apiKey, ok := raw["OPENAI_API_KEY"].(string); ok && strings.TrimSpace(apiKey) != "" {
			cfg.APIKey = strings.TrimSpace(apiKey)
			cfg.AuthToken = ""
		} else if apiKey, ok := raw["openai_api_key"].(string); ok && strings.TrimSpace(apiKey) != "" {
			cfg.APIKey = strings.TrimSpace(apiKey)
			cfg.AuthToken = ""
		}

		if strings.TrimSpace(cfg.AuthJSONPath) == "" {
			if onPersist != nil {
				return onPersist(cfg)
			}
			return nil
		}
		path := expandPath(cfg.AuthJSONPath)
		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
			return fmt.Errorf("create auth json directory: %w", err)
		}
		if err := os.WriteFile(path, append(data, '\n'), 0o600); err != nil {
			return fmt.Errorf("write auth json: %w", err)
		}
		if onPersist != nil {
			return onPersist(cfg)
		}
		return nil
	}
}