import argparse
import logging
from yandex.maps.wiki import config, db, release_stages

ADVISORY_LOCK_ID = 0x1000000000000020

STAT_QUERY = \
"""
INSERT INTO mapspro.release_metrics (
        task_id, started_at, start_date,
        branch_id, task_name, last_task_duration
    ) VALUES (
        {task_id}, '{started_at}', '{start_date}',
        {branch_id}, '{task_name}', {last_task_duration}
    )
    ON CONFLICT (task_name, start_date) DO
    UPDATE SET
        task_id=EXCLUDED.task_id,
        started_at=EXCLUDED.started_at,
        branch_id=EXCLUDED.branch_id,
        last_task_duration=EXCLUDED.last_task_duration
"""


def stages_to_metrics(stages):
    metrics = []
    for s in stages:
        if not s.is_completed:
            continue

        metrics.extend([
            (s.name,
             s.task_ids[-1],
             s.last_task_start,
             int((s.end - s.last_task_start).total_seconds()))])

    return metrics


def get_report_data(session, limit):
    data = {}
    branch_ids = release_stages.get_latest_branch_ids(session, limit=limit)

    for branch_id in branch_ids:
        if branch_id <= 1:
            continue

        stages = release_stages.get_all_stages(session, branch_id)
        for [task_name, task_id, started_at, value] in stages_to_metrics(stages):
            start_date = str(started_at)[0:10]
            key = start_date + ":" + task_name
            if key not in data or data[key][0] < value:
                data[key] = [value, branch_id, task_id, started_at]

    report_data = []
    for key, value in data.items():
        report_data.append({
            'start_date': key[0:10],
            'task_name': key[11:],
            'last_task_duration': value[0],
            'branch_id': value[1],
            'task_id': str(value[2]),
            'started_at': str(value[3]),
        })
    return report_data


def main():
    arg_parser = argparse.ArgumentParser(
        description='Calculate release metrics and send them to Stat')
    arg_parser.add_argument(
        '--config', type=str,
        help='full path to services.xml')
    arg_parser.add_argument(
        '--limit', type=int, default=2,
        help='limit of processing branches (default: 2)')
    arg_parser.add_argument(
        '--dry-run', action='store_true',
        help='print metrics instead of sending them to statface')
    args = arg_parser.parse_args()

    logger = logging.getLogger('release_metrics')
    logger.addHandler(logging.StreamHandler())
    logger.setLevel(logging.INFO)

    config.init_config(args.config)
    db.init_pool(['core', 'stat'], 'tasks')

    with db.get_write_session('stat') as stat_session:
        if (stat_session.execute(
            'SELECT (pg_is_in_recovery()=false and pg_try_advisory_xact_lock(' +
                str(ADVISORY_LOCK_ID) + '))::int').scalar() == 0):
            print('Database already locked')
            return

        report_data = []
        with db.get_read_session('core') as core_session:
            report_data = get_report_data(core_session, args.limit)

        if args.dry_run:
            for i in report_data:
                print(i)
        elif len(report_data) == 0:
            print('Empty data to upload')
        else:
            print('Uploading data')
            for task in report_data:
                stat_session.execute(STAT_QUERY.format(**task))
            print('Done')


if __name__ == '__main__':
    main()
