#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


"""
Reads the "expanded" version of .circleci/config.yml from an arbitrary
build of each commit from the "master" branch of the pytorch repo and
inserts workflow jobs and properties into the Dr. CI database.
"""

import psycopg2
import psycopg2.extras
import yaml
import json
import os
import requests
import subprocess
import hashlib
import argparse

import logan_db_config


CONFIG_YAML_SHA1_API_URL_TEMPLATE = "https://api.github.com/repos/pytorch/pytorch/contents/.circleci/config.yml?ref=%s"


def populate_config_yaml_lambda_handler(event, context):

    payload = run(25)

    return {
        "statusCode": 200,
        "body": json.dumps({
            "message": "hello world",
            "payload": payload,
        }),
    }


def get_config_yaml_content_sha1(local_repo_path, commit_sha1):
    """
    Uses auth token to avoid GitHub API rate limiting
    """

    if local_repo_path:
        try:
            file_hash = subprocess.check_output([
                "git",
                "--git-dir",
                local_repo_path,
                "rev-parse",
                "%s:.circleci/config.yml" % commit_sha1,
            ]).decode('utf-8').strip()

            return file_hash

        except Exception as e:
            print("Couldn't obtain file SHA1 from local repo:", str(e))

    print("\tFetching content SHA1 from GitHub...")

    repo_sha1_retrieval_url = CONFIG_YAML_SHA1_API_URL_TEMPLATE % commit_sha1

    headers = {
        'Authorization': 'token %s' % logan_db_config.repo_read_auth_token,
    }

    repo_request = requests.get(repo_sha1_retrieval_url, headers=headers)
    github_api_response_json = repo_request.json()

    sha1 = github_api_response_json.get("sha")

    if not sha1:
        print("PROBLEM:", github_api_response_json)

    return sha1


def populate_db_yaml_records(cur, build_number, repo_yaml_content_sha1):

    url = "https://circleci.com/api/v1.1/project/github/pytorch/pytorch/%d" % build_number

    r = requests.get(url)
    api_json_obj = r.json()

    yaml_text = api_json_obj.get("circle_yml").get("string")

    expanded_yaml_md5 = hashlib.md5(yaml_text.encode('utf-8')).hexdigest()

    yaml_obj = yaml.safe_load(yaml_text)
    workflows_dict = yaml_obj.get("workflows")

    cur.execute(
        'INSERT INTO circleci_config_yaml_hashes (expanded_yaml_content, expanded_yaml_md5, repo_yaml_sha1) VALUES (%s, %s, %s);',
        (yaml_text, expanded_yaml_md5, repo_yaml_content_sha1)
    )


    branch_filters_by_job_by_workflow = {}

    jobs_insertion_values = []
    # (workflow, dependent_job, required_job)
    job_dependency_tuples = []
    schedule_insertion_values = []

    for workflow_name, workflow_obj in filter(lambda x: x[0] != "version", workflows_dict.items()):

        if type(workflow_obj) is dict:


            cur.execute(
                'INSERT INTO circleci_workflows_by_yaml_file (yaml_content_sha1, name) VALUES (%s, %s) RETURNING id;',
                (repo_yaml_content_sha1, workflow_name)
            )
            workflow_id = cur.fetchone()[0]

            cron_values = []
            for trigger in workflow_obj.get("triggers", []):
                schedule_obj = trigger.get("schedule", {})

                for k, v in schedule_obj.items():
                    if k == "cron":
                        cron_values.append(v)

            for v in cron_values:
                schedule_insertion_values.append((workflow_id, v))


            branch_filters_by_job = branch_filters_by_job_by_workflow.setdefault(workflow_id, {})

            for job_obj in workflow_obj.get("jobs", []):
                if type(job_obj) is dict:
                    job_name = list(job_obj.keys())[0]

                    for key_job_name, job_value_obj in job_obj.items():
                        branch_filter_only_obj = job_value_obj.get("filters", {}).get("branches", {}).get("only")
                        if type(branch_filter_only_obj) is list:
                            branch_filters_by_job.setdefault(job_name, []).extend(branch_filter_only_obj)
                        elif type(branch_filter_only_obj) is str:
                            branch_filters_by_job.setdefault(job_name, []).append(branch_filter_only_obj)

                        for required_job in job_value_obj.get("requires", []):
                            job_dependency_tuples.append((workflow_id, key_job_name, required_job))

                else:
                    job_name = job_obj

                jobs_insertion_values.append((workflow_id, job_name))

    jobs_insert_query = 'INSERT INTO circleci_workflow_jobs (workflow, job_name) VALUES %s'
    psycopg2.extras.execute_values(
        cur, jobs_insert_query, jobs_insertion_values, template=None, page_size=100
    )

    dependencies_insert_query = 'INSERT INTO circleci_config_job_dependencies (workflow, dependent_job, required_job) VALUES %s'
    psycopg2.extras.execute_values(
        cur, dependencies_insert_query, job_dependency_tuples, template=None, page_size=100
    )

    filter_insertion_values = []
    for workflow_id, branch_filters_by_job in branch_filters_by_job_by_workflow.items():
        for job_name, filters_list in branch_filters_by_job.items():
            filter_insertion_values.extend([(workflow_id, job_name, branch, True) for branch in filters_list])

    insert_query2 = 'INSERT INTO circleci_job_branch_filters (workflow, job_name, branch, filter_include) VALUES %s'
    psycopg2.extras.execute_values(
        cur, insert_query2, filter_insertion_values, template=None, page_size=100
    )

    schedule_insert_query = 'INSERT INTO circleci_workflow_schedules (workflow, cron_schedule) VALUES %s'
    psycopg2.extras.execute_values(
        cur, schedule_insert_query, schedule_insertion_values, template=None, page_size=100
    )


def populate_config_info(local_repo_path, cur, commit_sha1, build_number):

    repo_yaml_content_sha1 = get_config_yaml_content_sha1(local_repo_path, commit_sha1)
    if repo_yaml_content_sha1:

        cur.execute(
            "SELECT repo_yaml_sha1 FROM circleci_config_yaml_hashes WHERE repo_yaml_sha1=%s LIMIT 1;",
            (repo_yaml_content_sha1,)
        )
        row = cur.fetchone()

        if not row:
            print("\tInserting workflow into database...")
            populate_db_yaml_records(cur, build_number, repo_yaml_content_sha1)
        else:
            print("\tWorkflow is already in database.")

        print("\tInserting git-commit/config.yaml association into database...")
        cur.execute(
            'INSERT INTO circleci_expanded_config_yaml_hashes_by_commit (commit_sha1, repo_yaml_sha1) VALUES (%s, %s);',
            (commit_sha1, repo_yaml_content_sha1)
        )
        print("\tInserted git-commit/config.yaml association into database.")

    else:
        print("Couldn't retrieve file content sha1 for commit %s!" % commit_sha1)


def run(commit_count, local_repo_path=None):

    if local_repo_path:
        return_code = subprocess.call([
            "git",
            "--git-dir",
            local_repo_path,
            "fetch",
            "origin",
            "master",
        ])

        print("Fetched local git repo with return code: %d" % return_code)

    conn = psycopg2.connect(
        host=logan_db_config.db_hostname,
        database=logan_db_config.db_name,
        user=logan_db_config.db_username,
        password=logan_db_config.db_password)

    with conn.cursor() as cur:

        cur.execute('SET SESSION lock_timeout = 3000;')  # 3 seconds
        cur.execute('SET SESSION statement_timeout = %d;' % (1000*60*3))  # 3 minutes

        cur.execute("SELECT sha1, build_num FROM master_commits_unpopulated_circleci_configs LIMIT %s;", (commit_count,))
        rows = cur.fetchall()

        enumerated_rows = list(enumerate(rows))

        def single_commit_populator(args_tuple):
            (i, (commit_sha1, build_number)) = args_tuple
            print("%d/%d: Populating CircleCI config for commit %s using build #%d..." % (i + 1, len(enumerated_rows), commit_sha1, build_number))

            populate_config_info(local_repo_path, cur, commit_sha1, build_number)

        # We don't allow concurrent actions here, since we don't want two Git commits
        # with the same config.yml hash to race in database insertion.
        #
        # We commit the transaction after every row, so that we can make incremental progress
        # even if the overall task fails.
        for x in enumerated_rows:
            single_commit_populator(x)
            conn.commit()

    return {
        "foo": "bar",
    }


def parse_args():
    parser = argparse.ArgumentParser(description='Parse config.yml files for revisions of the pytorch repo')
    parser.add_argument('--repo-path', dest='local_repo_path',
                        default=os.path.expanduser("~/github/pytorch-repos/pytorch/.git"),
                        help='Local filesystem path to pytorch repo .git directory')

    parser.add_argument('--commit-count', dest='commit_count',
                        type=int,
                        default=2,
                        help='How many commits to retrieve')

    return parser.parse_args()


if __name__ == "__main__":

    options = parse_args()

    payload = run(options.commit_count)
    print(payload)
