#!/usr/bin/env python3
# Copyright (c) Aptos
# SPDX-License-Identifier: Apache-2.0

import argparse
import getpass
import json
import os
import subprocess
import sys
import tempfile
import time
from datetime import datetime, timezone

from kube import (
    kube_init_context,
    kube_select_cluster,
    kube_ensure_cluster,
    get_cluster_context,
    kube_wait_job,
    create_forge_job,
    push_helm_charts,
)

TAG = ""
BASE_TAG = ""
WORKSPACE = ""
REPORT = ""
DEFL_TIMEOUT_SECS = 1800  # Default timeout is 30 mins
USER = getpass.getuser()  # Use the current user for naming
FORGE_OUTPUT_TEE = os.getenv("FGI_OUTPUT_LOG", tempfile.mkstemp()[1])

OPENSEARCH_TIME_FMT = '%Y-%m-%dT%H:%M:%S.000Z'
DEFL_FORGE_REPORT = "forge_report.json"

HEADER = "\033[95m"
OKBLUE = "\033[94m"
OKGREEN = "\033[92m"
WARNING = "\033[93m"
FAIL = "\033[91m"
RESTORE = "\033[0m"


def build_argparser():
    """Build the arg parser and return a tuple of (fgi args, forge args)"""
    parser = argparse.ArgumentParser(
        description="Entrypoint for the Forge unified testing framework"
    )
    parser.add_argument(
        "--timeout-secs",
        default=DEFL_TIMEOUT_SECS,
        help="Timeout for the Forge pod in seconds",
    )
    parser.add_argument(
        "--workspace",
        "-W",
        help="Workspace or kubernetes cluster to run Forge on, rather than picking one at random",
    )
    parser.add_argument(
        "--env",
        "-E",
        action="append",
        default=[],
        help="Extra environment variables to pass to Forge",
    )
    build_group = parser.add_mutually_exclusive_group(required=True)
    build_group.add_argument(
        "--tag", "-T", help="Image tag to use in kubernetes Forge tests")
    build_group.add_argument(
        "--pr", "-p", help="PR to build images from for kubernetes Forge tests")
    build_group.add_argument(
        "--local-swarm",
        "-L",
        action="store_true",
        help="Run Forge tests locally instead of on a kubernetes cluster",
    )
    parser.add_argument(
        "--base-image-tag",
        "-B",
        help="Base image tag to use in kubernetes Forge test"
    )
    parser.add_argument(
        "--report",
        "-R",
        default=DEFL_FORGE_REPORT,
        help="Dump report to file"
    )
    return parser.parse_known_args()


def get_grafana_url(cluster_name):
    return f"http://o11y.aptosdev.com/grafana/d/overview/overview?orgId=1&var-Datasource=Remote%20Prometheus%20Devinfra&var-chain_name=aptos-{cluster_name}&var-owner=All"


def cli_tool_installed(tool_name):
    ret = subprocess.run(
        ["which", tool_name], stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL
    )
    return ret.returncode == 0


def get_opensearch_time(timestamp_ms):
    return datetime.fromtimestamp(timestamp_ms/1000, timezone.utc).strftime(OPENSEARCH_TIME_FMT)


def print_streaming_o11y_resources(grafana_url, start_ts_ms):
    print("\n**********")
    print(
        f"{OKBLUE}Auto refresh Dashboard:{RESTORE} {grafana_url}&from={start_ts_ms}&to=now&refresh=10s")
    print("**********")


def get_final_o11y_resources(grafana_url, workspace, start_ts_ms, end_ts_ms):
    DASHBOARD_LINK = f"{grafana_url}&from={start_ts_ms}&to={end_ts_ms}"
    LOGGING_LINK = (
        f"https://es.devinfra.aptosdev.com/_dashboards/app/discover#/?"
        f"_g=(filters:!(),refreshInterval:(pause:!t,value:0),time:(from:'{get_opensearch_time(start_ts_ms)}',to:'{get_opensearch_time(end_ts_ms)}'))"
        f"&_a=(columns:!(_source),filters:!(('$state':(store:appState),"
        f"meta:(alias:!n,disabled:!f,index:d0bc5e20-badc-11ec-9a50-89b84ac337af,key:chain_name,negate:!f,params:(query:aptos-{workspace}),type:phrase),"
        f"query:(match_phrase:(chain_name:aptos-{workspace})))),"
        f"index:d0bc5e20-badc-11ec-9a50-89b84ac337af,interval:auto,query:(language:kuery,query:''),sort:!())"
    )

    print("\n**********")
    print(f"{OKBLUE}Dashboard snapshot:{RESTORE} {DASHBOARD_LINK}")
    print(f"{OKBLUE}Logs snapshot:{RESTORE} {LOGGING_LINK}")
    print("**********\n")
    return DASHBOARD_LINK, LOGGING_LINK


# ================ Parse the args ================
args, forge_args = build_argparser()

# build and push the images to be used, since an image tag
# was not specified explicitly
TAG = args.tag
BASE_TAG = args.base_image_tag
REPORT = args.report
if not args.tag:
    if args.pr:  # codebuild using a PR
        ret = subprocess.call(
            ["aws", "codebuild", "list-projects"], stdout=subprocess.DEVNULL
        )
        if ret != 0:
            print(f"{FAIL}Failed to access codebuild. Try aws-mfa?{RESTORE}")
            sys.exit(1)
        ret = subprocess.call(
            ["./docker/build-aws.sh", "--build-forge",
                "--version", f"pull/{args.pr}"]
        )
        if ret != 0:
            print(f"{FAIL}Failed to build forge.")
            sys.exit(1)
        TAG = f"dev_{USER}_pull_{args.pr}"
        print(
            f"**TIP Use ./scripts/fgi/run -T {TAG} <...> to restart this run with the same tag without rebuilding it"
        )
if not args.base_image_tag:
    BASE_TAG = "devnet"
# ================ Test setup ================
print(f"""
    {HEADER}______{OKBLUE}____  {OKGREEN}____  {WARNING}______{FAIL}______
   {HEADER}/ ____{OKBLUE}/ __ \{OKGREEN}/ __ \{WARNING}/ ____{FAIL}/ ____/
  {HEADER}/ /_  {OKBLUE}/ / / {OKGREEN}/ /_/ {WARNING}/ / __{FAIL}/ __/
 {HEADER}/ __/ {OKBLUE}/ /_/ {OKGREEN}/ _, _{WARNING}/ /_/ {FAIL}/ /___
{HEADER}/_/    {OKBLUE}\____{OKGREEN}/_/ |_|{WARNING}\____{FAIL}/_____/
{RESTORE}
""")

if args.local_swarm:
    print("Running Forge on backend: local swarm")
    ret = subprocess.call(
        ["cargo", "run", "-p", "forge-cli", "--", "test", "local-swarm"])
    sys.exit(ret)

print("Running Forge on backend: kubernetes testnet")

if not cli_tool_installed("kubectl"):
    print(
        f"{WARNING}kubectl is not installed. Please install kubectl. On mac, you can use: brew install kubectl{RESTORE}"
    )
    print(
        f"{WARNING}or install via dev setup: scripts/dev_setup.sh -i kubectl{RESTORE}"
    )
    sys.exit(1)

print("\nAttempting to reach Forge Kubernetes testnets...")
kube_init_context(args.workspace)
print("Grabbing a testnet...")
workspace = args.workspace
if not args.workspace:
    workspace = kube_select_cluster()
    if not workspace:
        print(f"{FAIL}Failed to select forge testnet cluster{RESTORE}")
        sys.exit(1)
else:
    ret = kube_ensure_cluster([workspace])
    if not ret:
        print(
            f"{FAIL}Failed to acquire specified forge testnet cluster {workspace}{RESTORE}")
        sys.exit(1)
context = get_cluster_context(workspace)
print(f"Running experiments on cluster: {workspace}")
push_helm_charts(workspace)
print()

job_name, template = create_forge_job(
    context, USER, TAG, BASE_TAG, args.timeout_secs, args.env, forge_args
)
if not template:
    print(f"{FAIL}Failed to create forge job template{RESTORE}")
    sys.exit(1)

_, specfile = tempfile.mkstemp(suffix=".json")
with open(specfile, "w") as f:
    f.write(json.dumps(template))
print(f"Specfile: {specfile}")

# ================ Create and run the job ================
print(f"Creating job: {job_name}")
ret = subprocess.call(
    ["kubectl", f"--context={context}", "apply", "-f", specfile])
if ret != 0:
    print(f"{FAIL}Failed to create forge job{RESTORE}")
    sys.exit(1)

ret = kube_wait_job(job_name, context)
if ret != 0:
    print(f"{FAIL}Failed to start forge job{RESTORE}")
    sys.exit(1)

# account for the time delta between querying pod status and finishing waiting
delta_ms = 1000
start_ts_ms = int(time.time() * 1000) - delta_ms
grafana_url = get_grafana_url(workspace)

print_streaming_o11y_resources(grafana_url, start_ts_ms)

print("==========begin-forge-pod-logs==========")
# TODO(rustielin): might have to retry this if kube reschedules it
subprocess.call(
    f"kubectl --context={context} logs -f -l job-name={job_name} | tee -a {FORGE_OUTPUT_TEE}",
    shell=True,
)
print("==========end-forge-pod-logs==========")
print(f"\nForge log output: {FORGE_OUTPUT_TEE}")

try:
    job_status = json.loads(
        subprocess.check_output(
            [
                "kubectl",
                f"--context={context}",
                "get",
                "job",
                job_name,
                "-o",
                "json",
            ],
            encoding="UTF-8",
        )
    )["status"]
except Exception as e:
    print(f"Failed to get job status for {job_name}, assuming failure: {e}")
    job_status = {"failed": 1}

end_ts_ms = int(time.time() * 1000)

DASHBOARD_LINK, LOGGING_LINK = get_final_o11y_resources(
    grafana_url, workspace, start_ts_ms, end_ts_ms)

test_res = 'failed'
# perf report
test_report = []
read_lines = False
read_error_caused_by = False
error_caused_by = "Forge test failure"
with open(f"{FORGE_OUTPUT_TEE}", 'r') as file:
    for line in file.readlines():
        # get the json report
        if 'json-report-end' in line:
            read_lines = False
        if read_lines:
            test_report.append(line)
        if 'json-report-begin' in line:
            read_lines = True
        if 'test result: ok' in line:
            test_res = 'passed'

        # attempt to get the error
        if read_error_caused_by:
            error_caused_by = line
            read_error_caused_by = False
        if 'Caused by:' in line:
            read_error_caused_by = True
if len(test_report) == 0:
    # If Forge emits no report (during test setup), return a generic error
    test_report.append("{\"text\": \"Forge test runner is terminated\"}")

temp_forge_report = json.loads(''.join(test_report))

# If Forge emits a report but no text, attempt to return the failure
if not temp_forge_report["text"]:
    temp_forge_report["text"] = error_caused_by

temp_forge_report["logs"] = LOGGING_LINK
temp_forge_report["dashboard"] = DASHBOARD_LINK
if args.report:
    with open(f"{REPORT}", 'w') as file_object:
        file_object.write(json.dumps(temp_forge_report))


if "failed" in job_status and job_status["failed"] == 1:
    print()
    print(f"{FAIL}Job {job_name} failed{RESTORE}")
else:
    print(f"{OKGREEN}Job {job_name} succeeded!{RESTORE}")


if test_res == 'failed':
    exit(1)
