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)
}