#!/usr/bin/python
#
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Library for using Logica in CoLab."""

from .common import color
from .common import concertina_lib

from .compiler import functors
from .compiler import rule_translate
from .compiler import universe

import IPython

from IPython.core.magic import register_cell_magic
from IPython.display import display

import os

import pandas

from .parser_py import parse
from .common import sqlite3_logica

BQ_READY = True  # By default.

try:
  from google.cloud import bigquery
except:
  BQ_READY = False
  print('Could not import google.cloud.bigquery.')

try:
  from google.colab import auth
except:
  BQ_READY = False
  print('Could not import google.cloud.auth.')

try:
  from google.colab import widgets
  WIDGETS_IMPORTED = True
except:
  WIDGETS_IMPORTED = False
  print('Could not import google.colab.widgets.')

PROJECT = None

# TODO: Should this be renamed to PSQL_ENGINE, PSQL_CONNECTION?
DB_ENGINE = None
DB_CONNECTION = None

USER_AUTHENTICATED = False

TABULATED_OUTPUT = True

SHOW_FULL_QUERY = True

PREAMBLE = None


def SetPreamble(preamble):
  global PREAMBLE
  PREAMBLE = preamble

def SetProject(project):
  global PROJECT
  PROJECT = project

def SetDbConnection(connection):
  global DB_CONNECTION
  DB_CONNECTION = connection

def EnsureAuthenticatedUser():
  global USER_AUTHENTICATED
  global PROJECT
  if USER_AUTHENTICATED:
    return
  auth.authenticate_user()
  if PROJECT is None:
    print("Please enter project_id to use for BigQuery queries.")
    PROJECT = input()
    print("project_id is set to %s" % PROJECT)
    print("You can change it with logica.colab_logica.SetProject command.")
  USER_AUTHENTICATED = True

def SetTabulatedOutput(tabulated_output):
  global TABULATED_OUTPUT
  global SHOW_FULL_QUERY
  TABULATED_OUTPUT = tabulated_output
  SHOW_FULL_QUERY = TABULATED_OUTPUT

if not WIDGETS_IMPORTED:
  SetTabulatedOutput(False)

def TabBar(*args):
  """Returns a real TabBar or a mock. Useful for UIs that don't support JS."""
  if TABULATED_OUTPUT:
    return widgets.TabBar(*args)
  class MockTab:
      def __init__(self):
          pass
      def __enter__(self):
          pass
      def __exit__(self, *x):
          pass
  class MockTabBar:
      def __init__(self):
          pass
      def output_to(self, x):
          return MockTab()
  return MockTabBar()

@register_cell_magic
def logica(line, cell):
  Logica(line, cell, run_query=True)


def ParseList(line):
  line = line.strip()
  if not line:
    predicates = []
  else:
    predicates = [p.strip() for p in line.split(',')]
  return predicates


def RunSQL(sql, engine, connection=None, is_final=False):
  if engine == 'bigquery':
    client = bigquery.Client(project=PROJECT)
    return client.query(sql).to_dataframe()
  elif engine == 'psql':
    if is_final:
      return pandas.read_sql(sql, connection)
    else:
      return connection.execute(sql)
  elif engine == 'sqlite':
    statements = parse.SplitRaw(sql, ';')
    connection.executescript(sql)
    if is_final:
      return pandas.read_sql(statements[-1], connection)
    else:
      pass
    return None
  else:
    raise Exception('Logica only supports BigQuery, PostgreSQL and SQLite '
                    'for now.')


class SqliteRunner(object):
  def __init__(self):
    self.connection = sqlite3_logica.SqliteConnect()
  
  # TODO: Sqlite runner should not be accepting an engine.
  def __call__(self, sql, engine, is_final):
    return RunSQL(sql, engine, self.connection, is_final)


class PostgresRunner(object):
  def __init__(self):
    global DB_CONNECTION
    global DB_ENGINE
    if DB_CONNECTION:
      self.engine = DB_ENGINE
      self.connection = DB_CONNECTION
    else:
      (self.engine, self.connection) = PostgresJumpStart()
      DB_ENGINE = self.engine
      DB_CONNECTION = self.connection
  
  def  __call__(self, sql, engine, is_final):
    return RunSQL(sql, engine, self.connection, is_final)


def ShowError(error_text):
  print(color.Format('[ {error}Error{end} ] ' + error_text))


def Logica(line, cell, run_query):
  """Running Logica predicates and storing results."""
  predicates = ParseList(line)
  if not predicates:
    ShowError('No predicates to run.')
    return
  try:
    program = ';\n'.join(s for s in [PREAMBLE, cell] if s)
    parsed_rules = parse.ParseFile(program)['rule']
  except parse.ParsingException as e:
    e.ShowMessage()
    return
  try:
    program = universe.LogicaProgram(parsed_rules)
  except functors.FunctorError as e:
    e.ShowMessage()
    return

  engine = program.annotations.Engine()

  if engine == 'bigquery' and not BQ_READY:
    ShowError(
      'BigQuery client and/or authentification is not installed. \n'
      'It is the easiest to run BigQuery requests from Google CoLab:\n'
      '  https://colab.research.google.com/.\n'
      'Note that running Logica on SQLite requires no installation.\n'
      'This could be a good fit for working with small data or learning Logica.\n'
      'Use {warning}@Engine("sqlite");{end} annotation in your program to use SQLite.')
    return

  bar = TabBar(predicates + ['(Log)'])
  logs_idx = len(predicates)

  executions = []
  sub_bars = []
  ip = IPython.get_ipython()
  for idx, predicate in enumerate(predicates):
    with bar.output_to(logs_idx):
      try:
        sql = program.FormattedPredicateSql(predicate)
        executions.append(program.execution)
        ip.push({predicate + '_sql': sql})
      except rule_translate.RuleCompileException as e:
        print('Encountered error when compiling %s.' % predicate)
        e.ShowMessage()
        return
    # Publish output to Colab cell.
    with bar.output_to(idx):
      sub_bar = TabBar(['SQL', 'Result'])
      sub_bars.append(sub_bar)
      with sub_bar.output_to(0):
        if SHOW_FULL_QUERY:
          print(
              color.Format(
                  'The following query is stored at {warning}%s{end} '
                  'variable.' % (
                      predicate + '_sql')))
          print(sql)
        else:
          print('Query is stored at %s variable.' %
                color.Warn(predicate + '_sql'))

  with bar.output_to(logs_idx):
    if engine == 'sqlite':
      sql_runner = SqliteRunner()
    elif engine == 'psql':
      sql_runner = PostgresRunner()
    elif engine == 'bigquery':
      EnsureAuthenticatedUser()
      sql_runner = RunSQL
    else:
      raise Exception('Logica only supports BigQuery, PostgreSQL and SQLite '
                      'for now.')   
                      
    result_map = concertina_lib.ExecuteLogicaProgram(
      executions, sql_runner=sql_runner, sql_engine=engine)

  for idx, predicate in enumerate(predicates):
    t = result_map[predicate]
    ip.push({predicate: t})
    with bar.output_to(idx):
      with sub_bars[idx].output_to(1): 
        if run_query:
          print(
              color.Format(
                  'The following table is stored at {warning}%s{end} '
                  'variable.' %
                  predicate))
          display(t)  
        else:
          print('The query was not run.')
      print(' ') # To activate the tabbar.

def PostgresJumpStart():
  # Install postgresql server.
  print("Installing and configuring an empty PostgreSQL database.")
  result = 0
  result += os.system('sudo apt-get -y -qq update')
  result += os.system('sudo apt-get -y -qq install postgresql')
  result += os.system('sudo service postgresql start')
  # Ignoring user creation error, as they may already exist.
  result += 0 * os.system(
    'sudo -u postgres psql -c "CREATE USER logica WITH SUPERUSER"')
  result += os.system(
    'sudo -u postgres psql -c "ALTER USER logica PASSWORD \'logica\';"')
  result += os.system(
    'sudo -u postgres psql -U postgres -c \'CREATE DATABASE logica;\'')
  if result != 0:
    print("""Installation failed. Please try the following manually:
# Install Logica.
!pip install logica

# Install postgresql server.
!sudo apt-get -y -qq update
!sudo apt-get -y -qq install postgresql
!sudo service postgresql start

# Prepare database for Logica.
!sudo -u postgres psql -c "CREATE USER logica WITH SUPERUSER"
!sudo -u postgres psql -c "ALTER USER logica PASSWORD 'logica';"
!sudo -u postgres psql -U postgres -c 'CREATE DATABASE logica;'

# Connect to the database.
from logica import colab_logica
from sqlalchemy import create_engine
import pandas
engine = create_engine('postgresql+psycopg2://logica:logica@127.0.0.1', pool_recycle=3600);
connection = engine.connect();
colab_logica.SetDbConnection(connection)""")
    return
  print('Installation succeeded. Connecting...')
  # Connect to the database.
  from logica import colab_logica
  from sqlalchemy import create_engine
  import pandas
  engine = create_engine('postgresql+psycopg2://logica:logica@127.0.0.1', pool_recycle=3600)
  connection = engine.connect()
  print('Connected.')
  return engine, connection

[((91, 2, 91, 26), 'google.colab.auth.authenticate_user', 'auth.authenticate_user', ({}, {}), '()', False, 'from google.colab import auth\n'), ((226, 7, 226, 28), 'IPython.get_ipython', 'IPython.get_ipython', ({}, {}), '()', False, 'import IPython\n'), ((288, 12, 288, 51), 'os.system', 'os.system', ({(288, 22, 288, 50): '"""sudo apt-get -y -qq update"""'}, {}), "('sudo apt-get -y -qq update')", False, 'import os\n'), ((289, 12, 289, 63), 'os.system', 'os.system', ({(289, 22, 289, 62): '"""sudo apt-get -y -qq install postgresql"""'}, {}), "('sudo apt-get -y -qq install postgresql')", False, 'import os\n'), ((290, 12, 290, 54), 'os.system', 'os.system', ({(290, 22, 290, 53): '"""sudo service postgresql start"""'}, {}), "('sudo service postgresql start')", False, 'import os\n'), ((294, 12, 295, 72), 'os.system', 'os.system', ({(295, 4, 295, 71): '"""sudo -u postgres psql -c "ALTER USER logica PASSWORD \'logica\';\\""""'}, {}), '(\'sudo -u postgres psql -c "ALTER USER logica PASSWORD \\\'logica\\\';"\')', False, 'import os\n'), ((296, 12, 297, 71), 'os.system', 'os.system', ({(297, 4, 297, 70): '"""sudo -u postgres psql -U postgres -c \'CREATE DATABASE logica;\'"""'}, {}), '("sudo -u postgres psql -U postgres -c \'CREATE DATABASE logica;\'")', False, 'import os\n'), ((326, 11, 326, 92), 'sqlalchemy.create_engine', 'create_engine', (), '', False, 'from sqlalchemy import create_engine\n'), ((111, 11, 111, 32), 'google.colab.widgets.TabBar', 'widgets.TabBar', ({(111, 26, 111, 31): '*args'}, {}), '(*args)', False, 'from google.colab import widgets\n'), ((142, 13, 142, 45), 'google.cloud.bigquery.Client', 'bigquery.Client', (), '', False, 'from google.cloud import bigquery\n'), ((292, 16, 293, 67), 'os.system', 'os.system', ({(293, 4, 293, 66): '"""sudo -u postgres psql -c "CREATE USER logica WITH SUPERUSER\\""""'}, {}), '(\'sudo -u postgres psql -c "CREATE USER logica WITH SUPERUSER"\')', False, 'import os\n'), ((146, 13, 146, 45), 'pandas.read_sql', 'pandas.read_sql', ({(146, 29, 146, 32): 'sql', (146, 34, 146, 44): 'connection'}, {}), '(sql, connection)', False, 'import pandas\n'), ((153, 13, 153, 56), 'pandas.read_sql', 'pandas.read_sql', ({(153, 29, 153, 43): 'statements[-1]', (153, 45, 153, 55): 'connection'}, {}), '(statements[-1], connection)', False, 'import pandas\n'), ((279, 10, 279, 20), 'IPython.display.display', 'display', ({(279, 18, 279, 19): 't'}, {}), '(t)', False, 'from IPython.display import display\n')]