#!/usr/bin/env python

# Copyright (C) 2012-2013 SINA, All rights reserved.

"""Command line client for SAE MySQL Service. """

import sys
import os
import logging
import optparse

import sae._restful_mysql
import sae._restful_mysql._mysql_exceptions
sys.modules['_mysql_exceptions'] = sae._restful_mysql._mysql_exceptions

from grizzled import db
from grizzled.db import mysql
import prettytable
import sqlcmd
from sqlcmd import config

logging.basicConfig(level=logging.WARNING)
sqlcmd.log = logging.getLogger('cloudsql')

sqlcmd.DEFAULT_CONFIG_DIR = os.path.expanduser('~/.saecloud')
sqlcmd.RC_FILE = os.path.join(sqlcmd.DEFAULT_CONFIG_DIR, 'cloudsql.config')
sqlcmd.HISTORY_FILE_FORMAT = os.path.join(sqlcmd.DEFAULT_CONFIG_DIR, '%s.hist')
sqlcmd.INTRO = 'SAE MySQL Client\n\nType "help" or "?" for help.\n'

DEFAULT_ENCODING = 'utf-8'
USAGE = '%prog [options] database_name'

DEFAULT_SAE_MYSQL_HOST = 'w.rdc.sae.sina.com.cn'
DEFAULT_SAE_MYSQL_PORT = 3307
DEFAULT_SAE_MYSQL_DB_PREFIX = 'app_'

class CloudSqlDriver(mysql.MySQLDriver):
    """Grizzled DB Driver for Cloud SAE MySQL Service."""

    NAME = 'cloudsql'

    def get_import(self):
        return sae._restful_mysql

    def get_display_name(self):
        return 'Cloud SQL'

    def do_connect(self, host, port, user, password, database):
        # Fix grizzled's mysql driver which omit the port argument when connecting.
        dbi = self.get_import()
        port = port and int(port) or 3306
        return dbi.connect(host=host, user=user, passwd=password, db=database, port=port)

class CloudSqlCmd(sqlcmd.SQLCmd):
    """The SQLCmd command interpreter for Cloud SQL."""

    sqlcmd.SQLCmd.MAIN_PROMPT = 'mysql> '
    sqlcmd.SQLCmd.CONTINUATION_PROMPT = '    -> '

    sqlcmd.SQLCmd.NO_SEMI_NEEDED.update(
        ['about', 'desc', 'describe', 'echo', 'exit', 'h', 'hist',
         'history', 'load', 'run', 'r', 'redo', 'set', 'show', 'var', 'vars'])

    for method in ['do_dot_connect', 'do_dot_desc', 'do_begin']:
        delattr(sqlcmd.SQLCmd, method)

    for cmd in ['show', 'describe', 'echo', 'load', 'run', 'exit', 'h',
                'hist', 'history', 'var', 'vars', 'about']:
        method = 'do_dot_' + cmd
        setattr(sqlcmd.SQLCmd, method.replace('dot_', ''), getattr(
            sqlcmd.SQLCmd, method))
        delattr(sqlcmd.SQLCmd, method)
        method = 'complete_dot_' + cmd
        if hasattr(sqlcmd.SQLCmd, method):
            setattr(sqlcmd.SQLCmd, method.replace('dot_', ''), getattr(
                sqlcmd.SQLCmd, method))
            delattr(sqlcmd.SQLCmd, method)

    def do_redo(self, args):
        # XXX: Fix global name 'do_r' is not defined problem in sqlcmd
        self.do_r(args)

    def _SQLCmd__set_setting(self, varname, value):
        # XXX: Fix bool object has no lower attribute in sqlcmd
        return sqlcmd.SQLCmd._SQLCmd__set_setting(self, varname, str(value))

    def do_desc(self, args):
        self.do_describe(args, cmd='.desc')
    complete_desc = sqlcmd.SQLCmd.complete_dot_desc

    def do_load(self, args):
        self.do_run(args)

    def preloop(self, *args, **kwargs):
        sqlcmd.SQLCmd.preloop(self, *args, **kwargs)
        # Just exit if the connect failed
        if self._SQLCmd__db is None: sys.exit(1)

    def __init__(self, *args, **kwargs):
        sqlcmd.SQLCmd.__init__(self, *args, **kwargs)
        self.prompt = sqlcmd.SQLCmd.MAIN_PROMPT
        self.output_encoding = DEFAULT_ENCODING

    def set_output_encoding(self, encoding):
        self.output_encoding = encoding

    def _build_table(self, cursor):
        """Builds an output PrettyTable from the results in the given cursor."""
        if not cursor.description:
            return None

        column_names = [column[0] for column in cursor.description]
        table = prettytable.PrettyTable(column_names)
        rows = cursor.fetchall()
        if not rows:
            return table
        for i, col in enumerate(rows[0]):
            table.align[column_names[i]] = isinstance(col, basestring) and 'l' or 'r'
        for row in rows: table.add_row(row)
        return table

    def _SQLCmd__handle_select(self, args, cursor, command='select'):
        """Overrides SQLCmd.__handle_select to display output with prettytable."""
        self._SQLCmd__exec_SQL(cursor, command, args)
        table = self._build_table(cursor)
        if table:
            output = table.get_string()
            if isinstance(output, unicode):
                print output.encode(self.output_encoding)
            else:
                print output

def _create_config_dir():
    """Creates the sqlcmd config directory if necessary."""
    directory = sqlcmd.DEFAULT_CONFIG_DIR
    if not os.access(directory, os.R_OK | os.W_OK | os.X_OK):
        old_umask = os.umask(077)
        os.makedirs(sqlcmd.DEFAULT_CONFIG_DIR)
        os.umask(old_umask)

def main(argv):
    parser = optparse.OptionParser(usage=USAGE)
    parser.add_option('-u', '--username', dest='username',
                      help='MySQL username to use when connecting to the server.')
    parser.add_option('-p', '--password', dest='password',
                      help='MySQL password to use when connecting to the server.')
    parser.add_option('-e', '--output_encoding', dest='output_encoding',
                      default=DEFAULT_ENCODING,
                      help='Output encoding. Defaults to %s.' % DEFAULT_ENCODING)

    (options, args) = parser.parse_args(argv[1:])

    if len(args) != 1:
        parser.print_help(sys.stderr)
        return 1

    if not options.username or not options.password:
        print >>sys.stderr, 'Error: username or password is missing.\n'
        return 1

    if args[0].startswith(DEFAULT_SAE_MYSQL_DB_PREFIX):
        database_name = args[0]
    else:
        database_name = DEFAULT_SAE_MYSQL_DB_PREFIX + args[0]
    instance_alias = database_name

    _create_config_dir()

    db.add_driver(CloudSqlDriver.NAME, CloudSqlDriver)
    sql_cmd_config = config.SQLCmdConfig(None)
    sql_cmd_config.add('__cloudsql__', instance_alias,
                       DEFAULT_SAE_MYSQL_HOST , DEFAULT_SAE_MYSQL_PORT, database_name,
                       CloudSqlDriver.NAME, options.username, options.password)
    sql_cmd = CloudSqlCmd(sql_cmd_config)
    sql_cmd.set_output_encoding(options.output_encoding)
    sql_cmd.set_database(instance_alias)
    sql_cmd.cmdloop()

    return 0

if __name__ == '__main__':
    sys.exit(main(sys.argv))