# -*- coding: utf-8 -*-
"""
Postgres-based processor

Author: Bartosz Broda, bartosz.broda@gmail.com
"""

import codecs, hashlib
import psycopg2
import psycopg2.extensions

from architect2.processors.memprocessor import MemProcessor

class PostgresProcessor(MemProcessor):
	"""
	Processes a given file and writes data to the (postgres) database.
	"""

	def __init__(self,matrix_id, user='boban', database = 'clarin'):
		"""
		Connects to a database, sets fields of the class
		"""
		super(PostgresProcessor, self).__init__()
		self.connection = psycopg2.connect(user=user, database=database)
		self.matrix_id = matrix_id
		self.cursor = self.connection.cursor()
		psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, self.cursor)
		self.data = {}

	def add_record(self,rel, head, modifier, k):
		"""
		Adds given record to a database. If such a record exists only the count column is updated.
		"""
		success = False
		forsha1= rel+head+modifier
		sha1 = hashlib.sha1(forsha1.encode('utf8')).hexdigest()
		while not success:
			try:
				self.cursor.execute("SELECT relation, head, modifier  FROM triples WHERE hash = %s AND matrix_id = %s", (sha1,self.matrix_id) )
				if self.cursor.rowcount == 0:
					q = "INSERT into triples VALUES(%s, %s, %s,%s, %s, %s);"
					self.cursor.execute(q, (self.matrix_id, rel, head, modifier, sha1, k) )
				else:
					found = False
					for (r,h,m) in self.cursor:
						if h == head and rel == r and m == modifier:
							found = True
							break
					if found:
						q = 'UPDATE triples SET amount = amount + %s WHERE matrix_id = %s AND relation = %s AND head = %s AND modifier = %s;'
						self.cursor.execute(q, (k,self.matrix_id, rel, head, modifier) )
					else:
						q = "INSERT into triples VALUES(%s, %s, %s,%s, %s, %s);"
						self.cursor.execute(q, (self.matrix_id, rel, head, modifier, sha1, k) )
				self.connection.commit()
				success = True
			except psycopg2.Error as e:
				print e
				self.connection.rollback()

	def commit(self):
		"""
		Commits all the gathered data to the database. Clears gathered data.
		"""
		for k  in self.data:
			h, r, m = k
			self.add_record(r,h,m, self.data[k])
		self.data = {}
