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"`
	KeyMaterial      map[string]realityKeyPair `json:"keyMaterial,omitempty"`
	// used records every keypair resolved during materialize (never
	// serialized); applyResourceLocked diffs it against the payload-provided
	// keys to compute the backfill reported to the panel.
	used map[string]realityKeyPair
}

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

// persistKeyMaterial merges panel-provided keypairs into keys.json before a
// resource is materialized. The panel is the source of truth: a provided
// pair replaces a divergent local entry, so a node restored from the panel
// store converges back to the hosted keypair. keys.json stays 0600 (saveKeys).
func persistKeyMaterial(material map[string]realityKeyPair) error {
	if len(material) == 0 {
		return nil
	}
	keys, err := loadKeys()
	if err != nil {
		return err
	}
	changed := false
	for keyID, pair := range material {
		if pair.PrivateKey == "" || pair.PublicKey == "" {
			continue
		}
		stored := pair
		if raw, ok := keys[keyID]; ok {
			var existing realityKeyPair
			if err := json.Unmarshal(raw, &existing); err == nil {
				if existing.PrivateKey == pair.PrivateKey && existing.PublicKey == pair.PublicKey {
					continue
				}
				if stored.CreatedAt == "" {
					stored.CreatedAt = existing.CreatedAt
				}
			}
		}
		if stored.CreatedAt == "" {
			stored.CreatedAt = isoNow()
		}
		encoded, err := json.Marshal(stored)
		if err != nil {
			return err
		}
		keys[keyID] = encoded
		changed = true
	}
	if !changed {
		return nil
	}
	return saveKeys(keys)
}

// backfillKeyMaterial returns the keypairs this apply actually used that the
// panel did not provide in the payload. The panel stores them on job
// completion, so keys that predate panel hosting migrate into the store.
func backfillKeyMaterial(used, provided map[string]realityKeyPair) map[string]realityKeyPair {
	out := map[string]realityKeyPair{}
	for keyID, pair := range used {
		if _, ok := provided[keyID]; !ok {
			out[keyID] = pair
		}
	}
	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
		if art.used == nil {
			art.used = map[string]realityKeyPair{}
		}
		art.used[match[1]] = pair
		return pair.PrivateKey, nil
	default:
		return value, nil
	}
}