package main

import (
        "context"
        "net"
        "net/netip"
        "time"

        "github.com/insomniacslk/dhcp/dhcpv4"
        "github.com/insomniacslk/dhcp/dhcpv6"
        "github.com/insomniacslk/dhcp/dhcpv6/nclient6"
        "github.com/insomniacslk/dhcp/iana"

	gen "golang.zx2c4.com/wireguard/android/gen"
        "google.golang.org/protobuf/types/known/durationpb"
)

const (
        ENABLE_PD = true
)

type dhcp struct {
        fqdn string
        hwAddr net.HardwareAddr
        conn *net.UDPConn
        serverAddr net.UDPAddr
        client *nclient6.Client
        linkAddr net.IP
        peerAddr net.IP
}

func newClientIDOpt(duid dhcpv6.DUID) dhcpv4.Option {
	iaid := []byte{0, 0, 0, 3}
	ident := []byte{255}                     // Type IAID+DUID
	ident = append(ident, iaid...)           // IAID
	ident = append(ident, duid.ToBytes()...) // DUID
	return dhcpv4.OptClientIdentifier(ident)
}

func getDuid(hwAddr net.HardwareAddr) dhcpv6.DUID {
	duid := &dhcpv6.DUIDLLT{
		HWType:        iana.HWTypeEthernet,
		Time:          uint32(time.Now().Unix()),
		LinkLayerAddr: hwAddr,
	}

        return duid
}

func (d *dhcp) encapSendAndRead(ctx context.Context, msg *dhcpv6.Message, match nclient6.Matcher) (*dhcpv6.Message, error) {
	packet, err := dhcpv6.EncapsulateRelay(msg, dhcpv6.MessageTypeRelayForward, d.linkAddr, d.peerAddr)
	if err != nil {
		return nil, err
	}

        relay, err := d.client.SendAndReadRelay(ctx, &d.serverAddr, packet, match)
	if err != nil {
		return nil, err
	}

        inner, err := relay.GetInnerMessage()
        if err != nil {
                return nil, err
        }

        return inner, nil
}

// isRelayMessageType returns a matcher that checks for the message type.
func isRelayMessageType(t dhcpv6.MessageType, tt ...dhcpv6.MessageType) nclient6.Matcher {
	return func(p dhcpv6.DHCPv6) bool {
                inner, err := p.GetInnerMessage()
                if err != nil {
                        return false
                }
		if inner.Type() == t {
			return true
		}
		for _, mt := range tt {
			if inner.Type() == mt {
				return true
			}
		}
		return false
	}
}

// func New() *dhcp {
// }

func RunDhcp(ctx context.Context, laddr, raddr netip.Addr) ([]*gen.Lease, error) {
        d := &dhcp{}
        
        d.linkAddr = net.ParseIP("fe80::101")
	d.peerAddr = net.ParseIP("::1")

        // TODO generate hostname and hwAddr from public key
	hostName := "foobar"
	d.fqdn = hostName + ".m7n.se"
        d.hwAddr = []byte{41, 42, 43, 44, 45, 46}


        src := net.UDPAddr{IP: laddr.AsSlice(),
                Port: 0,        // Use non-restrict UDP source port
        }


        d.serverAddr = net.UDPAddr{IP: raddr.AsSlice(),
                Port: 547,
        }

        err := d.Start(&src)
        if err != nil {
                return nil, err
        }

        defer d.Close()

        reply, err := d.ObtainLease(ctx) // Use reply
        if err != nil {
                return nil, err
        }

        return getAddressesFromReply(reply), nil
}

func getAddressesFromReply(reply *dhcpv6.Message) []*gen.Lease{
        var leases []*gen.Lease = make([]*gen.Lease, 0, 4)

        if opt := reply.GetOneOption(dhcpv6.OptionIANA); opt != nil {
                iana := opt.(*dhcpv6.OptIANA)
                ianaOpts := iana.Options.Get(dhcpv6.OptionIAAddr)

                for _, opt := range ianaOpts {
                        addrOpt := opt.(*dhcpv6.OptIAAddress)
                        lease := &gen.Lease{
                                Address: &gen.InetAddress{
                                        Address: addrOpt.IPv6Addr,
                                },
                                PreferredLifetime: durationpb.New(addrOpt.PreferredLifetime),
                                ValidLifetime: durationpb.New(addrOpt.ValidLifetime),
                        }
                        leases = append(leases, lease)
                }
        }

        return leases
}

func (d *dhcp) Start(localAddr *net.UDPAddr) error {
        conn, err := net.ListenUDP("udp6", localAddr)
	if err != nil {
                return err
        }
        
	d.client, err = nclient6.NewWithConn(conn, d.hwAddr, nclient6.WithDebugLogger())
	return err
}

func (d *dhcp) Close() error {
        err := d.client.Close()

        d.client = nil
        return err
}

func (d *dhcp) ObtainLease(ctx context.Context) (*dhcpv6.Message, error){
        duidOpt := dhcpv6.WithClientID(getDuid(d.hwAddr))
	fqdnOpt := dhcpv6.WithFQDN(0x1, d.fqdn)

        solicit, err := dhcpv6.NewSolicit(d.hwAddr, duidOpt, fqdnOpt, dhcpv6.WithRapidCommit)
        if err != nil {
                return nil, err
        }

        msg, err := d.encapSendAndRead(ctx, solicit, isRelayMessageType(dhcpv6.MessageTypeReply, dhcpv6.MessageTypeAdvertise))
        if err != nil {
                return nil, err
        }

        if msg.Type() == dhcpv6.MessageTypeReply {
		// We got RapidCommitted.
		return msg, nil
        }

        // We didn't get RapidCommitted. Request regular lease.
        req, err := dhcpv6.NewRequestFromAdvertise(msg, fqdnOpt)
        if err != nil {
                return nil, err
        }

        return d.encapSendAndRead(ctx, req, nil)
}