package main

import (
	"crypto/ecdh"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net/netip"
	"strings"
)

type vyosOpts struct {
	Hostname   string
	WANIf      string // VyOS interface facing the ISP (eth0)
	LANIf      string // VyOS interface facing tenant VMs (eth1)
	TenantCIDR string // e.g. 10.77.0.0/24
	DNS        string
	WireGuard  bool
	WGPort     int
	WGCIDR     string // e.g. 10.99.0.1/24 (router address)
}

type vyosResult struct {
	Commands []string
	WGPublic string // router public key (private key is only in Commands)
}

// vyosCommands renders `set` commands for VyOS 1.4 (sagitta).
func vyosCommands(o vyosOpts) (*vyosResult, error) {
	p, err := netip.ParsePrefix(o.TenantCIDR)
	if err != nil {
		return nil, fmt.Errorf("tenant cidr: %w", err)
	}
	if !p.Addr().Is4() || p.Bits() < 16 || p.Bits() > 24 {
		return nil, fmt.Errorf("tenant cidr must be IPv4 between /16 and /24, got %s", o.TenantCIDR)
	}
	p = p.Masked()
	gw := addrPlus(p.Addr(), 1)
	start, stop := addrPlus(p.Addr(), 100), addrPlus(p.Addr(), 250)
	sub := fmt.Sprintf("service dhcp-server shared-network-name TENANT subnet %s", p)

	c := []string{
		"system host-name " + o.Hostname,
		fmt.Sprintf("interfaces ethernet %s description WAN", o.WANIf),
		fmt.Sprintf("interfaces ethernet %s address dhcp", o.WANIf),
		fmt.Sprintf("interfaces ethernet %s description TENANT", o.LANIf),
		fmt.Sprintf("interfaces ethernet %s address %s/%d", o.LANIf, gw, p.Bits()),
		sub + " subnet-id 1",
		fmt.Sprintf("%s option default-router %s", sub, gw),
		fmt.Sprintf("%s option name-server %s", sub, o.DNS),
		fmt.Sprintf("%s range 0 start %s", sub, start),
		fmt.Sprintf("%s range 0 stop %s", sub, stop),
		// Tenant internet access via source NAT.
		fmt.Sprintf("nat source rule 100 outbound-interface name %s", o.WANIf),
		fmt.Sprintf("nat source rule 100 source address %s", p),
		"nat source rule 100 translation address masquerade",
		// Firewall: stateful, tenants may reach the internet but not private networks.
		"firewall global-options state-policy established action accept",
		"firewall global-options state-policy related action accept",
		"firewall global-options state-policy invalid action drop",
		"firewall group network-group PRIVATE network 10.0.0.0/8",
		"firewall group network-group PRIVATE network 172.16.0.0/12",
		"firewall group network-group PRIVATE network 192.168.0.0/16",
		"firewall ipv4 forward filter default-action drop",
		fmt.Sprintf("firewall ipv4 forward filter rule 10 inbound-interface name %s", o.LANIf),
		"firewall ipv4 forward filter rule 10 destination group network-group PRIVATE",
		"firewall ipv4 forward filter rule 10 action drop",
		fmt.Sprintf("firewall ipv4 forward filter rule 20 inbound-interface name %s", o.LANIf),
		fmt.Sprintf("firewall ipv4 forward filter rule 20 outbound-interface name %s", o.WANIf),
		"firewall ipv4 forward filter rule 20 action accept",
	}
	res := &vyosResult{}
	if o.WireGuard {
		priv, pub, err := wgKeypair()
		if err != nil {
			return nil, err
		}
		res.WGPublic = pub
		c = append(c,
			"interfaces wireguard wg0 address "+o.WGCIDR,
			fmt.Sprintf("interfaces wireguard wg0 port %d", o.WGPort),
			"interfaces wireguard wg0 private-key "+priv,
			"# add one peer per tenant:",
			"# interfaces wireguard wg0 peer TENANT1 public-key <CLIENT_PUBLIC_KEY>",
			"# interfaces wireguard wg0 peer TENANT1 allowed-ips <CLIENT_TUNNEL_IP>/32",
			fmt.Sprintf("firewall ipv4 input filter rule 30 protocol udp"),
			fmt.Sprintf("firewall ipv4 input filter rule 30 destination port %d", o.WGPort),
			fmt.Sprintf("firewall ipv4 input filter rule 30 inbound-interface name %s", o.WANIf),
			"firewall ipv4 input filter rule 30 action accept",
		)
	}
	res.Commands = c
	return res, nil
}

func addrPlus(a netip.Addr, n int) netip.Addr {
	b := a.As4()
	v := uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
	v += uint32(n)
	return netip.AddrFrom4([4]byte{byte(v >> 24), byte(v >> 16), byte(v >> 8), byte(v)})
}

// wgKeypair returns a WireGuard (Curve25519) private/public key pair in base64.
func wgKeypair() (priv, pub string, err error) {
	k, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return "", "", err
	}
	return base64.StdEncoding.EncodeToString(k.Bytes()),
		base64.StdEncoding.EncodeToString(k.PublicKey().Bytes()), nil
}

// vbashScript is pasted or run on the VyOS console after `install image`.
func vbashScript(cmds []string) string {
	var b strings.Builder
	b.WriteString("#!/bin/vbash\nsource /opt/vyatta/etc/functions/script-template\nconfigure\n")
	for _, c := range cmds {
		if strings.HasPrefix(c, "#") {
			b.WriteString(c + "\n")
		} else {
			b.WriteString("set " + c + "\n")
		}
	}
	b.WriteString("commit\nsave\nexit\n")
	return b.String()
}

// cloudInitUserData is for VyOS cloud-init capable images (NoCloud).
func cloudInitUserData(cmds []string) string {
	var b strings.Builder
	b.WriteString("#cloud-config\nvyos_config_commands:\n")
	for _, c := range cmds {
		if !strings.HasPrefix(c, "#") {
			b.WriteString("  - set " + c + "\n")
		}
	}
	return b.String()
}