/*
 *   Wi-Fi Detector opensource firmware
 *
 *   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 "main.h"
#include "util.h"
#include "lcd.h"
#include "zd1211.h"
#include "calibration.h"
#include "atomic.h"
#include "errno.h"

#include <avr/io.h>
#include <avr/interrupt.h>
#include <avr/sleep.h>
#include <stdint.h>

#include "battery_tab.h"


/* The battery-measure-enable pin */
#define BATTMEAS_EN_PORT	PORTD
#define BATTMEAS_EN_DDR		DDRD
#define BATTMEAS_EN_BIT		(1 << 3)

/* Battery charge status values. */
enum {
	CHARGESTAT_IDLE,	/* We are not charging. */
	CHARGESTAT_PROGRESS,	/* Charging is in progress. */
	CHARGESTAT_FINISHED,	/* Battery is fully charged. */
};

/* System jiffies counter. Overflows every 25.6 seconds at 10Hz. */
static uint8_t jiffies;
/* Number of jiffies the backlight is already ON.
 * (Also see FLAG_BACKL_ON) */
static uint8_t backl_on_time;
/* The number of APs that we already scanned. */
static uint8_t nr_scanned_aps;
/* The AP that's currently displayed in the LCD */
static uint8_t current_ap;
/* The battery voltage in percent */
static uint8_t battery_percent;
/* Some flags */
static uint8_t flags;
#define FLAG_SCANNING		0 /* We are scanning. */
#define FLAG_AP_LOCKED		1 /* Current AP is locked. */
#define FLAG_DO_RESCAN		2 /* Do a new scan. Used with FLAG_AP_LOCKED. */
#define FLAG_UPDATE_LCD		3 /* Update the LCD in the main event loop. */
#define FLAG_BACKL_ON		4 /* Backlight switched on. */
#define FLAG_USB_PLUGGED	5 /* USB is plugged in. */
#define FLAG_IN_INTERRUPT	6 /* This flag is set when executing in an IRQ handler. */
/* The BSSID we are locked to (if any) */
static uint8_t locked_bssid[ETH_ALEN];
/* The number of seconds-times-10 we've been idle. */
static uint8_t idle_10secs_count;


/* Entering an IRQ handler. */
static inline void irq_enter(void)
{
	__atomic_bit_set(&flags, FLAG_IN_INTERRUPT);
}

/* Exiting an IRQ handler. */
static inline void irq_exit(void)
{
	__atomic_bit_clear(&flags, FLAG_IN_INTERRUPT);
}

static inline uint8_t get_jiffies(void)
{
	mb();
	return jiffies;
}

/* Measure ADC6 with an 1.1V reference and a prescaler of 2 */
static uint16_t adc6_measure(void)
{
	while (ADCSRA & (1 << ADSC))
		;
	ADMUX = (1<<REFS0)|(1<<REFS1)|(1<<MUX1)|(1<<MUX2);
	ADCSRA = (1<<ADEN)|(1<<ADSC)|(1<<ADPS2);
	while (ADCSRA & (1 << ADSC))
		;

	return ADC;
}

static uint16_t adc6_average(void)
{
	uint16_t average;

	average = adc6_measure();
	average += adc6_measure();
	average += adc6_measure();
	average += adc6_measure();
	average /= 4;

	return average;
}

/* Measure ADC7 with an 1.1V reference and a prescaler of 2 */
static uint16_t adc7_measure(void)
{
	while (ADCSRA & (1 << ADSC))
		;
	ADMUX = (1<<REFS0)|(1<<REFS1)|(1<<MUX0)|(1<<MUX1)|(1<<MUX2);
	ADCSRA = (1<<ADEN)|(1<<ADSC)|(1<<ADPS2);
	while (ADCSRA & (1 << ADSC))
		;

	return ADC;
}

static uint16_t adc7_average(void)
{
	uint16_t average;

	average = adc7_measure();
	average += adc7_measure();
	average += adc7_measure();
	average += adc7_measure();
	average /= 4;

	return average;
}

static void battery_measure_on(void)
{
	BATTMEAS_EN_PORT |= BATTMEAS_EN_BIT;
	udelay(10);
}

static inline void battery_measure_off(void)
{
	BATTMEAS_EN_PORT &= ~BATTMEAS_EN_BIT;
}

static void battery_init(void)
{
	BATTMEAS_EN_DDR |= BATTMEAS_EN_BIT;
	battery_measure_off();
}

/* Measures the battery voltage and returns a percentage value */
static uint8_t measure_battery(void)
{
	uint16_t adc7;
	uint16_t percent;

	BUG_ON(atomic_bit_test(&flags, FLAG_USB_PLUGGED));

	battery_measure_on();
	adc7 = adc7_average();
	battery_measure_off();

	if (adc7 > ADC7_MAX)
		adc7 = ADC7_MAX;
	if (adc7 > ADC7_MIN)
		adc7 -= ADC7_MIN;
	else
		adc7 = 0;

	percent = pgm_read_byte(&battery_adc7_tab[adc7]);
	percent *= 100;
	percent /= pgm_read_byte(&battery_adc7_tab[ADC7_MAX - ADC7_MIN]);

	return percent;
}

static uint8_t battery_charge_status(void)
{
	uint16_t adc6;
	uint8_t stat;

	BUG_ON(!atomic_bit_test(&flags, FLAG_USB_PLUGGED));

	adc6 = adc6_average();
	if (adc6 <= 0xC0)
		stat = CHARGESTAT_IDLE;
	else if (adc6 <= 0x160)
		stat = CHARGESTAT_PROGRESS;
	else
		stat = CHARGESTAT_FINISHED;

lcd_charcursor(3, 0);
printhex(adc6 >> 8);
printhex(adc6);
lcd_charcursor(0, 0);

	return stat;
}

/* Trigger the LCD backlight ON. */
static void backlight_trigger(void)
{
	uint8_t sreg;

	sreg = irq_disable_save();
	backl_on_time = 0;
	if (!__atomic_bit_test(&flags, FLAG_BACKL_ON)) {
		__atomic_bit_set(&flags, FLAG_BACKL_ON);
		lcd_backlight_on();
	}
	irq_restore(sreg);
}

static void update_lcd_scanning(void)
{
	lcd_charcursor(0, 0);
	print(PSTR("scanning..."));
}

static void update_lcd_usb_plugged(void)
{
	uint8_t stat;

	lcd_charcursor(0, 0);
	stat = battery_charge_status();
	switch (stat) {
	case CHARGESTAT_IDLE:
		print(PSTR("USB WLAN stick"));
		lcd_charcursor(1, 0);
		print(PSTR("operation"));
		break;
	case CHARGESTAT_PROGRESS:
		print(PSTR("Charging battery"));
		break;
	case CHARGESTAT_FINISHED:
		print(PSTR("Battery fully"));
		lcd_charcursor(1, 0);
		print(PSTR("charged"));
		break;
	}
}

static void update_lcd_printresult(void)
{
	const struct scanned_ap *ap;
	uint8_t i;
	uint8_t nr_open_aps = 0;

	lcd_charcursor(0, 10);
	print(PSTR("B-"));
	printdec(battery_percent);
	lcd_char_out('%');

	if (!nr_scanned_aps && !atomic_bit_test(&flags, FLAG_AP_LOCKED)) {
		lcd_charcursor(0, 0);
		print(PSTR("0/0(0)"));
		lcd_charcursor(2, 0);
		print(PSTR("Press SCAN"));
		return;
	}

	/* Count the number of open APs. */
	for_each_ap(ap) {
		if (ap->flags_chan & AP_OPEN_NET)
			nr_open_aps++;
	}
	/* Get the currently active AP. */
	ap = zd1211_fetch_ap(current_ap);
	BUG_ON(!atomic_bit_test(&flags, FLAG_AP_LOCKED) && !ap);

	lcd_charcursor(0, 0);
	if (atomic_bit_test(&flags, FLAG_AP_LOCKED)) {
		print(PSTR("LOCKED"));
	} else {
		/* Print the counts. */
		printdec(current_ap + 1);
		lcd_char_out('/');
		printdec(nr_scanned_aps);
		lcd_char_out('(');
		printdec(nr_open_aps);
		lcd_char_out(')');
	}

	if (ap) {
		lcd_charcursor(1, 0);
		print(PSTR("CH-"));
		printdec(ap->flags_chan & AP_CHAN);

		lcd_charcursor(1, 6);
		if (ap->flags_chan & AP_OPEN_NET)
			lcd_char_out('O'); /* open unencrypted net */
		else if (ap->flags_chan & AP_WPA)
			lcd_char_out('C'); /* closed WPA net */
		else
			lcd_char_out('w'); /* "closed" WEP net */

		lcd_charcursor(1, 8);
		print(PSTR("Sig-"));
		printdec(ap->rssi);
		lcd_char_out('%');

		lcd_charcursor(2, 0);
		if (ap->flags_chan & AP_IS_IBSS)
			print(PSTR("AdH "));
		else
			print(PSTR("AP "));
		for (i = 0; i < ETH_ALEN; i++)
			printhex(ap->bssid[i]);

		lcd_charcursor(3, 0);
		if (ap_ssid_length(ap))
			printmem((const char *)(ap->ssid), ap_ssid_length(ap));
		else
			print(PSTR("<unknown>"));
	} else {
		/* We are locked to a BSSID, but we didn't find it anymore. */
		lcd_charcursor(2, 0);
		print(PSTR("<not found>"));
		lcd_charcursor(3, 0);
		print(PSTR("scanning..."));
	}
}

static void update_lcd(void)
{
	/* This must not be called from IRQ context. */
	BUG_ON(atomic_bit_test(&flags, FLAG_IN_INTERRUPT));

	lcd_clear_buffer();
	if (atomic_bit_test(&flags, FLAG_USB_PLUGGED)) {
		update_lcd_usb_plugged();
	} else {
		if (atomic_bit_test(&flags, FLAG_SCANNING) &&
		    !atomic_bit_test(&flags, FLAG_AP_LOCKED))
			update_lcd_scanning();
		else
			update_lcd_printresult();
	}
	lcd_commit();
}

static void watchdog_turn_off(void)
{
	__asm__ __volatile__("wdr" : : );
	MCUSR &= ~(1 << WDRF);
	WDTCSR |= (1 << WDCE) | (1 << WDE);
	WDTCSR = 0;
}

void emergency_shutdown(void)
{
	cli();
	watchdog_turn_off();
	battery_measure_off();
	ZD1211_POWER_PORT &= ~ZD1211_POWER_BIT; /* Disable scanning */
	lcd_backlight_off();
	TCCR1B = 0; /* Disable system timer */
	TIMSK1 = 0; /* Disable IRQs */
	UCSR0B = 0; /* Disable USART transceiver and IRQs */
}

/* Call this when the battery is critically low to stop the system */
static void __attribute__((noreturn)) battery_emergency(void)
{
	emergency_shutdown();
	lcd_clear_buffer();
	print(PSTR("*** WARNING ***"));
	lcd_charcursor(1, 0);
	print(PSTR("RECHARGE BATTERY"));
	lcd_charcursor(2, 2);
	print(PSTR("IMMEDIATELY"));
	lcd_charcursor(3, 4);
	print(PSTR("PLEASE"));
	lcd_commit();
	infinite_sleep();
}

static void adc_init(void)
{
	/* Throw the first result away */
	adc6_measure();
	adc7_measure();

	battery_percent = measure_battery();
}

/* Unlock scanning of the current AP.
 * Can be called from all contexts. */
static void unlock_current_ap(void)
{
	uint8_t sreg;
	bool in_interrupt;

	sreg = irq_disable_save();
	in_interrupt = __atomic_bit_test(&flags, FLAG_IN_INTERRUPT);

	__atomic_bit_clear(&flags, FLAG_AP_LOCKED);
	__atomic_bit_clear(&flags, FLAG_DO_RESCAN);
	if (in_interrupt)
		__atomic_bit_set(&flags, FLAG_UPDATE_LCD);

	irq_restore(sreg);

	if (!in_interrupt)
		update_lcd();
}

/* Lock scanning to the current AP.
 * Must not be called from IRQ context. */
static void lock_current_ap(void)
{
	const struct scanned_ap *ap;

	if (atomic_bit_test(&flags, FLAG_AP_LOCKED))
		return;
	if (!nr_scanned_aps)
		return;
	ap = zd1211_fetch_ap(current_ap);
	if (!ap)
		return;
	memcpy(locked_bssid, ap->bssid, ETH_ALEN);
	mb();
	atomic_bit_set(&flags, FLAG_AP_LOCKED);
	atomic_bit_set(&flags, FLAG_DO_RESCAN);
	update_lcd();
}

/* Called with IRQs disabled. */
static void usb_connection_changed(bool is_plugged_in)
{
	if (is_plugged_in) {
		BATTMEAS_EN_PORT &= ~BATTMEAS_EN_BIT;
		__atomic_bit_clear(&flags, FLAG_BACKL_ON);
		lcd_backlight_off();
	} else {
		__atomic_bit_set(&flags, FLAG_BACKL_ON);
		lcd_backlight_on();
	}
	__atomic_bit_set(&flags, FLAG_UPDATE_LCD);
}

/* Timer that triggers every second */
static void handle_timer_1sec(void)
{
	if (__atomic_bit_test(&flags, FLAG_USB_PLUGGED))
		return;

	/* If we are locked to an AP, rescan every second. */
	if (__atomic_bit_test(&flags, FLAG_AP_LOCKED)) {
		if (!__atomic_bit_test(&flags, FLAG_SCANNING))
			__atomic_bit_set(&flags, FLAG_DO_RESCAN);
	}
}

/* Timer that triggers every 10 seconds */
static void handle_timer_10sec(void)
{
	uint8_t old_batt;

	if (idle_10secs_count != 0xFF)
		idle_10secs_count++;
	if ((idle_10secs_count >= AP_LOCK_AUTORELEASE_TIMEOUT) &&
	    __atomic_bit_test(&flags, FLAG_AP_LOCKED)) {
		/* The user left us a long time idle and locked to an AP.
		 * Did he forget us? Auto-unlock to avoid draining the
		 * power too much, as scanning uses lots of power. */
		unlock_current_ap();
	}

	if (__atomic_bit_test(&flags, FLAG_USB_PLUGGED))
		return;

	/* Measure the battery voltage */
	old_batt = battery_percent;
	battery_percent = measure_battery();
	if (unlikely(battery_percent <= MIN_BATTERY_PERCENT))
		battery_emergency();
	if (battery_percent != old_batt)
		__atomic_bit_set(&flags, FLAG_UPDATE_LCD);
}

/* System timer. Triggers every centisecond. */
ISR(TIMER1_COMPA_vect)
{
	static uint8_t timer_1sec;
	static uint8_t timer_10sec;

	irq_enter();

	jiffies++;

	if (!__atomic_bit_test(&flags, FLAG_SCANNING)) {
		bool plugged;

		/* Check if we are plugged into USB. */
		plugged = zd1211_usb_is_plugged_in();
		if (unlikely(plugged != __atomic_bit_test(&flags, FLAG_USB_PLUGGED))) {
			__atomic_bit_flip(&flags, FLAG_USB_PLUGGED);
			usb_connection_changed(plugged);
		}
	}

	/* 1 second timer. */
	if (++timer_1sec == (JIFFIES_PER_SECOND * 1)) {
		timer_1sec = 0;
		handle_timer_1sec();
	}
	/* 10 second timer. */
	if (++timer_10sec == (JIFFIES_PER_SECOND * 10)) {
		timer_10sec = 0;
		handle_timer_10sec();
	}

	if (__atomic_bit_test(&flags, FLAG_USB_PLUGGED))
		goto out;

	/* Switch the backlight OFF, if it's ON for too long,
	 * but not if we are locked to an AP. */
	if (__atomic_bit_test(&flags, FLAG_AP_LOCKED)) {
		backl_on_time = 0;
	} else {
		if (__atomic_bit_test(&flags, FLAG_BACKL_ON)) {
			if (++backl_on_time > MAX_BACKL_ON_TIME) {
				__atomic_bit_clear(&flags, FLAG_BACKL_ON);
				lcd_backlight_off();
			}
		}
	}

out:
	irq_exit();
}

static void timer_init(void)
{
	/* Initialize the system timer */
	TCCR1B = (1 << WGM12) | SYSTIMER_TIMERFREQ; /* Speed */
	OCR1A = SYSTIMER_CMPVAL; /* CompareMatch value */
	TIMSK1 |= (1 << OCIE1A); /* IRQ mask */
}

static int8_t trigger_scan(void)
{
	const struct scanned_ap *ap;
	int16_t result;
	int8_t err = 0;

	atomic_bit_set(&flags, FLAG_SCANNING);
	mb();
	atomic_bit_clear(&flags, FLAG_DO_RESCAN);
	update_lcd();
	backlight_trigger();

	result = zd1211_scan();
	if (result < 0) {
		err = result;
		nr_scanned_aps = 0;
	}
	backlight_trigger();
	if (err == -EINTR)
		goto out; /* User interrupt */
	if (err) {
		lcd_clear_buffer();
		print(PSTR("Scanning failed"));
		lcd_charcursor(1, 0);
		print(PSTR("Error 0x"));
		printhex(-err);
		lcd_charcursor(2, 0);
		print(errno_string(err));
		lcd_commit();
		mdelay(10000);
		goto out;
	}
	nr_scanned_aps = result;

	if (atomic_bit_test(&flags, FLAG_AP_LOCKED)) {
		/* Search the AP we are locked to. */
		current_ap = 0xFF;
		for_each_ap(ap) {
			if (compare_ether_addr(locked_bssid, ap->bssid) == 0) {
				current_ap = ap->index;
				break;
			}
		}
	} else
		current_ap = 0;

out:
	mb();
	atomic_bit_clear(&flags, FLAG_SCANNING);
	update_lcd();

	return err;
}

static void keys_init(void)
{
	/* Configure port as input */
	KEYS_DDR &= ~KEYS_MASK;
	/* with pullups */
	KEYS_PORT |= KEYS_MASK;
}

void key_debounce(uint8_t key_mask, uint8_t delay)
{
	uint8_t i;

	idle_10secs_count = 0; /* Non-atomic access is OK. */

	mdelay(5);
	do {
		for (i = 35; i; i--) {
			if (!key_is_pressed(key_mask)) {
				/* Released */
				break;
			}
			mdelay(5);
		}
	} while (delay--);
}

/* The NEXT key was pressed. */
static void key_next_pressed(void)
{
	backlight_trigger();
	key_debounce(KEY_NEXT, 10);
	if (!nr_scanned_aps)
		return;
	if (atomic_bit_test(&flags, FLAG_AP_LOCKED))
		return;
	current_ap++;
	if (current_ap >= nr_scanned_aps)
		current_ap = 0;
	update_lcd();
}

/* The PREV key was pressed. */
static void key_prev_pressed(void)
{
	backlight_trigger();
	key_debounce(KEY_PREV, 10);
	if (!nr_scanned_aps)
		return;
	if (atomic_bit_test(&flags, FLAG_AP_LOCKED))
		return;
	if (current_ap == 0)
		current_ap = nr_scanned_aps - 1;
	else
		current_ap--;
	update_lcd();
}

/* The SCAN key was pressed. */
static void key_scan_pressed(void)
{
	backlight_trigger();
	key_debounce(KEY_SCAN, 6);
	if (key_is_pressed(KEY_SCAN)) {
		/* The key is still pressed?
		 * Ok, lock the current AP. */
		lock_current_ap();
		key_debounce(KEY_SCAN, 20);
		return;
	}
	if (atomic_bit_test(&flags, FLAG_AP_LOCKED)) {
		/* We were locked to an AP. Unlock. */
		unlock_current_ap();
		return;
	}
	BUG_ON(atomic_bit_test(&flags, FLAG_DO_RESCAN));

	/* Scan for APs */
	trigger_scan();
}

#if 0
static void print_hwstate(bool batt_enable)
{
	if (batt_enable)
		battery_measure_on();
	else
		battery_measure_off();
	lcd_clear_buffer();
	lcd_charcursor(0, 0);
	printdec(adc6_average());
	lcd_charcursor(0, 6);
	printdec(adc7_average());
	battery_measure_off();
	lcd_charcursor(1, 0);
	printhex(PINB);
	lcd_char_out(' ');
	printhex(PINC);
	lcd_char_out(' ');
	printhex(PIND);
	lcd_char_out(' ');
	lcd_charcursor(2, 0);
	if (zd1211_usb_is_plugged_in())
		print(PSTR("USB plugged"));
	else
		print(PSTR("USB unplugged"));
	lcd_charcursor(3, 0);
	if (batt_enable)
		print(PSTR("batt_enable = 1"));
	else
		print(PSTR("batt_enable = 0"));
	lcd_commit();
}

static void print_hwstate_loop(void)
{
	cli();
	while (1) {
		print_hwstate(1);
		mdelay(1000);
	}
}
#endif

int main(void)
{
	uint8_t last_update_jiffies;

	cli();
	flags = 0;
	battery_init();
	keys_init();
	lcd_init();
	adc_init();
	zd1211_init();
	timer_init();
	update_lcd();
	backlight_trigger();
	set_sleep_mode(SLEEP_MODE_IDLE);
	sei();

	last_update_jiffies = get_jiffies();
	while (1) {
		int8_t err;

		/* Save power. The system timer will periodically wake us. */
		sleep_mode();

		if (atomic_bit_test_and_clear(&flags, FLAG_UPDATE_LCD)) {
			update_lcd();
			last_update_jiffies = get_jiffies();
		}

		if (atomic_bit_test(&flags, FLAG_USB_PLUGGED)) {
			/* Periodically update the display to check for changes
			 * in the battery charge status. */
			if (time_before(last_update_jiffies + (JIFFIES_PER_SECOND * 1),
					get_jiffies())) {
				update_lcd();
				last_update_jiffies = get_jiffies();
			}
			continue;
		}
		/* USB not plugged. Running on battery. Normal operation... */

		if (key_is_pressed(KEY_NEXT))
			key_next_pressed();
		if (key_is_pressed(KEY_PREV))
			key_prev_pressed();
		if (key_is_pressed(KEY_SCAN))
			key_scan_pressed();

		if (atomic_bit_test_and_clear(&flags, FLAG_DO_RESCAN)) {
			err = trigger_scan();
			if (err == -EINTR) {
				/* User interrupt. Unlock it. */
				unlock_current_ap();
			}
		}
	}
}