from typing import Tuple, List, Optional
import json
import sys
import os
import shlex
import asyncio
import argparse
import logging
import tempfile
from urllib.parse import urlparse

logger = logging.getLogger(__name__)

def find_sqlmat_json() -> Optional[dict]:
    json_path = os.getenv('SQLMAT_JSON_PATH')
    if json_path:
        with open(json_path) as f:
            cfg = json.load(f)
            return cfg

    # iterate through the current dir up to the root dir "/" to find a
    # .sqlmat.json
    workdir = os.path.abspath(os.getcwd())
    while workdir:
        json_path = os.path.join(workdir, '.sqlmat.json')
        if os.path.exists(json_path):
            with open(json_path) as f:
                cfg = json.load(f)
                return cfg
        parentdir = os.path.abspath(os.path.join(workdir, '..'))
        if parentdir == workdir:
            break
        workdir = parentdir
    logger.warning('fail to find .sqlmat.json')
    return None

def find_dsn(prog: str, desc: str) -> Tuple[str, List[str]]:
    parser = argparse.ArgumentParser(
        prog=prog,
        description=desc)
    parser.add_argument('-d', '--dsn',
                        type=str,
                        help='postgresql dsn')

    parser.add_argument('-g', '--db',
                        type=str,
                        default='default',
                        help='postgresql db instance defined in .sqlmat.json')

    parser.add_argument('callee_args',
                        type=str,
                        nargs='*',
                        help='command line arguments of callee programs')

    # from arguments
    args = parser.parse_args()
    if args.dsn:
        return args.dsn, args.callee_args

    # find dsn from ./.sqlmat.json
    cfg = find_sqlmat_json()
    if cfg:
        dsn = cfg['databases'][args.db]['dsn']
        assert isinstance(dsn, str)
        return dsn, args.callee_args

    # default dsn using username
    user = os.getenv('USER', '')
    default_dsn = f'postgres://{user}@127.0.0.1:5432/{args.db}'

    logger.warning('no postgres dsn specified, use %s instead', default_dsn)
    return default_dsn, args.callee_args

def joinargs(callee_args: List[str]) -> str:
    if hasattr(shlex, 'join'):
        return shlex.join(callee_args)
    else:
        return ' '.join(shlex.quote(a) for a in callee_args)

# run psql client
async def run_shell(dsn: str, callee_args: List[str]) -> None:
    p = urlparse(dsn)
    username = p.username or ''
    password = p.password or ''
    dbname = p.path[1:]
    hostname = p.hostname
    port = p.port or 5432

    temp_pgpass = tempfile.NamedTemporaryFile(mode='w')
    print(
        '{}:{}:{}:{}:{}'.format(hostname, port, dbname, username, password),
            file=temp_pgpass,
            flush=True)
    os.environ['PGPASSFILE'] = temp_pgpass.name
    command = 'psql -h{} -p{} -U{} {} {}'.format(hostname, port, username, joinargs(callee_args), dbname)
    proc = await asyncio.create_subprocess_shell(command)
    await proc.communicate()

def cl_run_shell() -> None:
    dsn, callee_args = find_dsn('sqlmat-shell', 'run psql client shell')
    loop = asyncio.get_event_loop()
    loop.run_until_complete(run_shell(dsn, callee_args))

# run dbdump
async def run_dbdump(dsn: str, callee_args: List[str]) -> None:
    p = urlparse(dsn)
    username = p.username or ''
    password = p.password or ''
    dbname = p.path[1:]
    hostname = p.hostname
    port = p.port or 5432

    temp_pgpass = tempfile.NamedTemporaryFile(mode='w')
    print(
        '{}:{}:{}:{}:{}'.format(hostname, port, dbname, username, password),
        file=temp_pgpass,
        flush=True)
    os.environ['PGPASSFILE'] = temp_pgpass.name
    command = 'pg_dump -h{} -p{} -U{} {} {}'.format(hostname, port, username, joinargs(callee_args), dbname)
    proc = await asyncio.create_subprocess_shell(command)
    await proc.communicate()

def cl_run_dbdump() -> None:
    dsn, callee_args = find_dsn('sqlmat-dump', 'dump database')
    loop = asyncio.get_event_loop()
    loop.run_until_complete(run_dbdump(dsn, callee_args))

# generate alembic migrations
def gen_migrate(dsn: str) -> None:
    init_data = ALEMBIC_INIT.replace('{{dsn}}', dsn)
    with open('alembic.ini', 'w') as f:
        f.write(init_data)

def cl_gen_migrate() -> None:
    dsn, callee_args = find_dsn('sqlmat-genmigrate', 'generate alembic migration')
    gen_migrate(dsn)
    print('Wrote alembic.ini')

ALEMBIC_INIT = '''\
# A generic, single database configuration.

[alembic]
# path to migration scripts
script_location = migrations

# template used to generate migration files
# file_template = %%(rev)s_%%(slug)s

# timezone to use when rendering the date
# within the migration file as well as the filename.
# string value is passed to dateutil.tz.gettz()
# leave blank for localtime
# timezone =

# max length of characters to apply to the
# "slug" field
#truncate_slug_length = 40

# set to 'true' to run the environment during
# the 'revision' command, regardless of autogenerate
# revision_environment = false

# set to 'true' to allow .pyc and .pyo files without
# a source .py file to be detected as revisions in the
# versions/ directory
# sourceless = false

# version location specification; this defaults
# to migrations/versions.  When using multiple version
# directories, initial revisions must be specified with --version-path
# version_locations = %(here)s/bar %(here)s/bat migrations/versions

# the output encoding used when revision files
# are written from script.py.mako
# output_encoding = utf-8

#sqlalchemy.url = driver://user:pass@localhost/dbname
sqlalchemy.url = {{dsn}}

# Logging configuration
[loggers]
keys = root,sqlalchemy,alembic

[handlers]
keys = console

[formatters]
keys = generic

[logger_root]
level = WARN
handlers = console
qualname =

[logger_sqlalchemy]
level = WARN
handlers =
qualname = sqlalchemy.engine

[logger_alembic]
level = INFO
handlers =
qualname = alembic

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatter_generic]
format = %(levelname)-5.5s [%(name)s] %(message)s
datefmt = %H:%M:%S

'''

[((12, 9, 12, 36), 'logging.getLogger', 'logging.getLogger', ({(12, 27, 12, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((15, 16, 15, 45), 'os.getenv', 'os.getenv', ({(15, 26, 15, 44): '"""SQLMAT_JSON_PATH"""'}, {}), "('SQLMAT_JSON_PATH')", False, 'import os\n'), ((38, 13, 40, 25), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((68, 11, 68, 32), 'os.getenv', 'os.getenv', ({(68, 21, 68, 27): '"""USER"""', (68, 29, 68, 31): '""""""'}, {}), "('USER', '')", False, 'import os\n'), ((82, 8, 82, 21), 'urllib.parse.urlparse', 'urlparse', ({(82, 17, 82, 20): 'dsn'}, {}), '(dsn)', False, 'from urllib.parse import urlparse\n'), ((89, 18, 89, 55), 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', (), '', False, 'import tempfile\n'), ((101, 11, 101, 35), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((106, 8, 106, 21), 'urllib.parse.urlparse', 'urlparse', ({(106, 17, 106, 20): 'dsn'}, {}), '(dsn)', False, 'from urllib.parse import urlparse\n'), ((113, 18, 113, 55), 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', (), '', False, 'import tempfile\n'), ((125, 11, 125, 35), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((23, 30, 23, 41), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((25, 20, 25, 57), 'os.path.join', 'os.path.join', ({(25, 33, 25, 40): 'workdir', (25, 42, 25, 56): '""".sqlmat.json"""'}, {}), "(workdir, '.sqlmat.json')", False, 'import os\n'), ((26, 11, 26, 36), 'os.path.exists', 'os.path.exists', ({(26, 26, 26, 35): 'json_path'}, {}), '(json_path)', False, 'import os\n'), ((76, 15, 76, 38), 'shlex.join', 'shlex.join', ({(76, 26, 76, 37): 'callee_args'}, {}), '(callee_args)', False, 'import shlex\n'), ((96, 17, 96, 57), 'asyncio.create_subprocess_shell', 'asyncio.create_subprocess_shell', ({(96, 49, 96, 56): 'command'}, {}), '(command)', False, 'import asyncio\n'), ((120, 17, 120, 57), 'asyncio.create_subprocess_shell', 'asyncio.create_subprocess_shell', ({(120, 49, 120, 56): 'command'}, {}), '(command)', False, 'import asyncio\n'), ((18, 18, 18, 30), 'json.load', 'json.load', ({(18, 28, 18, 29): 'f'}, {}), '(f)', False, 'import json\n'), ((30, 36, 30, 63), 'os.path.join', 'os.path.join', ({(30, 49, 30, 56): 'workdir', (30, 58, 30, 62): '""".."""'}, {}), "(workdir, '..')", False, 'import os\n'), ((28, 22, 28, 34), 'json.load', 'json.load', ({(28, 32, 28, 33): 'f'}, {}), '(f)', False, 'import json\n'), ((78, 24, 78, 38), 'shlex.quote', 'shlex.quote', ({(78, 36, 78, 37): 'a'}, {}), '(a)', False, 'import shlex\n')]