#!/usr/bin/env python3
"sa_harness.py -- create sqlalchemy definitions from create table & index stmts"

import sqlparse, glob, collections, argparse
from . import wrappers, diffing

PREAMBLE = """# autogenerated by sa_harness.py
import sqlalchemy as sa"""

def read_glob_stmts(glob_pattern):
  "return wrapped stmts for files matching glob"
  stmts = []
  for fname in glob.glob(glob_pattern):
    stmts.extend(map(wrappers.wrap, sqlparse.parse(open(fname).read())))
  return stmts

def column_args(column, pkey_fields):
  "parse sql column in a horrible way and return args for column constructor"
  parsed = column.parse()
  assert parsed.success # todo: does parse() log something on failure?
  dets = {'name': parsed.name, 'type': parsed.type}
  if parsed.pkey or parsed.name in pkey_fields:
    dets['primary_key'] = True
  if parsed.default:
    # todo: figure out escaping here
    dets['server_default'] = "sa.text('%s')" % parsed.default
  if parsed.unique:
    raise NotImplementedError('todo: sa_harness unique keys')
  if parsed.not_null:
    raise NotImplementedError('todo: sa_harness not null')
  return dets.pop('name'), dets.pop('type'), dets

TYPES = {
  'int': 'Integer',
  'jsonb': 'JSONB',
  'json': 'JSON',
  'text': 'Text',
  'uuid': 'UUID',
  'timestamp': 'DateTime',
}

def render_col(col_name, col_type, details, composite_pkey):
  if composite_pkey and col_name in composite_pkey:
    details['primary_key'] = True
  return f"sa.Column('{col_name}', sa.{TYPES[col_type]}{', ' + ', '.join('%s=%s' % pair for pair in details.items()) if details else ''})"

def transform(table_stmts, delim=',\n  '):
  "return some representation of sqlalchemy tables"
  table_strings = []
  for tablename, stmts in diffing.group_by_table(table_stmts).items():
    indexes = []
    table = None
    cols = collections.OrderedDict()
    composite_pkey = None
    for stmt in stmts:
      if isinstance(stmt, wrappers.CreateTable):
        pkey_fields = stmt.pkey_fields()
        for col in stmt.columns():
          name, type_, dets = column_args(col, pkey_fields)
          cols[name, type_] = dets
      elif isinstance(stmt, wrappers.CreateIndex):
        paren = stmt.decl()[-1]
        assert isinstance(paren, sqlparse.sql.Parenthesis)
        index_cols = ['%s_table.c.%s' % (stmt.table, ident) for ident, in wrappers.split_pun(paren)]
        indexes.append(f"{stmt.index_name}_index = sa.Index('{stmt.index_name}', {', '.join(index_cols)})")
      else:
        raise TypeError('unhandled statement type', type(stmt))
    if cols:
      col_strings = [
        render_col(col_name, col_type, details, composite_pkey)
        for (col_name, col_type), details in cols.items()
      ]
      table_strings.append(f"{tablename}_table = sa.Table('{tablename}', META{delim}{delim.join(col_strings)})")
    table_strings.extend(indexes)
  return table_strings

PARSER = argparse.ArgumentParser(description="generate sqlalchemy models from create table stmts")
PARSER.add_argument('glob')
PARSER.add_argument('--target', help="optional output path; otherwise stdout")

def main():
  args = PARSER.parse_args()
  stmts = read_glob_stmts(args.glob)
  value = '\n\n'.join([PREAMBLE] + transform(stmts))
  if args.target:
    open(args.target, 'w').write(value)
  else:
    print(value)

if __name__ == '__main__':
  main()
