#!/usr/bin/env python
# Copyright 2014-2020 Scalyr Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Script which prints Circle CI usage data for the recent workflows.

It relies on the following environment variables being set:

- CIRCLE_CI_API_TOKEN
- MAILGUN_API_TOKEN (only when reports are sent to email)
"""

from __future__ import absolute_import
from __future__ import print_function
from scalyr_agent import compat

if False:  # NOSONAR
    from typing import Optional
    from typing import List

import sys
import argparse
import datetime
import json

from io import StringIO

import requests

# See https://circleci.com/pricing/#comparison-table
CREDIT_BUNDLE_PRICE = 15.0
CREDIT_BUNDLE_UNIT_COUNT = 25000.0

# Price for a single CPU usage credit
PRICE_PER_CREDIT = CREDIT_BUNDLE_PRICE / CREDIT_BUNDLE_UNIT_COUNT

GITHUB_API_PRS_URL = "https://api.github.com/repos/{project_slug}/pulls"
GITHUB_API_SEARCH_URL = "https://api.github.com/search/issues"
CIRCLE_CI_API_INSIGHTS_URL = (
    "https://circleci.com/api/v2/insights/{project_slug}/workflows/{workflow}"
)
CIRCLE_CI_WORKFLOW_VIEW_URL = "https://circleci.com/workflow-run/{workflow_id}"

CIRCLE_CI_API_TOKEN = compat.os_environ_unicode.get("CIRCLE_CI_API_TOKEN", None)
MAILGUN_API_TOKEN = compat.os_environ_unicode.get("MAILGUN_API_TOKEN", None)


def get_all_branches_for_repo(project_slug, limit=5):
    # type: (str, int) -> List[str]
    """
    Retrieve all the branches (including deleted ones) for a particular repo.

    NOTE: This function utilizes pulls API endpoint since /branches one doesn't include deleted
    branches.
    """
    # 1. Retrieve all the PRs for last week
    url = GITHUB_API_SEARCH_URL.format(project_slug=project_slug)

    now = datetime.datetime.utcnow().strftime("%Y-%m-%d")
    week_ago = (datetime.datetime.utcnow() - datetime.timedelta(days=7)).strftime(
        "%Y-%m-%d"
    )

    query = "repo:{project_slug} is:pr created:{week_ago}..{now}".format(
        project_slug=project_slug, week_ago=week_ago, now=now
    )
    params = {"q": query, "per_page": "1000", "sort": "created"}

    response = requests.get(url, params=params)
    data = response.json()
    pull_numbers = [item["number"] for item in data["items"]]

    # 2. Retrieve all the branches and filtered ones which haven't been updated during last week
    # (based on the PR activity)
    # 1. Retrieve all the branches
    url = GITHUB_API_PRS_URL.format(project_slug=project_slug)
    params = {"per_page": "100", "sort": "created", "direction": "desc", "state": "all"}
    response = requests.get(url, params=params)
    data = response.json()

    branches = []
    for item in data:
        branch_ref = item.get("head", {}).get("ref", None)
        pull_number = item["number"]

        if not branch_ref or pull_number not in pull_numbers:
            continue

        branches.append(branch_ref)

    if limit:
        return branches[:limit]

    return branches


def get_usage_data_for_branch_and_workflow(
    buff, project_slug, workflow, status="success", branch="master", limit=10
):
    # type: (StringIO, str, str, str, str, int) -> None
    """
    Get usage data for the provided project and write it to the provided buffer.

    :param status: Only return data for workflow runs with the provided status. "all" for all the
                   available statuses.
    :param branch: Branch to return data work.
    :param limit: Maximum number of workflow runs we want to see for a particular branch. NOTE: This
                  is different from the API limit parameter because API returns data for all the
                  workflow run and not only the ones we are interested in and perform late filtering
                  on.
    """
    url = CIRCLE_CI_API_INSIGHTS_URL.format(
        project_slug=project_slug, workflow=workflow
    )
    params = {"limit": "100", "branch": branch}
    headers = {"Circle-Token": CIRCLE_CI_API_TOKEN}

    response = requests.get(url, params=params, headers=headers)

    if response.status_code == 404:
        raise ValueError("Invalid CIRCLE_CI_API_TOKEN or project_slug")

    try:
        data = response.json()
    except json.decoder.JSONDecodeError as e:
        # This likely indicates issue with Circle CI returning invalid response
        print(
            "Failed to parse Circle CI response as json: %s. Got response: %s"
            % (str(e), str(response.text))
        )
        sys.exit(1)

    items = data.get("items", [])
    items = [item for item in items if (status == "all" or item["status"] == status)]

    buff.write(
        'Usage data for recent workflow runs for workflow "%s" with status="%s" and branch "%s"\n\n'
        % (workflow, status, branch)
    )

    if not items:
        buff.write("No recent runs which match this criteria.\n\n")
        buff.write("=" * 100 + "\n\n")
        return

    count = 0
    for item in items:
        workflow_view_url = CIRCLE_CI_WORKFLOW_VIEW_URL.format(workflow_id=item["id"])
        price_estimate = round(item["credits_used"] * PRICE_PER_CREDIT, 2)

        buff.write("Workflow ID: %s (%s)\n" % (item["id"], workflow_view_url))
        buff.write("Start time: %s\n" % (item["created_at"]))
        buff.write("Duration: %s seconds\n" % item["duration"])
        buff.write("Status: %s\n" % (item["status"]))
        buff.write("Total duration: %s seconds\n" % item["duration"])
        buff.write(
            "Credits used: %s (~%s USD)\n\n" % (item["credits_used"], price_estimate)
        )

        count += 1
        if count >= limit:
            break

    buff.write("=" * 100 + "\n\n")


def print_usage_data(
    project_slug,
    workflows,
    status="success",
    branch="master",
    limit=5,
    branch_limit=5,
    emails=None,
):
    # type: (str, List[str], str, str, int, int, Optional[List[str]]) -> None
    if branch == "all":
        branches = get_all_branches_for_repo(
            project_slug=project_slug.replace("gh/", ""), limit=branch_limit
        )
        # We always include master branch
        branches.insert(0, "master")
    else:
        branches = [branch]

    buff = StringIO()

    buff.write("Circle CI usage and cost report\n\n")
    buff.write(
        "Pricing estimate assumes %s credits cost %s$\n\n"
        % (CREDIT_BUNDLE_UNIT_COUNT, CREDIT_BUNDLE_PRICE)
    )

    for branch in branches:
        for workflow in workflows:
            get_usage_data_for_branch_and_workflow(
                buff=buff,
                project_slug=project_slug,
                workflow=workflow,
                status=status,
                branch=branch,
                limit=limit,
            )

    value = buff.getvalue()

    if emails:
        now = datetime.datetime.utcnow().strftime("%Y-%m-%d")
        week_ago = (datetime.datetime.utcnow() - datetime.timedelta(days=7)).strftime(
            "%Y-%m-%d"
        )

        subject = "Circle CI Usage Report For Period %s - %s" % (week_ago, now)
        send_email(to=emails, subject=subject, text=value)
        print(("Sent email report to %s" % (",".join(emails))))

    print(value)


def send_email(to, subject, text):
    # type: (List[str], str, str) -> None
    response = requests.post(
        "https://api.mailgun.net/v3/kami.mailgun.org/messages",
        auth=("api", MAILGUN_API_TOKEN),
        data={
            "from": "Circle CI Usage Reporter <mailgun@kami.mailgun.org>",
            "to": to,
            "subject": subject,
            "text": text,
        },
    )

    if response.status_code == 401:
        raise ValueError("Invalid value for MAILGUN_API_TOKEN environment variable")
    elif response.status_code != 200:
        raise ValueError("Failed to send email: %s" % (response.text))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description=(
            "Retrieve and print / email Circle CI usage data "
            "for a particular project and workflow."
        )
    )
    parser.add_argument(
        "--project_slug",
        help=(
            "Project slug in the following format: <vcs>/<org name>/<repo name> "
            "(e.g. gh/scalyr/scalyr-agent-2)"
        ),
        default="gh/scalyr/scalyr-agent-2",
    )
    parser.add_argument(
        "--workflows",
        help="Name of the workflows to print the usage data for (e.g. unittest,benchmarks,e2e).",
        default="unittest",
    )
    parser.add_argument(
        "--status",
        help='Workflow status to filter on. "all" for all available statuses.',
        default="success",
    )
    parser.add_argument(
        "--branch",
        help='Branch to print the usage data for. "all" for all the active branches.',
        default="master",
    )
    parser.add_argument(
        "--limit",
        help="Maximum number of workflow runs per branch to print data for.",
        type=int,
        default=5,
    )
    parser.add_argument(
        "--branch-limit",
        help="Maximum number of most recently active branches to print data for.",
        type=int,
        default=5,
    )

    parser.add_argument(
        "--emails",
        help="If provided, report will also be emailed to those addresses (comma separated list).",
        type=str,
        default=None,
        required=False,
    )

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

    if args.status not in ["all", "success", "failed"]:
        raise ValueError("Invalid status: %s" % (args.status))

    if not CIRCLE_CI_API_TOKEN:
        raise ValueError("CIRCLE_CI_API_TOKEN environment variable is not set")

    emails = args.emails.split(",") if args.emails else []
    if emails and not MAILGUN_API_TOKEN:
        raise ValueError("MAILGUN_API_TOKEN environment variable is not set")

    workflows = args.workflows.split(",")

    print_usage_data(
        project_slug=args.project_slug,
        workflows=workflows,
        status=args.status,
        branch=args.branch,
        limit=args.limit,
        branch_limit=args.branch_limit,
        emails=emails,
    )
