package main

import (
	"crypto/ecdh"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"regexp"
)

// keys.json stores Reality keypairs by key id:
// { "<keyId>": { "privateKey": "...", "publicKey": "...", "createdAt": "..." } }
// Entries are kept as raw JSON so keys written by the Node agent (or any
// future fields) survive a save by this agent untouched.

type realityKeyPair struct {
	PrivateKey string `json:"privateKey"`
	PublicKey  string `json:"publicKey"`
	CreatedAt  string `json:"createdAt"`
}

func loadKeys() (map[string]json.RawMessage, error) {
	data, err := os.ReadFile(cfg.KeyFile)
	if err != nil {
		if os.IsNotExist(err) {
			return map[string]json.RawMessage{}, nil
		}
		return nil, err
	}
	keys := map[string]json.RawMessage{}
	if err := json.Unmarshal(data, &keys); err != nil {
		return nil, err
	}
	if keys == nil {
		keys = map[string]json.RawMessage{}
	}
	return keys, nil
}

func saveKeys(keys map[string]json.RawMessage) error {
	body, err := marshalIndentJSON(keys)
	if err != nil {
		return err
	}
	tmp := cfg.KeyFile + ".tmp"
	if err := os.WriteFile(tmp, []byte(body+"\n"), 0o600); err != nil {
		return err
	}
	if err := os.Rename(tmp, cfg.KeyFile); err != nil {
		return err
	}
	return os.Chmod(cfg.KeyFile, 0o600)
}

// realityKey returns the stored keypair for keyId or generates and persists a
// new X25519 pair. Encoding is base64url WITHOUT padding (base64.RawURLEncoding
// of the raw 32-byte keys) — byte-identical to the JWK `d`/`x` values Node's
// crypto.generateKeyPairSync('x25519') produces, and to the output of
// `sing-box generate reality-keypair`. Converting to padded standard base64
// here was a production bug; do not reintroduce it.
func realityKey(keyID string) (realityKeyPair, error) {
	keys, err := loadKeys()
	if err != nil {
		return realityKeyPair{}, err
	}
	if raw, ok := keys[keyID]; ok {
		var pair realityKeyPair
		if err := json.Unmarshal(raw, &pair); err == nil && pair.PrivateKey != "" {
			return pair, nil
		}
	}
	priv, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return realityKeyPair{}, fmt.Errorf("X25519 密钥生成失败: %w", err)
	}
	pair := realityKeyPair{
		PrivateKey: base64.RawURLEncoding.EncodeToString(priv.Bytes()),
		PublicKey:  base64.RawURLEncoding.EncodeToString(priv.PublicKey().Bytes()),
		CreatedAt:  isoNow(),
	}
	encoded, err := json.Marshal(pair)
	if err != nil {
		return realityKeyPair{}, err
	}
	keys[keyID] = encoded
	if err := saveKeys(keys); err != nil {
		return realityKeyPair{}, err
	}
	return pair, nil
}

type artifacts struct {
	RealityPublicKey *string `json:"realityPublicKey,omitempty"`
}

func (a *artifacts) toMap() map[string]interface{} {
	out := map[string]interface{}{}
	if a.RealityPublicKey != nil {
		out["realityPublicKey"] = *a.RealityPublicKey
	}
	return out
}

var realityPlaceholderRe = regexp.MustCompile(`^\$\{REALITY_PRIVATE:(.+)}$`)

// materialize deep-walks a job resource replacing whole-string
// "${REALITY_PRIVATE:<keyId>}" placeholders with the stored private key,
// recording the matching public key in artifacts — agent.js materialize().
func materialize(value interface{}, art *artifacts) (interface{}, error) {
	switch v := value.(type) {
	case []interface{}:
		out := make([]interface{}, len(v))
		for i, item := range v {
			replaced, err := materialize(item, art)
			if err != nil {
				return nil, err
			}
			out[i] = replaced
		}
		return out, nil
	case map[string]interface{}:
		out := make(map[string]interface{}, len(v))
		for key, item := range v {
			replaced, err := materialize(item, art)
			if err != nil {
				return nil, err
			}
			out[key] = replaced
		}
		return out, nil
	case string:
		match := realityPlaceholderRe.FindStringSubmatch(v)
		if match == nil {
			return v, nil
		}
		pair, err := realityKey(match[1])
		if err != nil {
			return nil, err
		}
		pub := pair.PublicKey
		art.RealityPublicKey = &pub
		return pair.PrivateKey, nil
	default:
		return value, nil
	}
}