// Package agent is the control loop: pull an envelope, prove it came from a paired phone, decrypt
// it, decide, execute, and answer.
//
// The ordering in Handle is the security property and is asserted by tests rather than described in
// comments alone: nothing that depends on the payload happens before the signature verifies, and
// nothing that touches SimAdmin happens before the policy says yes.
package agent

import (
	"context"
	"crypto/ed25519"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"fun.chick/msa-agent/internal/policy"
	"fun.chick/msa-agent/internal/runtime"
	"fun.chick/msa-agent/internal/wire"
)

// Keys is the agent's own key material, loaded once at startup.
type Keys struct {
	// Sign signs responses and the handshake challenge. The phone verifies responses against its
	// copy of the public half, so a relay cannot forge an answer.
	Sign ed25519.PrivateKey
	// Kx is the raw 32-byte X25519 private key the phone encrypts requests to. It is raw on
	// purpose: clamping is the curve library's job, and a stored key that has already been clamped
	// would be a second, non-standard encoding to keep in sync across three languages.
	Kx []byte
}

// PublicKx derives the X25519 public key to publish at pairing time.
//
// Called rather than stored, so the published value cannot disagree with the private key it was
// derived from. A mismatched pair would fail as an unreadable request with no useful cause.
func (k Keys) PublicKx() ([]byte, error) { return wire.PublicFrom(k.Kx) }

// Devices answers "is this phone allowed to command this node".
//
// It is an interface because the first implementation is a static list from the config file and the
// intended one is a snapshot pulled from the relay, and the loop must not change when that arrives.
type Devices interface {
	// SigningKey returns the Ed25519 public key for app, or nil if it is not paired.
	SigningKey(app string) ed25519.PublicKey
	// Len reports how many devices are trusted. Logged at startup and on reload.
	Len() int
}

// StaticDevices is a paired-device list fixed at startup.
type StaticDevices map[string]ed25519.PublicKey

func (d StaticDevices) SigningKey(app string) ed25519.PublicKey { return d[app] }
func (d StaticDevices) Len() int                                { return len(d) }

// Error classes the agent reports to the phone, inside the encrypted response so the relay cannot
// read or rewrite them.
const (
	OutcomeBadEnvelope = "bad_envelope"
	OutcomeBadSig      = "bad_signature"
	OutcomeNoDevice    = "unknown_device"
	OutcomeExpired     = "expired"
	OutcomeReplay      = "replay"
	OutcomePolicy      = "policy_denied"
	OutcomeNoConfirm   = "confirmation_missing"
	OutcomeBadConfirm  = "confirmation_mismatch"
	OutcomeRateLimited = "rate_limited"
	OutcomeTooLarge    = "payload_too_large"
	OutcomeUpstream    = "upstream_error"
)

// Agent owns the loop over one node.
type Agent struct {
	nodeID  string
	keys    Keys
	devices Devices
	pol     *policy.Policy
	exec    *runtime.Executor
	now     func() time.Time
}

// New wires the pieces. It returns an error rather than starting degraded: an agent that runs with
// no policy or no devices would accept nothing and look healthy while doing so.
func New(nodeID string, keys Keys, devices Devices, pol *policy.Policy, exec *runtime.Executor) (*Agent, error) {
	if nodeID == "" {
		return nil, errors.New("agent: node id is empty")
	}
	if len(keys.Sign) != ed25519.PrivateKeySize {
		return nil, errors.New("agent: signing key is the wrong size")
	}
	if len(keys.Kx) != wire.KeyBytes {
		return nil, fmt.Errorf("agent: x25519 key must be %d bytes", wire.KeyBytes)
	}
	if devices == nil || devices.Len() == 0 {
		return nil, errors.New("agent: no paired devices, so nothing could ever be authorised")
	}
	if pol == nil || exec == nil {
		return nil, errors.New("agent: policy and executor are required")
	}
	return &Agent{nodeID: nodeID, keys: keys, devices: devices, pol: pol, exec: exec, now: time.Now}, nil
}

// Handle processes one request envelope and returns the response envelope ready to post back.
//
// It never returns an error for a bad request. A request that fails any check is answered with a
// refusal inside the encrypted response, for two reasons: the relay must not be able to read why a
// command failed, and the phone must be able to tell the user what to do about it. A transport-level
// error would put that explanation in the relay's hands instead.
func (a *Agent) Handle(ctx context.Context, raw []byte) (json.RawMessage, error) {
	var envelope wire.Envelope
	if err := json.Unmarshal(raw, &envelope); err != nil {
		return nil, fmt.Errorf("agent: envelope: %w", err)
	}
	// The node is checked before the signature, and this is the one exception to "verify first": an
	// envelope addressed to another node is not for us at all, and answering it would leak that this
	// agent is alive and knows the difference.
	if envelope.Node != a.nodeID {
		return nil, fmt.Errorf("agent: envelope is addressed to %q, not %q", envelope.Node, a.nodeID)
	}

	device := a.devices.SigningKey(envelope.App)
	if device == nil {
		return a.answer(envelope, OutcomeNoDevice, 403, nil)
	}
	if err := wire.VerifyRequest(envelope, device); err != nil {
		return a.answer(envelope, OutcomeBadSig, 401, nil)
	}
	// Only now is the payload touched. Decrypting before verifying would let an unauthenticated
	// caller choose the bytes the AEAD processes, which is exactly what an unauthenticated AEAD
	// oracle is.
	shared, err := wire.SharedSecret(a.keys.Kx, envelope.EK)
	if err != nil {
		return a.answer(envelope, OutcomeBadEnvelope, 400, nil)
	}
	key, iv, err := wire.DeriveKey(shared, envelope.Node, envelope.Req, false)
	if err != nil {
		return a.answer(envelope, OutcomeBadEnvelope, 400, nil)
	}
	opened, err := wire.Open(key, iv, envelope.P)
	if err != nil {
		return a.answer(envelope, OutcomeBadEnvelope, 400, nil)
	}
	var inner wire.InnerRequest
	if err := json.Unmarshal(opened, &inner); err != nil {
		return a.answer(envelope, OutcomeBadEnvelope, 400, nil)
	}

	result, refusal := a.exec.Execute(ctx, envelope, inner)
	if refusal != nil {
		return a.answer(envelope, refusal.Class, statusFor(refusal.Class), nil)
	}
	return a.answer(envelope, "ok", result.Status, result.Body)
}

// answer encrypts, signs and serialises the response.
//
// The ephemeral key pair is generated per response. The request direction is ephemeral-static, so
// generating a fresh one here is what keeps the two directions from sharing a key.
func (a *Agent) answer(request wire.Envelope, outcome string, status int, body []byte) (json.RawMessage, error) {
	ephemeral, err := wire.NewEphemeral()
	if err != nil {
		return nil, err
	}
	// The response is encrypted to the phone's per-request ephemeral public key, the one that
	// arrived in ek and is covered by the signature just verified. That keeps this direction
	// ephemeral-ephemeral rather than falling back to a key the phone published once at pairing
	// time, and it needs no extra round trip: ek is right there in the request.
	shared, err := ephemeral.Exchange(request.EK)
	if err != nil {
		return nil, err
	}
	key, iv, err := wire.DeriveKey(shared, request.Node, request.Req, true)
	if err != nil {
		return nil, err
	}
	inner, err := wire.MarshalInner(wire.InnerResponse{Status: status, Body: body, Outcome: outcome})
	if err != nil {
		return nil, err
	}
	sealed, err := wire.Seal(key, iv, inner)
	if err != nil {
		return nil, err
	}
	response := wire.Envelope{
		V:    wire.Version,
		Node: request.Node,
		Req:  request.Req,
		TS:   uint64(a.now().Unix()),
		EK:   ephemeral.Public(),
		Alg:  wire.Suite,
		P:    sealed,
	}
	if err := wire.SignResponse(&response, a.keys.Sign); err != nil {
		return nil, err
	}
	encoded, err := json.Marshal(response)
	if err != nil {
		return nil, err
	}
	return json.RawMessage(encoded), nil
}

func statusFor(class string) int {
	switch class {
	case runtime.ClassExpired, runtime.ClassReplay:
		return 408
	case runtime.ClassRateLimited:
		return 429
	case runtime.ClassTooLarge:
		return 413
	case runtime.ClassUpstream:
		return 502
	default:
		return 403
	}
}