// Package wifi provides WiFi packet capture and parsing
package wifi

import (
	"encoding/binary"
	"fmt"
	"net"
	"os/exec"
	"syscall"
	"time"
)

// 802.11 Frame Control field
const (
	FrameTypeManagement = 0x00
	FrameSubtypeProbe   = 0x04
)

// Radiotap constants
const (
	RadiotapPresent_TSFT      = 1 << 0
	RadiotapPresent_Flags     = 1 << 1
	RadiotapPresent_Rate      = 1 << 2
	RadiotapPresent_Channel   = 1 << 3
	RadiotapPresent_FHSS      = 1 << 4
	RadiotapPresent_DBM_Antsignal = 1 << 5
)

// CaptureSocket represents a raw packet capture socket
type CaptureSocket struct {
	fd    int
	iface string
}

// SetMonitorMode enables monitor mode on a WiFi interface
func SetMonitorMode(iface string) error {
	// Bring interface down
	if err := exec.Command("ip", "link", "set", iface, "down").Run(); err != nil {
		return fmt.Errorf("ip link down: %w", err)
	}

	// Set monitor mode
	if err := exec.Command("iw", "dev", iface, "set", "type", "monitor").Run(); err != nil {
		return fmt.Errorf("iw set monitor: %w", err)
	}

	// Set monitor flags
	exec.Command("iw", "dev", iface, "set", "monitor", "control", "otherbss").Run()

	// Bring interface up
	if err := exec.Command("ip", "link", "set", iface, "up").Run(); err != nil {
		return fmt.Errorf("ip link up: %w", err)
	}

	return nil
}

// SetManagedMode restores managed mode on a WiFi interface
func SetManagedMode(iface string) error {
	exec.Command("ip", "link", "set", iface, "down").Run()
	exec.Command("iw", "dev", iface, "set", "type", "managed").Run()
	exec.Command("ip", "link", "set", iface, "up").Run()
	return nil
}

// SetChannel sets the WiFi channel
func SetChannel(iface string, channel int) error {
	return exec.Command("iw", "dev", iface, "set", "channel", fmt.Sprintf("%d", channel)).Run()
}

// NewCaptureSocket creates a raw packet capture socket
func NewCaptureSocket(iface string) (*CaptureSocket, error) {
	// Get interface index
	ifi, err := net.InterfaceByName(iface)
	if err != nil {
		return nil, fmt.Errorf("interface %s: %w", iface, err)
	}

	// Create raw socket (ETH_P_ALL = 0x0003)
	fd, err := syscall.Socket(syscall.AF_PACKET, syscall.SOCK_RAW, int(htons(0x0003)))
	if err != nil {
		return nil, fmt.Errorf("socket: %w", err)
	}

	// Bind to interface
	addr := syscall.SockaddrLinklayer{
		Protocol: htons(0x0003),
		Ifindex:  ifi.Index,
	}
	if err := syscall.Bind(fd, &addr); err != nil {
		syscall.Close(fd)
		return nil, fmt.Errorf("bind: %w", err)
	}

	// Note: promiscuous mode not needed when interface is in monitor mode
	// Monitor mode already captures all packets

	return &CaptureSocket{fd: fd, iface: iface}, nil
}

// Read reads a packet from the socket
func (c *CaptureSocket) Read(buf []byte) (int, error) {
	return syscall.Read(c.fd, buf)
}

// Close closes the capture socket
func (c *CaptureSocket) Close() error {
	return syscall.Close(c.fd)
}

// Fd returns the file descriptor
func (c *CaptureSocket) Fd() int {
	return c.fd
}

// htons converts a short (uint16) from host to network byte order
func htons(i uint16) uint16 {
	return (i<<8)&0xff00 | i>>8
}

// ProbeRequest represents a parsed WiFi probe request
type ProbeRequest struct {
	SourceMAC string
	RSSI      int8
	SSID      string
	Timestamp time.Time
}

// ParseRadiotapRSSI extracts RSSI from radiotap header
func ParseRadiotapRSSI(data []byte) (headerLen int, rssi int8, ok bool) {
	if len(data) < 8 {
		return 0, 0, false
	}

	// Radiotap header: version(1) + pad(1) + length(2) + present_flags(4+)
	// version := data[0]
	headerLen = int(binary.LittleEndian.Uint16(data[2:4]))
	if headerLen > len(data) {
		return 0, 0, false
	}

	presentFlags := binary.LittleEndian.Uint32(data[4:8])
	offset := 8

	// Skip extended present flags if any
	for presentFlags&(1<<31) != 0 && offset+4 <= headerLen {
		presentFlags = binary.LittleEndian.Uint32(data[offset : offset+4])
		offset += 4
	}

	// Re-read first present flags
	presentFlags = binary.LittleEndian.Uint32(data[4:8])

	// Walk through present fields to find dBm antenna signal
	if presentFlags&RadiotapPresent_TSFT != 0 {
		offset = (offset + 7) &^ 7 // Align to 8 bytes
		offset += 8
	}
	if presentFlags&RadiotapPresent_Flags != 0 {
		offset += 1
	}
	if presentFlags&RadiotapPresent_Rate != 0 {
		offset += 1
	}
	if presentFlags&RadiotapPresent_Channel != 0 {
		offset = (offset + 1) &^ 1 // Align to 2 bytes
		offset += 4
	}
	if presentFlags&RadiotapPresent_FHSS != 0 {
		offset += 2
	}
	if presentFlags&RadiotapPresent_DBM_Antsignal != 0 {
		if offset < headerLen {
			rssi = int8(data[offset])
			return headerLen, rssi, true
		}
	}

	return headerLen, -100, true // Default RSSI if not found
}

// ParseProbeRequest parses a WiFi frame for probe requests
func ParseProbeRequest(data []byte) (*ProbeRequest, bool) {
	// Parse radiotap header
	rtLen, rssi, ok := ParseRadiotapRSSI(data)
	if !ok || rtLen >= len(data) {
		return nil, false
	}

	frame := data[rtLen:]
	if len(frame) < 24 {
		return nil, false
	}

	// 802.11 header
	// Frame Control: type/subtype(2) + duration(2) + addr1(6) + addr2(6) + addr3(6) + seq(2)
	frameControl := binary.LittleEndian.Uint16(frame[0:2])

	// Check if it's a Probe Request (type=0, subtype=4)
	frameType := (frameControl >> 2) & 0x03
	frameSubtype := (frameControl >> 4) & 0x0F

	if frameType != FrameTypeManagement || frameSubtype != FrameSubtypeProbe {
		return nil, false
	}

	// Extract Source Address (addr2 at offset 10)
	srcMAC := fmt.Sprintf("%02x:%02x:%02x:%02x:%02x:%02x",
		frame[10], frame[11], frame[12], frame[13], frame[14], frame[15])

	// Parse management frame body (after 24-byte header)
	body := frame[24:]

	// Parse information elements to find SSID
	ssid := ""
	offset := 0
	for offset+2 <= len(body) {
		elementID := body[offset]
		elementLen := int(body[offset+1])
		if offset+2+elementLen > len(body) {
			break
		}
		if elementID == 0 { // SSID
			ssidBytes := body[offset+2 : offset+2+elementLen]
			ssid = string(ssidBytes) // Best-effort UTF-8
		}
		offset += 2 + elementLen
	}

	return &ProbeRequest{
		SourceMAC: srcMAC,
		RSSI:      rssi,
		SSID:      ssid,
		Timestamp: time.Now(),
	}, true
}

// ChannelHopper cycles through WiFi channels
type ChannelHopper struct {
	iface    string
	channels []int
	dwell    time.Duration
	stop     chan struct{}
}

// NewChannelHopper creates a new channel hopper
func NewChannelHopper(iface string, channels []int, dwell time.Duration) *ChannelHopper {
	return &ChannelHopper{
		iface:    iface,
		channels: channels,
		dwell:    dwell,
		stop:     make(chan struct{}),
	}
}

// Start starts channel hopping
func (h *ChannelHopper) Start() {
	go func() {
		i := 0
		for {
			select {
			case <-h.stop:
				return
			default:
				ch := h.channels[i%len(h.channels)]
				SetChannel(h.iface, ch)
				i++
				time.Sleep(h.dwell)
			}
		}
	}()
}

// Stop stops channel hopping
func (h *ChannelHopper) Stop() {
	close(h.stop)
}