package runtime

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strconv"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/automation/dsl"
	"github.com/tidbcloud/AgentTeam/internal/core/secrets"
	"github.com/tidbcloud/AgentTeam/internal/store/automation"
)

var serverRunnerMaskReplacement = []byte("***")

type serverRunnerResolvedEnv struct {
	jobEnv  map[string]string
	stepEnv []map[string]string
	secrets map[string]string
}

func (env serverRunnerResolvedEnv) envForStep(step int) map[string]string {
	if len(env.jobEnv) == 0 && (step < 0 || step >= len(env.stepEnv) || len(env.stepEnv[step]) == 0) {
		return nil
	}
	merged := make(map[string]string, len(env.jobEnv))
	for key, value := range env.jobEnv {
		merged[key] = value
	}
	if step >= 0 && step < len(env.stepEnv) {
		for key, value := range env.stepEnv[step] {
			merged[key] = value
		}
	}
	return merged
}

func (r *ServerRunner) resolveShellJobEnv(ctx context.Context, workspaceID string, runID string, spec dsl.JobSpec) (serverRunnerResolvedEnv, error) {
	refs, err := dsl.JobEnvReferences(spec)
	if err != nil {
		return serverRunnerResolvedEnv{}, fmt.Errorf("automation job env is invalid")
	}
	var run automation.AutomationRun
	if len(refs.SecretNames) > 0 || refs.UsesInputs {
		run, err = r.store.GetRunEnvContext(ctx, workspaceID, runID)
		if err != nil {
			return serverRunnerResolvedEnv{}, fmt.Errorf("automation run for job is not available")
		}
	}
	secretsByName, err := ResolveAutomationSecrets(ctx, r.store, r.cfg.CredentialKey, run, refs.SecretNames)
	if err != nil {
		return serverRunnerResolvedEnv{}, err
	}
	var inputs map[string]any
	if refs.UsesInputs {
		inputs, err = serverRunnerAutomationInputs(run.InputJSON)
		if err != nil {
			return serverRunnerResolvedEnv{}, err
		}
	}
	jobEnv, err := serverRunnerResolveEnvMap(spec.Env, inputs, secretsByName)
	if err != nil {
		return serverRunnerResolvedEnv{}, err
	}
	stepEnv := make([]map[string]string, len(spec.Steps))
	for idx, step := range spec.Steps {
		resolved, err := serverRunnerResolveEnvMap(step.Env, inputs, secretsByName)
		if err != nil {
			return serverRunnerResolvedEnv{}, err
		}
		stepEnv[idx] = resolved
	}
	return serverRunnerResolvedEnv{jobEnv: jobEnv, stepEnv: stepEnv, secrets: secretsByName}, nil
}

// ResolveAutomationSecrets loads and decrypts the explicit automation secret names
// referenced by one automation run.
func ResolveAutomationSecrets(ctx context.Context, store *automation.Store, credentialKey []byte, run automation.AutomationRun, names []string) (map[string]string, error) {
	if len(names) == 0 {
		return nil, nil
	}
	codec, err := secrets.NewAESGCMCodec(credentialKey)
	if err != nil {
		return nil, fmt.Errorf("automation secret storage is not configured on server: %w", err)
	}
	secrets, err := store.GetAutomationSecretCiphertextsByNames(ctx, run.WorkspaceID, run.RepositoryOwner, run.RepositoryName, names)
	if err != nil {
		return nil, err
	}
	values := make(map[string]string, len(names))
	for _, name := range names {
		ciphertext, ok := secrets[name]
		if !ok {
			return nil, fmt.Errorf("automation secret %s is not available", name)
		}
		plaintext, err := codec.Open(ciphertext)
		if err != nil {
			return nil, fmt.Errorf("automation secret %s cannot be decrypted", name)
		}
		values[name] = plaintext
	}
	return values, nil
}

func serverRunnerResolveEnvMap(env map[string]any, inputs map[string]any, secrets map[string]string) (map[string]string, error) {
	if len(env) == 0 {
		return nil, nil
	}
	keys := make([]string, 0, len(env))
	for key := range env {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	resolved := make(map[string]string, len(env))
	for _, key := range keys {
		value, err := dsl.ParseEnvValue(env[key])
		if err != nil {
			return nil, fmt.Errorf("automation env %s is invalid", key)
		}
		switch value.Kind {
		case dsl.EnvValueLiteral:
			resolved[key] = value.Value
		case dsl.EnvValueSecret:
			secret, ok := secrets[value.Name]
			if !ok {
				return nil, fmt.Errorf("automation secret %s is not available", value.Name)
			}
			resolved[key] = secret
		case dsl.EnvValueInput:
			input, err := serverRunnerInputValue(inputs, value.Name)
			if err != nil {
				return nil, err
			}
			resolved[key] = input
		default:
			return nil, fmt.Errorf("automation env %s is invalid", key)
		}
	}
	return resolved, nil
}

func serverRunnerAutomationInputs(raw string) (map[string]any, error) {
	if strings.TrimSpace(raw) == "" {
		return nil, nil
	}
	decoder := json.NewDecoder(strings.NewReader(raw))
	decoder.UseNumber()
	var payload any
	if err := decoder.Decode(&payload); err != nil {
		return nil, fmt.Errorf("automation run input_json is invalid")
	}
	inputs, ok := payload.(map[string]any)
	if !ok {
		return nil, nil
	}
	return inputs, nil
}

func serverRunnerInputValue(inputs map[string]any, name string) (string, error) {
	if len(inputs) == 0 {
		return "", nil
	}
	value, ok := inputs[name]
	if !ok || value == nil {
		return "", nil
	}
	switch typed := value.(type) {
	case string:
		return typed, nil
	case bool:
		return strconv.FormatBool(typed), nil
	case json.Number:
		return typed.String(), nil
	default:
		return "", fmt.Errorf("automation input %s cannot be used as env", name)
	}
}

type serverRunnerSecretMatcher struct {
	secrets [][]byte
}

type serverRunnerSecretMasker struct {
	matcher *serverRunnerSecretMatcher
	buffer  []byte
}

// SecretRedactor masks resolved automation secret values from runner-controlled
// callback surfaces before they are logged or persisted.
type SecretRedactor struct {
	matcher *serverRunnerSecretMatcher
}

// NewSecretRedactor returns a redactor for the provided secret values.
func NewSecretRedactor(secretsByName map[string]string) *SecretRedactor {
	matcher := newServerRunnerSecretMatcher(secretsByName)
	if matcher == nil {
		return nil
	}
	return &SecretRedactor{matcher: matcher}
}

// RedactBytes replaces each known secret value in payload with a fixed mask.
func (r *SecretRedactor) RedactBytes(payload []byte) []byte {
	if r == nil || r.matcher == nil || len(payload) == 0 {
		return payload
	}
	masker := newServerRunnerSecretMasker(r.matcher)
	return masker.Mask(payload, true)
}

// RedactString replaces each known secret value in value with a fixed mask.
func (r *SecretRedactor) RedactString(value string) string {
	if r == nil || r.matcher == nil || value == "" {
		return value
	}
	return string(r.RedactBytes([]byte(value)))
}

func newServerRunnerSecretMatcher(secretsByName map[string]string) *serverRunnerSecretMatcher {
	if len(secretsByName) == 0 {
		return nil
	}
	seen := map[string]struct{}{}
	secrets := make([][]byte, 0, len(secretsByName))
	for _, value := range secretsByName {
		if value == "" {
			continue
		}
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		secrets = append(secrets, []byte(value))
	}
	if len(secrets) == 0 {
		return nil
	}
	sort.Slice(secrets, func(i, j int) bool {
		return len(secrets[i]) > len(secrets[j])
	})
	return &serverRunnerSecretMatcher{secrets: secrets}
}

func newServerRunnerSecretMasker(matcher *serverRunnerSecretMatcher) *serverRunnerSecretMasker {
	if matcher == nil || len(matcher.secrets) == 0 {
		return nil
	}
	return &serverRunnerSecretMasker{matcher: matcher}
}

func (m *serverRunnerSecretMasker) Mask(payload []byte, final bool) []byte {
	if m == nil || m.matcher == nil || len(m.matcher.secrets) == 0 {
		return payload
	}
	if len(payload) > 0 {
		m.buffer = append(m.buffer, payload...)
	}
	var out []byte
	for len(m.buffer) > 0 {
		if secret, ok := m.fullSecretPrefix(); ok {
			out = append(out, serverRunnerMaskReplacement...)
			m.buffer = m.buffer[len(secret):]
			continue
		}
		if !final && m.possibleSecretPrefix() {
			break
		}
		out = append(out, m.buffer[0])
		m.buffer = m.buffer[1:]
	}
	return out
}

func (m *serverRunnerSecretMasker) fullSecretPrefix() ([]byte, bool) {
	for _, secret := range m.matcher.secrets {
		if bytes.HasPrefix(m.buffer, secret) {
			return secret, true
		}
	}
	return nil, false
}

func (m *serverRunnerSecretMasker) possibleSecretPrefix() bool {
	for _, secret := range m.matcher.secrets {
		if len(m.buffer) < len(secret) && bytes.HasPrefix(secret, m.buffer) {
			return true
		}
	}
	return false
}