#!/usr/bin/env python2
# coding: utf-8

from contextlib import closing
from datetime import timedelta
import opster
import psycopg2
import sys
import time


ROBOT_UID = 82282794


def chunked(items, chunk_size):
    for i in xrange(0, len(items), chunk_size):
        yield items[i:i + chunk_size]


def fetch_existing_uids(mpro_conn):
    with closing(mpro_conn.cursor()) as cursor:
        cursor.execute('SELECT uid FROM acl.user')
        return set(r[0] for r in cursor)


def edits_level(edits):
    level = 0
    while edits >= 10:
        edits /= 10
        level += 2
    if edits >= 5:
        level += 1
    return min(11, level)


@opster.command()
def cs_exp(outpath,
           wiki_conn_str=('', '', 'wiki9 connection string')):
    """Export active users changeset statistics (uid, total, min id) (wiki9)"""
    print >>sys.stderr, 'Exporting active user stats...',
    with closing(open(outpath, 'w')) as outfile:
        with closing(psycopg2.connect(wiki_conn_str)) as wiki_conn:
            with closing(wiki_conn.cursor()) as cursor:
                cursor.execute(
                    'SELECT created_by, count(*), min(id) '
                        'FROM core.changesets '
                        'GROUP BY created_by')
                for r in cursor:
                    print >>outfile, r[0], r[1], r[2]
    print >>sys.stderr, 'done'


@opster.command()
def users_exp(cspath,
              outpath,
              wiki_conn_str=('', '', 'wiki9 connection string')):
    """Export users info (uid, registration timestamp, login) (lback)"""
    import blackbox

    print >>sys.stderr, 'Loading active users...',
    with closing(open(cspath, 'r')) as csfile:
        min_csid_by_uid = {}
        for r in csfile:
            uid, _, min_csid = map(long, r.split())
            if uid / (10 ** 13) == 112: # y-t uid
                continue
            min_csid_by_uid[uid] = min_csid
    print >>sys.stderr, 'done'

    print >>sys.stderr, 'Fetching registration dates...',
    with closing(psycopg2.connect(wiki_conn_str)) as wiki_conn:
        reg_ts_by_uid = {}
        missing_ts_uids = set()
        for uid_chunk in chunked(min_csid_by_uid.keys(), 1000):
            with closing(wiki_conn.cursor()) as cursor:
                cursor.execute(
                    'SELECT uid, registration_date '
                        'FROM social.user_settings '
                        'WHERE uid IN (%s)' % ','.join(map(str, uid_chunk)))
                for r in cursor:
                    if r[1] is None:
                        missing_ts_uids.add(long(r[0]))
                    else:
                        reg_ts_by_uid[long(r[0])] = \
                            long(time.mktime(r[1].timetuple()))

        # for missing reg dates, take the ts of first changeset
        with closing(wiki_conn.cursor()) as cursor:
            cursor.execute(
                'SELECT DISTINCT ON (changeset_id) created_by, created '
                    'FROM core.revisions '
                    'WHERE changeset_id IN (%s)' % ','.join(
                        str(min_csid_by_uid[m_uid]) for m_uid in missing_ts_uids))
            for r in cursor:
                uid = long(r[0])
                missing_ts_uids.remove(uid)
                reg_ts_by_uid[uid] = long(time.mktime(r[1].timetuple()))
    print >>sys.stderr, 'done (%s users not registered!)' % (len(min_csid_by_uid) - len(reg_ts_by_uid))

    print >>sys.stderr, 'Fetching logins from blackbox...',
    login_by_uid = {}
    bb = blackbox.JsonBlackbox()
    for uid_chunk in chunked(reg_ts_by_uid.keys(), 100):
        userinfos = bb.userinfo(uid=uid_chunk, userip='127.0.0.1', dbfields=['accounts.login.uid'])
        if 'users' not in userinfos:
            raise RuntimeError(str(userinfos))
        for info in userinfos['users']:
            if not info['uid']: # deleted user
                continue
            login = str(info['dbfields']['accounts.login.uid'])
            if not login or ' ' in login:
                raise RuntimeError('invalid login "' + login + '" for uid ' + info['id'])
            login_by_uid[long(info['id'])] = login

    print >>sys.stderr, 'done (%s deleted users)' % (len(reg_ts_by_uid) - len(login_by_uid))

    print >>sys.stderr, 'Writing results...',
    with closing(open(outpath, 'w')) as outfile:
        for uid, login in login_by_uid.iteritems():
            print >>outfile, uid, reg_ts_by_uid[uid], login
    print >>sys.stderr, 'done'


@opster.command()
def users_imp(upath,
              mpro_conn_str=('', '', 'mpro connection string')):
    """Import users from file in 'deleted' state with 'nmaps' group (um)"""
    print >>sys.stderr, 'Importing users...',
    with closing(psycopg2.connect(mpro_conn_str)) as mpro_conn:
        existing_uids = fetch_existing_uids(mpro_conn)

        with closing(mpro_conn.cursor()) as cursor:
            cursor.execute("SELECT id FROM acl.group WHERE name = 'nmaps'")
            r = cursor.fetchone()
            if r is None:
                raise RuntimeError("missing 'nmaps' group")
            nmaps_group_id = r[0]

        with closing(mpro_conn.cursor()) as cursor:
            cursor.execute(
                'PREPARE UpdateUser (bigint, bigint) AS '
                    'UPDATE acl.user '
                        'SET created_by = $1, '
                            'created = LEAST(created, to_timestamp($2)) '
                        'WHERE uid = $1 '
                        'RETURNING id')
            cursor.execute(
                'PREPARE CreateUser (bigint, text, bigint) AS '
                    'INSERT INTO acl.user '
                        '(uid, login, created, created_by, modified_by, status) '
                        "VALUES ($1, $2, to_timestamp($3), $1, %s, 'deleted') "
                        'RETURNING id', (ROBOT_UID,))
            cursor.execute(
                'PREPARE AssignNmaps (bigint) AS '
                    'WITH newrows AS ('
                        'SELECT %s, $1 WHERE NOT EXISTS ('
                            'SELECT * FROM acl.group_user '
                                'WHERE group_id = %s AND user_id = $1)'
                    ') INSERT INTO acl.group_user SELECT * FROM newrows',
                (nmaps_group_id, nmaps_group_id))

        with closing(mpro_conn.cursor()) as cursor:
            with closing(open(upath, 'r')) as ufile:
                for r in ufile:
                    uid, reg_ts, login = r.split()
                    uid = long(uid)
                    reg_ts = long(reg_ts)
                    if long(uid) in existing_uids:
                        cursor.execute('EXECUTE UpdateUser (%s, %s)', (uid, reg_ts))
                    else:
                        cursor.execute('EXECUTE CreateUser (%s, %s, %s)', (uid, login, reg_ts))
                    acl_id = cursor.fetchone()[0]
                    cursor.execute('EXECUTE AssignNmaps (%s)', (acl_id,))
        mpro_conn.commit()
    print >>sys.stderr, 'done'


@opster.command()
def bans(wiki_conn_str=('', '', 'wiki9 connection string'),
         mpro_conn_str=('', '', 'mpro connection string')):
    """Convert ban history (wiki9)"""
    with closing(psycopg2.connect(mpro_conn_str)) as mpro_conn:
        print >>sys.stderr, 'Fetching current bans summary...',
        existing_uids = fetch_existing_uids(mpro_conn)
        latest_ban_dt_by_uid = {}

        with closing(mpro_conn.cursor()) as mpro_cursor:
            mpro_cursor.execute(
                'SELECT br_uid, MAX(br_created) '
                    'FROM acl.ban_record '
                    'GROUP BY br_uid')
            for r in mpro_cursor:
                latest_ban_dt_by_uid[long(r[0])] = r[1]
        print >>sys.stderr, 'done'

        print >>sys.stderr, 'Converting ban records...',
        with closing(psycopg2.connect(wiki_conn_str)) as wiki_conn:
            with closing(mpro_conn.cursor()) as mpro_cursor:
                mpro_cursor.execute(
                    'PREPARE CreateBan AS '
                        'INSERT INTO acl.ban_record '
                            '(br_uid, br_action, br_created_by, br_created, br_expires) '
                            'VALUES ($1, $2, $3, $4, $5) '
                            'RETURNING br_id')
                mpro_cursor.execute(
                    'PREPARE UpdateCurrentBan AS '
                        'UPDATE acl.user '
                            'SET current_ban_id = $2 '
                            'WHERE uid = $1')

                with closing(wiki_conn.cursor()) as wiki_cursor:
                    wiki_cursor.execute(
                        'SELECT uid, start_date, days, moderator '
                            'FROM social.user_ban_history '
                            'ORDER BY start_date ASC')
                    for r in wiki_cursor:
                        uid = long(r[0])
                        ban_dt = r[1]
                        days = r[2]
                        moderator = long(r[3])
                        if uid not in existing_uids or moderator not in existing_uids:
                            continue
                        # skip possibly imported bans
                        latest_ban_dt = latest_ban_dt_by_uid.get(uid)
                        if latest_ban_dt is not None and latest_ban_dt >= ban_dt:
                            continue

                        mpro_cursor.execute(
                            'EXECUTE CreateBan (%s, %s, %s, %s, %s)',
                            (uid,
                             'ban' if days >= 0 else 'unban',
                             moderator,
                             ban_dt,
                             ban_dt + timedelta(days=days) if days > 0 else None))
                        ban_id = mpro_cursor.fetchone()[0]
                        mpro_cursor.execute(
                            'EXECUTE UpdateCurrentBan (%s, %s)',
                            (uid, ban_id))
        mpro_conn.commit()
    print >>sys.stderr, 'done'


@opster.command()
def badges(cspath,
           mpro_conn_str=('', '', 'mpro connection string')):
    """Award badges (um)"""
    with closing(psycopg2.connect(mpro_conn_str)) as mpro_conn:
        existing_uids = fetch_existing_uids(mpro_conn)

        with closing(mpro_conn.cursor()) as cursor:
            cursor.execute(
                'PREPARE AwardBadge (bigint, text, integer) AS '
                    'WITH newrows AS ('
                        'SELECT $1, $2, $3, %s WHERE NOT EXISTS ('
                            'SELECT * FROM social.badge '
                                'WHERE uid = $1 AND badge_id = $2 AND '
                                    '(level IS NULL OR $3 IS NULL OR level >= $3))'
                    ') INSERT INTO social.badge '
                        '(uid, badge_id, level, awarded_by) '
                        'SELECT * FROM newrows', (ROBOT_UID,))

        print >>sys.stderr, 'Awarding old edits badges...',
        with closing(mpro_conn.cursor()) as cursor:
            with closing(open(cspath, 'r')) as csfile:
                for r in csfile:
                    uid, old_edits, _ = map(long, r.split())
                    level = edits_level(old_edits)
                    if level and uid in existing_uids:
                        cursor.execute(
                            'EXECUTE AwardBadge (%s, %s, %s)',
                            (uid, 'old_edits', level))
        mpro_conn.commit() # fix badges order
        print >>sys.stderr, 'done'

        print >>sys.stderr, 'Awarding active years badges...',
        with closing(mpro_conn.cursor()) as write_cursor:
            with closing(mpro_conn.cursor()) as read_cursor:
                read_cursor.execute(
                    'SELECT uid, EXTRACT(YEAR FROM created) '
                        'FROM acl.user')
                for r in read_cursor:
                    uid, reg_year = map(long, r)
                    level = min(5, 2015 - reg_year)
                    if level:
                        write_cursor.execute(
                            'EXECUTE AwardBadge (%s, %s, %s)',
                            (uid, 'active_years', level))
        mpro_conn.commit() # fix badges order
        print >>sys.stderr, 'done'

        print >>sys.stderr, 'Awarding edits badges...',
        with closing(mpro_conn.cursor()) as write_cursor:
            with closing(mpro_conn.cursor()) as read_cursor:
                read_cursor.execute(
                    'SELECT uid, total_edits '
                        'FROM social.stats')
                for r in read_cursor:
                    uid, edits = map(long, r)
                    level = edits_level(edits)
                    if level:
                        write_cursor.execute(
                            'EXECUTE AwardBadge (%s, %s, %s)',
                            (uid, 'edits', level))
        mpro_conn.commit() # fix badges order
        print >>sys.stderr, 'done'


@opster.command()
def ratings(wiki_conn_str=('', '', 'wiki9 connection string'),
            mpro_conn_str=('', '', 'mpro connection string')):
    """Convert rating positions (wiki9)"""
    with closing(psycopg2.connect(wiki_conn_str)) as wiki_conn:
        print >>sys.stderr, "Fetching 'created' positions...",
        created_pos_by_uid = {}
        with closing(wiki_conn.cursor()) as wiki_cursor:
            # carefully carved out of wiki-social
            wiki_cursor.execute(
                'SELECT social.user_stats.uid, social.user_stats.objects_created '
                    'FROM social.user_stats, social.user_settings '
                    'WHERE social.user_stats.uid = social.user_settings.uid '
                        'AND social.user_settings.banned = 0 '
                        'AND social.user_settings.uid NOT IN (96000691, 96000877, 108600248) '
                        'ORDER BY social.user_stats.objects_created DESC')

            cur_val = None
            cur_pos = 1
            for pos, r in enumerate(wiki_cursor, 1):
                if r[1] != cur_val:
                    cur_val = r[1]
                    cur_pos = pos
                created_pos_by_uid[long(r[0])] = cur_pos
        print >>sys.stderr, 'done'

        print >>sys.stderr, "Fetching 'edits' positions...",
        edits_pos_by_uid = {}
        with closing(wiki_conn.cursor()) as wiki_cursor:
            # carefully carved out of wiki-social
            wiki_cursor.execute(
                'SELECT social.user_stats.uid, social.user_stats.objects_modified '
                    'FROM social.user_stats, social.user_settings '
                    'WHERE social.user_stats.uid = social.user_settings.uid '
                        'AND social.user_settings.banned = 0 '
                        'AND social.user_settings.uid NOT IN (96000691, 96000877, 108600248) '
                        'ORDER BY social.user_stats.objects_modified DESC')

            cur_val = None
            cur_pos = 1
            for pos, r in enumerate(wiki_cursor, 1):
                if r[1] != cur_val:
                    cur_val = r[1]
                    cur_pos = pos
                edits_pos_by_uid[long(r[0])] = cur_pos
        print >>sys.stderr, 'done'

    print >>sys.stderr, 'Updating stats...',
    with closing(psycopg2.connect(mpro_conn_str)) as mpro_conn:
        existing_uids = fetch_existing_uids(mpro_conn)

        with closing(mpro_conn.cursor()) as mpro_cursor:
            mpro_cursor.execute('SELECT uid, old_rating_pos_created, old_rating_pos_edits FROM social.stats')
            existing_stats_uids = dict((long(r[0]), (r[1], r[2])) for r in mpro_cursor)

        with closing(mpro_conn.cursor()) as mpro_cursor:
            for uid_chunk in chunked(created_pos_by_uid.keys(), 100):
                update_vals = []
                create_vals = []

                for uid in uid_chunk:
                    if uid not in existing_uids:
                        continue

                    created_pos = created_pos_by_uid[uid]
                    edits_pos = edits_pos_by_uid[uid]
                    if uid in existing_stats_uids:
                        if existing_stats_uids[uid] != (created_pos, edits_pos):
                            update_vals.append('(%s, %s, %s)' % (uid, created_pos, edits_pos))
                    else:
                        create_vals.append('(%s, %s, %s)' % (uid, created_pos, edits_pos))

                if update_vals:
                    mpro_cursor.execute(
                        'UPDATE social.stats st '
                            'SET old_rating_pos_created = v.created_pos, '
                                'old_rating_pos_edits = v.edits_pos '
                            'FROM (VALUES %s) AS v (uid, created_pos, edits_pos) '
                            'WHERE st.uid = v.uid' % ','.join(update_vals))
                if create_vals:
                    mpro_cursor.execute(
                        'INSERT INTO social.stats '
                            '(uid, old_rating_pos_created, old_rating_pos_edits) '
                            'VALUES %s' % ','.join(create_vals))

        mpro_conn.commit()

    print >>sys.stderr, 'done'


if __name__ == '__main__':
    exit(opster.dispatch())
