package main

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"syscall"
	"time"

	"github.com/soypat/cyw43439"
	"github.com/soypat/lneto"
	"github.com/soypat/lneto/ethernet"
	"github.com/soypat/lneto/x/xnet"
)

const (
	pollTime     = 5 * time.Millisecond
	protoTimeout = 5 * time.Second
	protoRetries = 3

	tcpBufsize            = 2048
	tcpPacketQueueSize    = 4
	tcpConnPoolSize       = 5
	tcpEstablishedTimeout = 4 * time.Second
	tcpCloseTimeout       = protoTimeout
)

var nanotime = func() int64 {
	return time.Now().UnixNano()
}

// CywAdapter bridges the cyw43439 driver to the lneto interface requirements.
type CywAdapter struct {
	dev    *cyw43439.Device
	rxBuf  [1514]byte
	rxLen  int
	hasPkt bool
}

func NewCywAdapter(dev *cyw43439.Device) *CywAdapter {
	adapter := &CywAdapter{dev: dev}

	dev.RecvEthHandle(func(pkt []byte) error {
		if !adapter.hasPkt && len(pkt) <= len(adapter.rxBuf) {
			copy(adapter.rxBuf[:], pkt)
			adapter.rxLen = len(pkt)
			adapter.hasPkt = true
		}
		return nil
	})

	return adapter
}

// In CywAdapter:
func (a *CywAdapter) PollHardware() error {
	_, err := a.dev.TryPoll()
	return err
}

func (a *CywAdapter) SendEth(frame []byte) error {
	return a.dev.SendEth(frame)
}

func (a *CywAdapter) RecvEth(dst []byte) (int, error) {
	if !a.hasPkt {
		return 0, nil
	}
	n := copy(dst, a.rxBuf[:a.rxLen])
	a.hasPkt = false
	a.rxLen = 0
	return n, nil
}

func (a *CywAdapter) HardwareAddress6() ([6]byte, error) {
	return a.dev.HardwareAddr6()
}

func (a *CywAdapter) MaxFrameLength() (int, error) {
	return 1514, nil
}

type NetworkStack struct {
	adapter *CywAdapter
	stack   *xnet.StackAsync
	gostack xnet.StackGo
	localIP netip.Addr
}

func InitNetwork(ssid, pass string) (*NetworkStack, error) {
	println("Initializing CYW43439 Wi-Fi hardware...")
	dev := cyw43439.NewPicoWDevice()
	cfg := cyw43439.DefaultWifiConfig()
	if err := dev.Init(cfg); err != nil {
		return nil, fmt.Errorf("device init failed: %w", err)
	}

	println("Associating with SSID:", ssid)
	if err := dev.JoinWPA2(ssid, pass); err != nil {
		return nil, fmt.Errorf("wifi association failed: %w", err)
	}
	println("Wi-Fi associated.")

	adapter := NewCywAdapter(dev)
	hwaddr, err := adapter.HardwareAddress6()
	if err != nil {
		return nil, fmt.Errorf("read MAC: %w", err)
	}
	framelen, err := adapter.MaxFrameLength()
	if err != nil {
		return nil, fmt.Errorf("max frame len: %w", err)
	}

	stack := &xnet.StackAsync{}
	err = stack.Reset(xnet.StackConfig{
		Hostname:          "pico2w-bme280",
		RandSeed:          time.Now().UnixNano(),
		MaxActiveTCPPorts: 2,
		MTU:               uint16(framelen - ethernet.MaxOverheadSize),
		HardwareAddress:   hwaddr,
	})
	if err != nil {
		return nil, fmt.Errorf("stack config reset: %w", err)
	}

	// Start background frame pump
	ctx := context.Background()
	go stackLoop(ctx, stack, adapter)

	println("Acquiring IP via DHCP...")
	rstack := stack.StackRetrying(stackBackoff)
	results, err := rstack.DoDHCPv4([4]byte{}, protoTimeout, protoRetries)
	if err != nil {
		return nil, fmt.Errorf("DHCP failed: %w", err)
	}

	err = stack.AssimilateDHCPResults(results)
	if err != nil {
		return nil, fmt.Errorf("assimilate DHCP failed: %w", err)
	}

	println("Resolving router MAC...")
	gateway, err := rstack.DoResolveHardwareAddress6(results.Router, protoTimeout, protoRetries)
	if err != nil {
		return nil, fmt.Errorf("resolving router MAC failed: %w", err)
	}
	stack.SetGatewayHardwareAddr(gateway)

	localIP := netip.AddrFrom4(results.AssignedAddr4)
	println("DHCP lease assigned! IP:", localIP.String())

	gostack := stack.StackBlocking(stackBackoff).StackGo(xnet.StackGoConfig{
		ListenerPoolConfig: xnet.TCPPoolConfig{
			PoolSize:           tcpConnPoolSize,
			QueueSize:          tcpPacketQueueSize,
			TxBufSize:          tcpBufsize,
			RxBufSize:          tcpBufsize,
			NanoTime:           nanotime,
			EstablishedTimeout: tcpEstablishedTimeout,
			ClosingTimeout:     tcpCloseTimeout,
			NewBackoff:         func() lneto.BackoffStrategy { return tcpBackoff },
		},
	})

	return &NetworkStack{
		adapter: adapter,
		stack:   stack,
		gostack: gostack,
		localIP: localIP,
	}, nil
}

func (ns *NetworkStack) DialTCP(ctx context.Context, hostIP string, port uint16) (net.Conn, error) {
	rIP, err := netip.ParseAddr(hostIP)
	if err != nil {
		return nil, fmt.Errorf("invalid host IPv4: %w", err)
	}

	laddr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(ns.localIP, 0))
	raddr := net.TCPAddrFromAddrPort(netip.AddrPortFrom(rIP, port))

	const sockstream = 0x1
	c, err := ns.gostack.Socket(ctx, "tcp", syscall.AF_INET, sockstream, laddr, raddr)
	if err != nil {
		return nil, fmt.Errorf("socket dial: %w", err)
	}

	conn, ok := c.(net.Conn)
	if !ok {
		return nil, errors.New("socket did not return a stream connection")
	}

	return conn, nil
}

func stackLoop(ctx context.Context, stack *xnet.StackAsync, adapter *CywAdapter) {
	frameLength, _ := adapter.MaxFrameLength()
	buf := make([]byte, frameLength)

	for ctx.Err() == nil {
		// 1. Pump the CYW43439 hardware over SPI to trigger RecvEthHandle
		_ = adapter.PollHardware()

		// 2. Ingress: read from adapter into lneto
		nread, err := adapter.RecvEth(buf[:])
		if err != nil {
			println("recv err:", err.Error())
		} else if nread > 0 {
			err = stack.IngressEthernet(buf[:nread])
			if err != nil && err != lneto.ErrPacketDrop {
				println("ingress err:", err.Error())
			}
		}

		// 3. Egress: send out any frames generated by lneto
		nwrite, err := stack.EgressEthernet(buf[:])
		if err != nil {
			println("egress err:", err.Error())
		} else if nwrite > 0 {
			if err := adapter.SendEth(buf[:nwrite]); err != nil {
				println("send eth err:", err.Error())
			}
		}

		if nwrite == 0 && nread == 0 {
			time.Sleep(pollTime)
		}
	}
}

func stackBackoff(consecutiveBackoffs uint) time.Duration {
	if consecutiveBackoffs < 10 {
		return time.Millisecond
	}
	return 10 * time.Millisecond
}

func tcpBackoff(consecutiveBackoffs uint) time.Duration {
	const (
		minWait  = uint32(time.Microsecond)
		maxWait  = 5 * uint32(time.Millisecond)
		maxShift = 22
	)
	shifted := minWait << min(consecutiveBackoffs, maxShift)
	wait := min(shifted, maxWait)
	return time.Duration(wait)
}

// FetchNTP requests the current time from an NTP server IPv4 address over UDP (port 123).
func (ns *NetworkStack) FetchNTP(serverIP string) (time.Time, error) {
	ip, err := netip.ParseAddr(serverIP)
	if err != nil || !ip.Is4() {
		return time.Time{}, fmt.Errorf("invalid IPv4 address: %w", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), protoTimeout)
	defer cancel()

	laddr := net.UDPAddrFromAddrPort(netip.AddrPortFrom(ns.localIP, 0))
	raddr := net.UDPAddrFromAddrPort(netip.AddrPortFrom(ip, 123))

	const sockDgram = 0x2
	c, err := ns.gostack.Socket(ctx, "udp", syscall.AF_INET, sockDgram, laddr, raddr)
	if err != nil {
		return time.Time{}, fmt.Errorf("udp socket dial: %w", err)
	}
	conn, ok := c.(net.Conn)
	if !ok {
		return time.Time{}, errors.New("socket did not return a net.Conn")
	}
	defer conn.Close()

	_ = conn.SetDeadline(time.Now().Add(protoTimeout))

	// Standard NTP request packet (48 bytes, leap indicator 0, version 4, mode 3 client = 0x23)
	req := make([]byte, 48)
	req[0] = 0x23

	if _, err := conn.Write(req); err != nil {
		return time.Time{}, fmt.Errorf("ntp write error: %w", err)
	}

	resp := make([]byte, 48)
	n, err := conn.Read(resp)
	if err != nil {
		return time.Time{}, fmt.Errorf("ntp read error: %w", err)
	}
	if n < 48 {
		return time.Time{}, fmt.Errorf("ntp response too short: %d bytes", n)
	}

	// Transmit timestamp seconds are at bytes 40-43 (big-endian)
	secs := uint32(resp[40])<<24 | uint32(resp[41])<<16 | uint32(resp[42])<<8 | uint32(resp[43])
	fraction := uint32(resp[44])<<24 | uint32(resp[45])<<16 | uint32(resp[46])<<8 | uint32(resp[47])

	// NTP epoch begins Jan 1, 1900. Unix epoch begins Jan 1, 1970.
	// Difference is 2,208,988,800 seconds (70 years).
	const ntpEpochOffset = 2208988800
	if secs < ntpEpochOffset {
		return time.Time{}, errors.New("invalid NTP timestamp received")
	}

	unixSecs := int64(secs - ntpEpochOffset)
	nanos := (int64(fraction) * 1e9) >> 32

	return time.Unix(unixSecs, nanos), nil
}