import grakn
import xml.etree.cElementTree as et

def build_phone_call_graph(inputs):
  client = grakn.Grakn(uri = "localhost:48555")
  with client.session(keyspace = "phone_calls") as session:
    for input in inputs:
      print("Loading from [" + input["data_path"] + "] into Grakn ...")
      load_data_into_grakn(input, session)

def load_data_into_grakn(input, session):
  items = parse_data_to_dictionaries(input)

  for item in items:
    with session.transaction(grakn.TxType.WRITE) as tx:
      graql_insert_query = input["template"](item)
      print("Executing Graql Query: " + graql_insert_query)
      tx.query(graql_insert_query)
      tx.commit()

  print("\nInserted " + str(len(items)) + " items from [ " + input["data_path"] + "] into Grakn.\n")

def company_template(company):
  return 'insert $company isa company has name "' + company["name"] + '";'

def person_template(person):
  # insert person
  graql_insert_query = 'insert $person isa person has phone-number "' + person["phone_number"] + '"'
  if "first_name" in person:
    # person is a customer
    graql_insert_query += " has is-customer true"
    graql_insert_query += ' has first-name "' + person["first_name"] + '"'
    graql_insert_query += ' has last-name "' + person["last_name"] + '"'
    graql_insert_query += ' has city "' + person["city"] + '"'
    graql_insert_query += " has age " + str(person["age"])
  else:
    # person is not a customer
    graql_insert_query += " has is-customer false"
  graql_insert_query += ";"
  return graql_insert_query

def contract_template(contract):
  # match company
  graql_insert_query = 'match $company isa company has name "' + contract["company_name"] + '";'
  # match person
  graql_insert_query += ' $customer isa person has phone-number "' + contract["person_id"] + '";'
  # insert contract
  graql_insert_query += " insert (provider: $company, customer: $customer) isa contract;"
  return graql_insert_query

def call_template(call):
  # match caller
  graql_insert_query = 'match $caller isa person has phone-number "' + call["caller_id"] + '";'
  # match callee
  graql_insert_query += ' $callee isa person has phone-number "' + call["callee_id"] + '";'
  # insert call
  graql_insert_query += " insert $call(caller: $caller, callee: $callee) isa call; $call has started-at " + call["started_at"] + "; $call has duration " + str(call["duration"]) + ";"
  return graql_insert_query

def parse_data_to_dictionaries(input):
  items = []
  with open(input["data_path"] + ".xml", "rb") as inputfile:
    append = False
    for line in inputfile:
      if "<" + input["selector"] + ">" in str(line):
        ## start of a new xml tag
        buffer = line
        append = True
      elif "</" + input["selector"]  + ">" in str(line):
        ## we got a complete xml tag
        buffer += line
        append = False
        tnode = et.fromstring(buffer)
        ## parse the tag to a dictionary and append to tiems
        item = {}
        for element in tnode.getchildren():
          item[element.tag] = element.text
        items.append(item)
        ## delete the buffer to free the memory
        del buffer
      elif append:
        ## inside the current xml tag
        buffer += line
  return items

inputs = [
  {
    "data_path": "./data/companies",
    "template": company_template,
    "selector": "company"
  },
  {
    "data_path": "./data/people",
    "template": person_template,
    "selector": "person"
  },
  {
    "data_path": "./data/contracts",
    "template": contract_template,
    "selector": "contract"
  },
  {
    "data_path": "./data/calls",
    "template": call_template,
    "selector": "call"
  }
]

build_phone_call_graph(inputs)