#!/usr/bin/env python3
'''
exmaple use:

./initdb.py sql db_name < schema.in > db.sql
mysql -h x.x.x.x -pxxxxxxx -u user -s < db.sql

./initdb.py sql db_name < schema.in | tee db.sql | mysql -h x.x.x.x -pxxxxxx -u user -s


# (HOST, USER, PASS) need to be set in config.py
# config.py should be in your PYTHONPATH or same dir with initdb.py
./initdb.py python db_name < schema.in | python
'''
import re
import sys


#################### utils ####################
def remove_cmt(line):
    return line.split('#')[0]


'''
col_info:
    name: column name
    type: (str) ex:
        int
        text
        float
        boolean
        varchar(255)
        datetime
        enum(...)
    null: True/False
    foreign: None | tuple(table_name, table_col)
'''
def parse_input(schema_raw):
    table_name, col_infos, primary_keys = '', [], []
    ret = []    # list of (table_name, col_infos, primary_keys)

    for line in schema_raw.splitlines():

        line = remove_cmt(line)
        if not line.strip():    # empty line
            continue

        if line[0] != ' ':      # new table
            if table_name:
                ret.append((table_name, col_infos, primary_keys))
                table_name, col_infos, primary_keys = '', [], []
            table_name = line.strip()

        else:                   # col info
            col_name, col_desc = line.strip().split(':')

            # find out true 'col_name', pick out primary key.
            if col_name[0] == '*':
                col_name = col_name[1:]
                primary_keys.append(col_name)

            col_desc = col_desc.strip()
            null = False


            # find out 'null'
            # assume after '|' only 'null'
            if len(col_desc.split('|')) > 1:
                col_desc = col_desc.split('|')[0].strip()
                null = True


            # find out 'foreign', 'type_'
            ptn = 'foreign key to (?P<tbl>.*?)(| *\((?P<col>.*)\))$'
            # foreign key to Feature
            # foreign key to Feature (f_id)
            r = re.search(ptn, col_desc)
            if r:
                d = r.groupdict()
                foreign = (d['tbl'], d['col'] if d['col'] else col_name)
                type_ = 'int'
            else:
                foreign = None
                type_ = col_desc


            # fix 'type_' if type is enum
            r = re.search('enum\((.*)\)', type_)
            if r:
                enum_keys = ['"%s"' % i.strip() for i in r.group(1).split(',')]
                type_ = 'enum(%s)' % ', '.join(enum_keys)
                

            col_infos.append({
                'name': col_name,
                'type': type_,
                'null': null,
                'foreign': foreign,
            })

    # add the last table.
    if table_name:
        ret.append((table_name, col_infos, primary_keys))
    return ret


#################### sql part ####################
def gen_sql_header():
    return '-- THIS FILE IS GENERATED BY initdb.py --\n\n'


def gen_sql_database(db_name):
    return '''
DROP DATABASE IF EXISTS {db_name};

CREATE DATABASE {db_name}
    DEFAULT CHARACTER SET utf8 COLLATE utf8_bin;
USE {db_name};
'''.format(db_name=db_name)


def gen_sql_schema(table_name, col_infos, primary_keys):
    '''
        table_name: str
        col_infos: [col_info, ...]
        primary_keys: [col_name, ...]
    '''
    col_lines = []
    for col in col_infos:

        line = '    %s %s' % (col['name'], col['type'])
        if not col['null']:
            line += ' NOT NULL'

        if len(primary_keys) == 1 and primary_keys[0] == col['name']:
            line += ' AUTO_INCREMENT'

        line += ','
        if col['foreign']:
            line += ' FOREIGN KEY ({}) REFERENCES {} ({}),'.format(
                    col['name'], col['foreign'][0], col['foreign'][1]
            )

        col_lines.append(line)

    output = '''
CREATE TABLE {table_name} (
{schema}
{primary_keys}
)
    ENGINE=InnoDB DEFAULT CHARACTER SET utf8 COLLATE utf8_bin;'''

    return output.format(
        table_name = table_name,
        schema = '\n'.join(col_lines),
        primary_keys = '    PRIMARY KEY (%s)' % ', '.join(primary_keys),
    )

def gen_sql_footer():
    return ''


#################### python part ####################
def gen_python_header():
    return '''#!/usr/bin/env python3
"""THIS FILE IS GENERATED BY initdb.py"""

from sqlalchemy import Column, ForeignKey
from sqlalchemy import Integer, Text, Float, Boolean, String, DateTime, Enum
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()
'''
 
    
def gen_python_database(db_name):
    try:
        from config import HOST, USER, PASS
    except ImportError:
        try: from .config import HOST, USER, PASS
        except ImportError:
            print('You should have config.py in your PYTHONPATH or same dir with initdb.py.')
            exit(1)

    return '''
from sqlalchemy import create_engine
import subprocess

HOST, USER, PASS = '{HOST}', '{USER}', '{PASS}'

# use `mysql` command to create/drop database.
print('[initdb] deleting old database...')
sql = """DROP DATABASE IF EXISTS {db_name};
CREATE DATABASE {db_name}
    DEFAULT CHARACTER SET utf8 COLLATE utf8_bin;
USE {db_name};
"""
subprocess.Popen(
    ['mysql', '-s', '-h', HOST, '-u', USER, '-p' + PASS],
    stdin=subprocess.PIPE,
).communicate(input=sql.encode())


print('[initdb] creating new database...')
url = 'mysql+mysqlconnector://{{user}}:{{pass_}}@{{host}}/{db_name}'.format(
    user=USER, pass_=PASS, host=HOST
)
engine = create_engine(url)
'''.format(HOST=HOST, USER=USER, PASS=PASS, db_name=db_name)


def gen_python_schema(table_name, col_infos, primary_keys):
    type_map = {
        'int': 'Integer',
        'text': 'Text',
        'float': 'Float',
        'boolean': 'Boolean',
        'char': 'String',
        'datetime': 'DateTime',
        'enum': 'Enum',
    }

    col_lines = []
    for col in col_infos:

        # dont replace thins in enum
        if 'enum' in col['type']:
            col['type'] = col['type'].replace('enum', 'Enum')
        else:
            for ori, new in type_map.items():
                col['type'] = col['type'].replace(ori, new)

        line = '    %s = Column(%s' % (col['name'], col['type'])
        if col['foreign']:
            line += ', ForeignKey("{}.{}")'.format(
                    col['foreign'][0], col['foreign'][1]
            )

        if col['name'] in primary_keys:
            line += ', primary_key=True'

            if len(primary_keys) != 1:
                line += ', autoincrement=False'

        if not col['null']:
            line += ', nullable=False'

        col_lines.append(line + ')')


    output = '''
class {table_name}(Base):
    __tablename__ = '{table_name}'
{schema}
'''
    return output.format(table_name=table_name, schema='\n'.join(col_lines))


def gen_python_footer():
    return '''
from sqlalchemy.orm import sessionmaker
session = sessionmaker()
session.configure(bind=engine)
Base.metadata.create_all(engine)
'''


#################### main ####################
def main(gen_type, db_name, schema_raw):
    ret = []

    if gen_type == 'sql':
        ret.append(gen_sql_header())
        ret.append(gen_sql_database(db_name))
        ret.append('\n-- -------------------- TABLES --------------------')
        for table in parse_input(schema_raw):
            ret.append(gen_sql_schema(*table))
        ret.append(gen_sql_footer())

    elif gen_type == 'python':
        ret.append(gen_python_header())
        ret.append(gen_python_database(db_name))
        for table in parse_input(schema_raw):
            ret.append(gen_python_schema(*table))
        ret.append(gen_python_footer())

    else:
        print('Unknown type. It should be one of "sql", "python"')
        exit(0)

    return '\n'.join(ret)


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print('Usage: %s sql|python db_name < schema.in' % sys.argv[0])
        exit(1)

    schema_raw = []
    while True:
        try:
            schema_raw.append(input())
        except EOFError:
            break

    output = main(sys.argv[1], sys.argv[2], '\n'.join(schema_raw))
    print(output)
