__author__ = 'Tom James Holub'


import sqlite3
import json

class Db(object):

	def __init__(self, path, foundation_block_id, foundation_block_height):
		self.connection = sqlite3.connect(path)
		self.setup(foundation_block_id, foundation_block_height)

	def setup(self, foundation_block_id, foundation_block_height):
		c = self.connection.cursor()
		c.execute("CREATE TABLE IF NOT EXISTS `transaction` (\
			`id` VARCHAR(64) NOT NULL,\
			`sent` TINYINT NOT NULL DEFAULT '0', \
			`data` TEXT NOT NULL \
		);")
		c.execute("CREATE UNIQUE INDEX IF NOT EXISTS `id_idx` ON `transaction` (`id`);")
		c.execute("CREATE INDEX IF NOT EXISTS `sent_idx` ON `transaction` (`sent`);")
		c.execute("CREATE TABLE IF NOT EXISTS `address` (`address` VARCHAR(64) NOT NULL, `account_id` VARCHAR(64) NULL);")
		c.execute("CREATE UNIQUE INDEX IF NOT EXISTS `address_idx` ON `address` (`address`);")
		c.execute("CREATE INDEX IF NOT EXISTS `account_id_idx` ON `address` (`account_id`);")
		c.execute("CREATE TABLE IF NOT EXISTS `block` (\
			`id` VARCHAR(64) NOT NULL, \
			`parent_id` VARCHAR(64) NOT NULL, \
			`height` INT(32) NULL, \
			`have_parent` INT(1) NULL, \
			`timestamp` INT(32) NULL \
		);")
		c.execute("CREATE UNIQUE INDEX IF NOT EXISTS `id_idx` ON `block` (`id`);")
		c.execute("CREATE UNIQUE INDEX IF NOT EXISTS `parent_id_idx` ON `block` (`id`);")
		c.execute("CREATE UNIQUE INDEX IF NOT EXISTS `height_idx` ON `block` (`height`);")
		try:
			c.execute("INSERT INTO `block` (`id`, `parent_id`, `have_parent`, `height`) VALUES (?, 'foundation', 1, ?);", (foundation_block_id, foundation_block_height))
		except sqlite3.IntegrityError:
			pass
		self.connection.commit()
		c.close()

	def fetch_all(self, table):
		c = self.connection.cursor()
		c.execute("SELECT * FROM `%s`;" % table)
		rows = c.fetchall()
		c.close()
		return rows

	def add_address(self, address, account_id=None):
		c = self.connection.cursor()
		try:
			c.execute("INSERT INTO `address` (`address`, `account_id`) VALUES (?, ?);", (address, account_id))
			self.connection.commit()
		except sqlite3.IntegrityError:
			pass
		c.close()

	def get_addresses(self):
		c = self.connection.cursor()
		c.execute("SELECT `address`, `account_id` FROM `address`;")
		rows = c.fetchall()
		c.close()
		addresses = {}
		for row in rows:
			addresses[row[0]] = row[1]  # {address: account_id}
		return addresses

	def add_tx(self, tx, watch_addresses):
		data = {'in': [], 'out': []}
		input_addresses = []
		for input in tx.txs_in:
			subscribed = input.bitcoin_address() in watch_addresses
			input_addresses.append(input.bitcoin_address())
			data['in'].append({
				'address': input.bitcoin_address(),
				'subscribed': subscribed,
				'account_id': watch_addresses[input.bitcoin_address()] if subscribed else None,
			})
		for out in tx.txs_out:
			subscribed = out.bitcoin_address() in watch_addresses
			data['out'].append({
				'address': out.bitcoin_address(),
				'subscribed': subscribed,
				'amount': out.coin_value,
				'account_id': watch_addresses[out.bitcoin_address()] if subscribed else None,
				'change': out.bitcoin_address() in input_addresses
			})
		c = self.connection.cursor()
		try:
			c.execute("INSERT INTO `transaction` (`id`, `data`) VALUES (?, ?);", (tx.id(), json.dumps(data)))
			self.connection.commit()
		except sqlite3.IntegrityError:
			pass
		c.close()

	def add_block(self, id, parent_id, timestamp):
		c = self.connection.cursor()

		c.execute("SELECT `id` FROM `block` where `id`=?;", (id,))
		id_taken = bool(c.fetchone())
		c.execute("SELECT `id` FROM `block` where `parent_id`=?;", (parent_id,))
		parent_id_taken = bool(c.fetchone())
		c.execute("SELECT `height` FROM `block` where `id`=?;", (parent_id,))
		parent = c.fetchone()
		have_parent = bool(parent)
		height = None if not have_parent or not parent[0] else parent[0] + 1
		c.execute("SELECT `id`, `height` FROM `block` where `parent_id`=?;", (id,))
		child = c.fetchone()
		child_id = child[0] if child else None
		if height is None and child is not None:
			height = None if not child[1] else child[1] - 1
			if height is not None and have_parent and parent[0] +1 != height - 1:
				raise AssertionError('Parent height +1 != child height -1 for block %s' % id)

		if not parent_id_taken and not id_taken:
			c.execute("INSERT INTO `block` (`id`, `parent_id`, `have_parent`, `height`, `timestamp`) VALUES (?, ?, ?, ?, ?);", (id, parent_id, have_parent, height, timestamp))
		elif not parent_id_taken and id_taken:  # we had incorrect parent_id
			c.execute("UPDATE `block` SET `parent_id`=?, `have_parent`=?, `height`=?, `timestamp`=? WHERE `id`=?;", (parent_id, have_parent, height, timestamp, id))
		elif parent_id_taken and not id_taken:  # we had incorrenct id
			c.execute("UPDATE `block` SET `id`=?, `have_parent`=?, `height`=? WHERE `parent_id`=?;", (id, have_parent, height, parent_id))
		else:  # exact same block exists, update height/have_parent
			c.execute("UPDATE `block` SET `have_parent`=?, `height`=?, `timestamp`=? WHERE `id`=?;", (have_parent, height, timestamp, id))
		if child_id is not None:
			c.execute("UPDATE `block` SET `have_parent`=1 WHERE `id`=?;", (child_id,))
		self.connection.commit()
		c.close()

	def get_missing_blocks(self):
		c = self.connection.cursor()
		c.execute("SELECT `parent_id` FROM `block` WHERE `have_parent`=0;")
		rows = c.fetchall()
		c.close()
		return [str(row[0]) for row in rows]

	def mark_txs_sent(self, ids):
		if ids:
			c = self.connection.cursor()
			c.execute("UPDATE `transaction` SET `sent`=1 WHERE `id` IN ('%s');" % "', '".join(ids))
			self.connection.commit()
			c.close()

	def get_unsent_txs(self, limit=50):
		c = self.connection.cursor()
		c.execute("SELECT `id`, `data` FROM `transaction` WHERE `sent`=0 LIMIT %d;" % limit)
		rows = c.fetchall()
		c.close()
		result = {}
		for row in rows:
			result[row[0]] = json.loads(row[1])
		return result
