/*
 *   Wi-Fi Detector opensource firmware
 *   ZD1211 through UART support
 *
 *   Copyright (C) 2007 Michael Buesch <mb@bu3sch.de>
 *
 *   This program is free software; you can redistribute it and/or
 *   modify it under the terms of the GNU General Public License
 *   as published by the Free Software Foundation; either version 2
 *   of the License, or (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 */

#include "zd1211.h"
#include "util.h"
#include "lcd.h"
#include "main.h"
#include "errno.h"

#include <string.h>
#include <avr/io.h>
#include <avr/interrupt.h>


/* Size of a scan data packet for one AP */
#define ZD1211_SCAN_PACK_SIZE		127

/* The packet data structure that we receive from the ZD1211 chip */
struct scan_packet {
	struct {
		uint8_t unknown;
		uint8_t channel;
		uint8_t rssi;
	} __attribute__((packed)) header;
	uint8_t __padding;
	uint8_t bssid[6];
	uint8_t seqctl[2];
	struct {
		uint8_t timestamp[8];
		uint8_t beaconint[2];
		uint8_t capab[2];
		uint8_t ie[0];
	} __attribute__((packed)) beacon;
} __attribute__((packed));

/* This is a hack to save RAM.
 * We don't need the LCD mem while receiving the data bits,
 * so we use it as temporary USART RX memory. */
struct usart_rx_data {
	/* The RX ring buffer */
#define USART_RX_RING_SIZE		90
	uint8_t rx_ring[USART_RX_RING_SIZE];
	int8_t rx_ring_first;
	int8_t rx_ring_last;
	uint8_t rx_ring_used;
	/* Error code from the RX IRQ handler */
	int8_t rx_irq_error;
	/* Buffer for the received struct scan_packet data */
#define USART_PACKET_BUFFER_SIZE	127
	uint8_t packet_buffer[USART_PACKET_BUFFER_SIZE];
};
#define get_rx_ring()		(((struct usart_rx_data *)lcd_buffer)->rx_ring)
#define get_rx_ring_first()	(&((struct usart_rx_data *)lcd_buffer)->rx_ring_first)
#define get_rx_ring_last()	(&((struct usart_rx_data *)lcd_buffer)->rx_ring_last)
#define get_rx_ring_used()	(&((struct usart_rx_data *)lcd_buffer)->rx_ring_used)
#define get_rx_irq_error()	(&((struct usart_rx_data *)lcd_buffer)->rx_irq_error)
#define get_packet_buffer()	(((struct usart_rx_data *)lcd_buffer)->packet_buffer)


/* Dynamic list for the scanned APs.
 * The data layout for this list is struct scanned_ap. */
#define AP_LIST_SIZE		300
uint8_t scanned_ap_list[AP_LIST_SIZE + 1];


static void usart_tx(uint8_t data)
{
	while (!(UCSR0A & (1 << UDRE0)))
		mb();
	UDR0 = data;
}

static int8_t usart_rx(uint8_t *data)
{
	uint8_t status;

	status = UCSR0A;
	if (!(status & (1 << RXC0)))
		return -ENODATA;
	if (unlikely(status & ((1 << FE0) | (1 << UPE0) | (1 << DOR0)))) {
		if (status & (1 << FE0))
			return -ERXFE;
		if (status & (1 << UPE0))
			return -ERXPE;
		if (status & (1 << DOR0))
			return -ERXOV;
	}
	*data = UDR0;

	return 0;
}

static void usart_init(void)
{
	uint8_t dummy;

	BUILD_BUG_ON(sizeof(struct usart_rx_data) > LCD_BUFFER_SIZE);
	lcd_clear_buffer();
	/* Set baud rate */
	UBRR0 = 0x000C;
	UCSR0A = (1 << U2X0);
	/* 8 Data bits, 1 Stop bit, No parity */
	UCSR0C = (1 << UCSZ00) | (1 << UCSZ01);
	/* Enable transceiver and RX IRQs */
	UCSR0B = (1 << RXEN0) | (1 << TXEN0) | (1 << RXCIE0);
	/* Drain the RX buffer */
	while (usart_rx(&dummy) != -ENODATA)
		mb();
}

static void usart_exit(void)
{
	/* Disable transceiver */
	UCSR0B = 0;
	lcd_clear_buffer();
}

ISR(USART_RX_vect)
{
	uint8_t data = data;
	uint8_t *ringbuf;
	int8_t last;
	uint8_t used;
	int8_t err;

	ringbuf = get_rx_ring();
	while (1) {
		err = usart_rx(&data);
		if (err == -ENODATA)
			break;
		if (unlikely(err)) {
			*get_rx_irq_error() = err;
			break;
		}
		used = *get_rx_ring_used();
		if (unlikely(used == USART_RX_RING_SIZE)) {
			/* ring buffer overflow */
			*get_rx_irq_error() = -EOVERFLOW;
			break;
		}
		*get_rx_ring_used() = used + 1;
		last = *get_rx_ring_last();
		ringbuf[last] = data;
		last++;
		if (last == USART_RX_RING_SIZE - 1)
			last = 0;
		*get_rx_ring_last() = last;
	}
}

/* Returns the length of the IE payload (zero if not found)
 * and puts a pointer to the IE payload into ie_ptr */
static uint8_t ieee80211_find_ie(const struct scan_packet *p, uint8_t ie_id,
				 const uint8_t **ie_ptr)
{
	uint8_t nrbytes;
	const uint8_t *d;
	uint8_t id, len;
	uint8_t found_ie_length = 0;

	nrbytes = ZD1211_SCAN_PACK_SIZE - offsetof(struct scan_packet, beacon.ie);
	d = p->beacon.ie;

	while (1) {
		if (nrbytes < 2)
			break;
		id = d[0];
		len = d[1];
		d += 2;
		nrbytes -= 2;
		if (id == ie_id) {
			*ie_ptr = d;
			found_ie_length = min(len, nrbytes);
			break;
		}
		/* Go to the next IE */
		if (len > nrbytes)
			break;
		d += len;
		nrbytes -= len;
	}

	return found_ie_length;
}

/* Read the next byte from the RX ringbuffer.
 * IRQs must be enabled.
 * timeout is in milliseconds. */
static int8_t rxring_read_next(uint8_t *data, uint8_t timeout)
{
	const uint8_t *ringbuf;
	int8_t first;
	uint8_t used, i;
	int8_t err;

	ringbuf = get_rx_ring();
	while (1) {
		cli();
		mb();
		err = *get_rx_irq_error();
		if (unlikely(err)) {
			/* There was some error in the lowlevel
			 * RX IRQ handler */
			sei();
			return err;
		}
		used = *get_rx_ring_used();
		if (!used) {
			sei();
			if (timeout == 0)
				return -ETIMEDOUT;
			for (i = 0; i < 200; i++) {
				cli();
				mb();
				used = *get_rx_ring_used();
				sei();
				if (used)
					break;
				if (key_is_pressed(KEY_SCAN)) {
					/* User interrupt */
					key_debounce(KEY_SCAN, 5);
					return -EINTR;
				}
				udelay(5);
			}
			if (timeout != 0xFF)
				timeout--;
			continue;
		}
		*get_rx_ring_used() = used - 1;
		first = *get_rx_ring_first();
		*data = ringbuf[first];
		first++;
		if (first == USART_RX_RING_SIZE - 1)
			first = 0;
		*get_rx_ring_first() = first;
		sei();
		break;
	}

	return 0;
}

void zd1211_init(void)
{
	ZD1211_POWER_DDR |= ZD1211_POWER_BIT;
	ZD1211_POWER_PORT &= ~ZD1211_POWER_BIT;
}

static void zd1211_start_scan(void)
{
	usart_init();
	ZD1211_POWER_PORT |= ZD1211_POWER_BIT;
	/* I'm not sure why we send 0xEE to the ZD here... */
	usart_tx(0xEE);
}

static void zd1211_disable(void)
{
	ZD1211_POWER_PORT &= ~ZD1211_POWER_BIT;
	mdelay(100);
	usart_exit();
}

static int8_t zd1211_wait_for_packet(uint8_t timeout)
{
	uint8_t data;
	int8_t err;

	while (1) {
		err = rxring_read_next(&data, timeout);
		if (err)
			return err;
		if (data == 0xAA)
			break;
	}

	return 0;
}

/* Check if we already have this AP in our list of scanned APs. */
static bool is_duplicate(const struct scan_packet *packet)
{
	const struct scanned_ap *ap;

	for_each_ap(ap) {
		if (compare_ether_addr(ap->bssid, packet->bssid) == 0)
			return 1;
	}

	return 0;
}

static int8_t zd1211_parse_packet(const struct scan_packet *packet,
				  uint8_t **aplist)
{
	uint8_t ie_len;
	const uint8_t *ie_ptr;
	struct scanned_ap *ap;

	if (!is_valid_ap_ptr(*aplist))
		return -EOVERFLOW;
	ap = (struct scanned_ap *)(*aplist);
	ap->size = 0;

	if (is_duplicate(packet))
		return -EIGNORE;

	ie_len = ieee80211_find_ie(packet, 0 /* SSID */, &ie_ptr);
	if (ie_len) {
		/* Found an SSID IE.
		 * First truncate the SSID to a line in the LCD.
		 * (We don't need looong cryptic SSIDs, as we have the BSSID). */
		ie_len = min(ie_len, LCD_NR_CHARCOLUMNS);
		/* Check if we have enough buffer space. */
		if (is_valid_ap_ptr(*aplist + ie_len)) {
			/* Copy it into the APlist SSID field */
			memcpy(ap->ssid, ie_ptr, ie_len);
			ap->size += ie_len;
		}
	}
	ap->size += sizeof(struct scanned_ap);
	/* Copy the BSSID */
	memcpy(ap->bssid, packet->bssid, ETH_ALEN);
	/* Get the channel number */
	ie_len = ieee80211_find_ie(packet, 3 /* DS channel */, &ie_ptr);
	if (ie_len == 1)
		ap->flags_chan = *ie_ptr & AP_CHAN;
	else
		ap->flags_chan = packet->header.channel & AP_CHAN;
	/* The the RSN IE and extract crypto information from it. */
	ie_len = ieee80211_find_ie(packet, 48 /* RSN IE */, &ie_ptr);
	if (ie_len >= 12) {
		const uint8_t *pairwise, *group;
		uint8_t suite_type;

		/* We only support RSN IE version 1 */
		if ((ie_ptr[0] != 0x01) || (ie_ptr[1] != 0x00))
			goto rsn_unsupported;
		/* Get a pointer to the pairwise and group cipher suite. */
		group = ie_ptr + 2;
		pairwise = ie_ptr + 8;
		/* Check if this is a standard OUI. */
		if ((pairwise[0] != 0x00) ||
		    (pairwise[1] != 0x0F) ||
		    (pairwise[2] != 0xAC))
			goto rsn_unsupported;
		suite_type = pairwise[3];
		if (suite_type == 0) {
			/* Use group suite. */
			if ((group[0] != 0x00) ||
			    (group[1] != 0x0F) ||
			    (group[2] != 0xAC))
				goto rsn_unsupported;
			suite_type = group[3];
		}
		if ((suite_type != 1 /* WEP-40 */) &&
		    (suite_type != 5 /* WEP-104 */))
			ap->flags_chan |= AP_WPA;
    rsn_unsupported:;
	}
	/* Update flags */
	if (!(packet->beacon.capab[0] & 0x10))
		ap->flags_chan |= AP_OPEN_NET;
	if (packet->beacon.capab[0] & 0x02)
		ap->flags_chan |= AP_IS_IBSS;
	ap->rssi = packet->header.rssi;

	*aplist += ap->size;

	return 0;
}

static int16_t zd1211_read_scan_result(void)
{
	int8_t err;
	uint8_t i;
	uint8_t data;
	uint8_t *packet_buffer;
	uint8_t *aplist = scanned_ap_list;
	uint8_t nr_aps_found = 0;

	memset(aplist, 0, AP_LIST_SIZE);
	packet_buffer = get_packet_buffer();
	err = zd1211_wait_for_packet(0xFF);
	if (err)
		goto out;
	while (1) {
		/* We are receiving a scan result packet.
		 * Get and parse it. */
		for (i = 0; i < ZD1211_SCAN_PACK_SIZE; i++) {
			err = rxring_read_next(&data, 100);
			if (err)
				goto out;
			packet_buffer[i] = data;
		}
		err = zd1211_parse_packet((const struct scan_packet *)packet_buffer,
					  &aplist);
		if (err == -EOVERFLOW) {
			/* Scan list full. Ignore the rest (if any). */
			err = 0;
			goto out;
		}
		if (err && (err != -EIGNORE))
			goto out;
		if (err != -EIGNORE)
			nr_aps_found++;
		err = zd1211_wait_for_packet(200);
		if (err) {
			if (err == -ETIMEDOUT)
				err = 0;
			break;
		}
	}
out:
	BUG_ON(err > 0);

	return err ? err : nr_aps_found;
}

static uint8_t do_sort_aps(uint8_t cur_index, bool only_open)
{
	struct scanned_ap *ap, *best_ap;

	while (1) {
		best_ap = NULL;
		for_each_ap(ap) {
			if (only_open &&
			    !(ap->flags_chan & AP_OPEN_NET))
				continue;
			if (ap->flags_chan & __AP_SORTED)
				continue;
			if (!best_ap || (best_ap->rssi < ap->rssi))
				best_ap = ap;
		}
		if (!best_ap)
			break;
		best_ap->index = cur_index++;
		best_ap->flags_chan |= __AP_SORTED;
	}

	return cur_index;
}

/* Sort the scanned APs. First come open and best. */
static void zd1211_sort_aps(void)
{
	uint8_t cur_index;

	/* The best open networks will come first. */
	cur_index = do_sort_aps(0, 1);
	/* Closed networks will come last. */
	do_sort_aps(cur_index, 0);
}

int16_t zd1211_scan(void)
{
	int16_t nr_aps_found;

	zd1211_start_scan();
	nr_aps_found = zd1211_read_scan_result();
	zd1211_disable();
	zd1211_sort_aps();

	return nr_aps_found;
}

bool zd1211_usb_is_plugged_in(void)
{
	bool connected;

	/* Does not work while scanning. Assert this here. */
	BUG_ON(UCSR0B & (1 << TXEN0));
	/* If USART TXD pin is high, USB is connected. */
	DDRD &= ~(1 << 1);
	PORTD &= ~(1 << 1);
	connected = !!(PIND & (1 << 1));

	return connected;
}

const struct scanned_ap * zd1211_fetch_ap(uint8_t index)
{
	const struct scanned_ap *ap;

	for_each_ap(ap) {
		if (ap->index == index)
			return ap;
	}

	return NULL;
}