package lightnode

// Upstream authentication. Reading streams at a relay requires a session,
// and sessions require an AuthAssertion signed by *some* key. The light node
// therefore holds one ephemeral Ed25519 key used purely as a transport
// credential: it signs nothing but its own login, grants nothing, and is
// regenerated on every start unless a seed directory is configured.
//
// The handshake must survive relays that run proof-of-work admission
// control: when the challenge endpoint answers 429 "proof of work required",
// the node solves the keyed BLAKE3 puzzle for the zero target exactly like
// any other client.

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"git.n1ko.dev/Niko/niko_trust/internal/identity/signer"
	"git.n1ko.dev/Niko/niko_trust/internal/pow"
	"git.n1ko.dev/Niko/niko_trust/internal/protocol"
)

type upstreamAuth struct {
	client *http.Client

	// key is the node's ephemeral transport credential. All private-key
	// handling stays inside internal/identity/signer, keeping the
	// source-scanned invariant intact: this key signs exactly one thing,
	// the node's own login assertion.
	key    *signer.Signer
	pubHex string

	// per-peer session state, guarded by the node mutex in practice; each
	// peer gets its own token.
	tokens map[string]string
}

func newUpstreamAuth(seedDir string, client *http.Client) (*upstreamAuth, error) {
	if seedDir != "" {
		path := seedDir + "/transport_key.seed"
		// 32 bytes = ed25519 seed size; signer.FromSeed validates the rest.
		if seed, err := os.ReadFile(path); err == nil && len(seed) == 32 {
			signerKey, err := signer.FromSeed(seed)
			if err != nil {
				return nil, err
			}
			return &upstreamAuth{
				client: client,
				key:    signerKey,
				pubHex: hex.EncodeToString(signerKey.Public()),
				tokens: make(map[string]string),
			}, nil
		}
	}
	gen, err := signer.Generate()
	if err != nil {
		return nil, err
	}
	if seedDir != "" {
		if err := os.WriteFile(seedDir+"/transport_key.seed", gen.Seed(), 0o600); err != nil {
			return nil, err
		}
	}
	return &upstreamAuth{
		client: client,
		key:    gen,
		pubHex: hex.EncodeToString(gen.Public()),
		tokens: make(map[string]string),
	}, nil
}

var (
	errAuthFailed = errors.New("lightnode: upstream authentication failed")
)

// session returns a live bearer token for the peer, performing (and caching)
// the challenge/assert dance when needed. A 401 downstream invalidates the
// cached token once.
func (a *upstreamAuth) session(ctx context.Context, peer string) (string, error) {
	if tok := a.tokens[peer]; tok != "" {
		return tok, nil
	}
	tok, err := a.authenticate(ctx, peer)
	if err != nil {
		return "", err
	}
	a.tokens[peer] = tok
	return tok, nil
}

func (a *upstreamAuth) invalidate(peer string) { delete(a.tokens, peer) }

func (a *upstreamAuth) authenticate(ctx context.Context, peer string) (string, error) {
	body, status, err := postJSON(ctx, a.client, peer+"/v1/auth/challenge", nil)
	if err == nil && status == http.StatusTooManyRequests && bytes.Contains(body, []byte("proof of work")) {
		// Admission control is on: fetch a PoW challenge, solve for the
		// zero target, retry with the proof attached.
		powBody, pstatus, perr := postJSON(ctx, a.client, peer+"/v1/pow/challenge", map[string]string{"purpose": "auth"})
		if perr != nil || pstatus != http.StatusOK {
			return "", fmt.Errorf("%w: pow challenge", errAuthFailed)
		}
		var pc struct {
			Key        string `json:"key"`
			Difficulty int    `json:"difficulty"`
		}
		json.Unmarshal(powBody, &pc)
		key, err := pow.ParseKey(pc.Key)
		if err != nil {
			return "", errAuthFailed
		}
		counter, ok := pow.Solve(key, [pow.TargetSize]byte{}, pc.Difficulty)
		if !ok {
			return "", errAuthFailed
		}
		body, status, err = postJSON(ctx, a.client, peer+"/v1/auth/challenge",
			map[string]any{"pow": map[string]any{"key": pc.Key, "counter": counter}})
	}
	if err != nil || status != http.StatusOK {
		return "", fmt.Errorf("%w: challenge status %d", errAuthFailed, status)
	}
	var ch struct {
		Challenge string `json:"challenge"`
	}
	json.Unmarshal(body, &ch)
	chBytes, err := hex.DecodeString(ch.Challenge)
	if err != nil || len(chBytes) != 32 {
		return "", errAuthFailed
	}

	audience, err := a.audience(ctx, peer)
	if err != nil {
		return "", err
	}
	a2 := &protocol.AuthAssertion{
		PubKey:    a.key.Public(),
		Challenge: chBytes,
		Scope:     "read",
		Audience:  audience,
		CreatedAt: uint64(time.Now().Unix()),
	}
	tb, err := protocol.EncodeAuthAssertion(a2)
	if err != nil {
		return "", errAuthFailed
	}
	sig := a.key.Sign(tb)
	envBody, _, err := postJSON(ctx, a.client, peer+"/v1/auth/assert",
		map[string]any{"tce": base64.StdEncoding.EncodeToString(tb),
			"signature": base64.StdEncoding.EncodeToString(sig)})
	if err != nil || !bytes.Contains(envBody, []byte("session_token")) {
		return "", errAuthFailed
	}
	var out struct {
		SessionToken string `json:"session_token"`
	}
	json.Unmarshal(envBody, &out)
	if out.SessionToken == "" {
		return "", errAuthFailed
	}
	return out.SessionToken, nil
}

func postJSON(ctx context.Context, client *http.Client, url string, body any) ([]byte, int, error) {
	var rd io.Reader
	if body != nil {
		raw, _ := json.Marshal(body)
		rd = bytes.NewReader(raw)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, rd)
	if err != nil {
		return nil, 0, err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<16))
	return raw, resp.StatusCode, err
}

var _ = base64.StdEncoding

// audience fetches and caches the relay's audience binding.
func (a *upstreamAuth) audience(ctx context.Context, peer string) (string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, peer+"/v1/config", nil)
	if err != nil {
		return "", err
	}
	resp, err := a.client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	var cfg struct {
		Audience string `json:"audience"`
	}
	json.NewDecoder(io.LimitReader(resp.Body, 1<<16)).Decode(&cfg)
	if cfg.Audience == "" {
		return "", errAuthFailed
	}
	return cfg.Audience, nil
}