package peer

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/netip"
	"os"
	"path/filepath"

	"golang.org/x/crypto/nacl/sign"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"vppn/m"
)

// LocalState is the persisted identity for this peer, written on first run and
// loaded on every subsequent run.
type LocalState struct {
	PrivKey     wgtypes.Key
	SignKey     [64]byte // nacl/sign Ed25519 private key
	VPNIP       netip.Addr
	VPNNet      netip.Prefix
	WGPort      uint16
	IsRelay     bool
	IsPublic    bool
	LocalDomain string
}

// localStateJSON is the on-disk representation.
type localStateJSON struct {
	PrivKey     string       `json:"priv_key"`     // standard base64
	SignKey     string       `json:"sign_key"`     // standard base64
	VPNIP       netip.Addr   `json:"vpn_ip"`
	VPNNet      netip.Prefix `json:"vpn_net"`
	WGPort      uint16       `json:"wg_port"`
	IsRelay     bool         `json:"is_relay"`
	IsPublic    bool         `json:"is_public"`
	LocalDomain string       `json:"local_domain"`
}

// LoadOrInit loads LocalState from path, or registers with the hub and creates
// the file if it doesn't exist.
func LoadOrInit(statePath, hubURL, apiKey string) (LocalState, error) {
	if data, err := os.ReadFile(statePath); err == nil {
		return parseLocalState(data)
	}

	privKey, err := wgtypes.GeneratePrivateKey()
	if err != nil {
		return LocalState{}, fmt.Errorf("generate key: %w", err)
	}

	state, err := initFromHub(hubURL, apiKey, privKey)
	if err != nil {
		return LocalState{}, err
	}

	if err := saveLocalState(statePath, state); err != nil {
		return LocalState{}, fmt.Errorf("save state: %w", err)
	}
	return state, nil
}

func initFromHub(hubURL, apiKey string, privKey wgtypes.Key) (LocalState, error) {
	wgPubKey := privKey.PublicKey()

	signPubKey, signPrivKey, err := sign.GenerateKey(rand.Reader)
	if err != nil {
		return LocalState{}, fmt.Errorf("generate sign key: %w", err)
	}

	body, _ := json.Marshal(m.PeerInitArgs{
		WGPubKey:   wgPubKey[:],
		SignPubKey: signPubKey[:],
	})

	req, err := http.NewRequest(http.MethodPost, hubURL+"/peer/init/", bytes.NewReader(body))
	if err != nil {
		return LocalState{}, err
	}
	req.SetBasicAuth("", apiKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return LocalState{}, fmt.Errorf("hub init: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return LocalState{}, fmt.Errorf("hub init: HTTP %d", resp.StatusCode)
	}

	var r m.PeerInitResp
	if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
		return LocalState{}, fmt.Errorf("hub init decode: %w", err)
	}

	if len(r.Network) != 4 {
		return LocalState{}, fmt.Errorf("hub init: invalid network %v", r.Network)
	}

	netAddr := netip.AddrFrom4([4]byte(r.Network))
	octets := netAddr.As4()
	octets[3] = r.PeerIP
	vpnIP := netip.AddrFrom4(octets)
	vpnNet := netip.PrefixFrom(netAddr, 24)

	var isRelay, isPublic bool
	var wgPort uint16
	if self := r.NetworkState.Peers[r.PeerIP]; self != nil {
		isRelay = self.Relay
		isPublic = len(self.Addr4) > 0 || len(self.Addr6) > 0
		wgPort = self.Port
	}

	return LocalState{
		PrivKey:     privKey,
		SignKey:     *signPrivKey,
		VPNIP:       vpnIP,
		VPNNet:      vpnNet,
		WGPort:      wgPort,
		IsRelay:     isRelay,
		IsPublic:    isPublic,
		LocalDomain: r.LocalDomain,
	}, nil
}

func parseLocalState(data []byte) (LocalState, error) {
	var j localStateJSON
	if err := json.Unmarshal(data, &j); err != nil {
		return LocalState{}, fmt.Errorf("parse state: %w", err)
	}
	keyBytes, err := base64.StdEncoding.DecodeString(j.PrivKey)
	if err != nil {
		return LocalState{}, fmt.Errorf("decode key: %w", err)
	}
	key, err := wgtypes.NewKey(keyBytes)
	if err != nil {
		return LocalState{}, fmt.Errorf("invalid key: %w", err)
	}
	signKeyBytes, err := base64.StdEncoding.DecodeString(j.SignKey)
	if err != nil {
		return LocalState{}, fmt.Errorf("decode sign key: %w", err)
	}
	if len(signKeyBytes) != 64 {
		return LocalState{}, fmt.Errorf("invalid sign key length: %d", len(signKeyBytes))
	}
	var signKey [64]byte
	copy(signKey[:], signKeyBytes)
	return LocalState{
		PrivKey:     key,
		SignKey:     signKey,
		VPNIP:       j.VPNIP,
		VPNNet:      j.VPNNet,
		WGPort:      j.WGPort,
		IsRelay:     j.IsRelay,
		IsPublic:    j.IsPublic,
		LocalDomain: j.LocalDomain,
	}, nil
}

func saveLocalState(path string, s LocalState) error {
	j := localStateJSON{
		PrivKey:     base64.StdEncoding.EncodeToString(s.PrivKey[:]),
		SignKey:     base64.StdEncoding.EncodeToString(s.SignKey[:]),
		VPNIP:       s.VPNIP,
		VPNNet:      s.VPNNet,
		WGPort:      s.WGPort,
		IsRelay:     s.IsRelay,
		IsPublic:    s.IsPublic,
		LocalDomain: s.LocalDomain,
	}
	data, err := json.MarshalIndent(j, "", "  ")
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	return os.WriteFile(path, data, 0600)
}