package main

import (
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"os"
	"sync"

	"baalvpn/shared"
)

type Peer struct {
	RealIP     string
	InternalIP string
}

var (
	peers      = map[string]*Peer{}
	peersMutex sync.Mutex
	nextIP     = 3
	conn       *net.UDPConn
)

func main() {
	if os.Geteuid() != 0 {
		panic("root permissions needed")
	}

	receivePackets()
}

func receivePackets() {
	log.Println("Listening for packets...")

	listenAddr, err := net.ResolveUDPAddr("udp", ":38000")
	if err != nil {
		panic(err)
	}
	conn, err = net.ListenUDP("udp", listenAddr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	for {
		buffer := make([]byte, 50000)
		n, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Println(err)
			continue
		}
		if n == 0 {
			continue
		}

		go receivePacket(buffer[:n], addr)
	}
}

func receivePacket(pkt []byte, addr *net.UDPAddr) {
	if len(pkt) < 4 {
		log.Println("packet too short")
		return
	}
	version := binary.LittleEndian.Uint16(pkt[0:])
	if version != shared.PROTO_VERSION {
		log.Println("mismatched packet version")
		return
	}

	pktType := binary.LittleEndian.Uint16(pkt[2:])

	peerIP := addr.IP.String()

	switch pktType {
	case shared.REQ_REGISTER:
		peersMutex.Lock()
		defer peersMutex.Unlock()

		// TODO
		internalIP := fmt.Sprintf("10.1.0.%d", nextIP)

		if peer, ok := peers[addr.IP.String()]; ok {
			internalIP = peer.InternalIP
		} else {
			peers[addr.IP.String()] = &Peer{
				RealIP:     peerIP,
				InternalIP: internalIP,
			}
			nextIP++
		}

		resp := shared.BuildPkt(shared.RESP_REGISTER, net.ParseIP(internalIP).To4())
		if _, err := conn.WriteToUDP(resp, addr); err != nil {
			log.Println(err)
		}
	case shared.REQ_DATA:
		log.Println("received packet from " + peerIP + ":")

		frame := pkt[4:]
		etherType := binary.BigEndian.Uint16(frame[12:14])
		payload := frame[14:]

		switch etherType {
		case 0x0800: // IPv4
			destIP := net.IP(payload[16:20]).String()
			if peer := getPeerByInternalIP(destIP); peer != nil {
				log.Println("relaying to", peer.RealIP)
				resp := shared.BuildPkt(shared.RESP_DATA, frame)
				sendTo(peer.RealIP, resp)
			} else {
				log.Println("unknown dest IP: " + destIP)
			}
		case 0x0806: // ARP
			log.Println("broadcasting ARP")
			resp := shared.BuildPkt(shared.RESP_DATA, frame)
			broadcast(resp, peerIP)
		case 0x86dd: // IPv6
			log.Println("IPv6") // TODO
		default:
			log.Printf("unhandled EtherType: 0x%x\n", etherType)
		}
	default:
		log.Println("unknown packet type")
	}
}

func broadcast(pkt []byte, except string) {
	peersMutex.Lock()
	defer peersMutex.Unlock()
	for _, p := range peers {
		if p.RealIP == except {
			continue
		}
		sendTo(p.RealIP, pkt)
	}
}

func sendTo(ip string, pkt []byte) {
	addr, err := net.ResolveUDPAddr("udp", ip+":38001")
	if err != nil {
		log.Println(err)
		return
	}
	if _, err := conn.WriteToUDP(pkt, addr); err != nil {
		log.Println(err)
	}
}

func getPeerByInternalIP(internalIP string) *Peer {
	peersMutex.Lock()
	defer peersMutex.Unlock()
	for _, p := range peers {
		if p.InternalIP == internalIP {
			return p
		}
	}
	return nil
}