#!/usr/bin/env python3

import httpx
import argparse
import os
import getpass
import sys
from prettytable import PrettyTable

GREEN = "\033[92m"
RED = "\033[91m"
GRAY = "\033[90m"
YELLOW = "\033[93m"
RESET = "\033[0m"
BLUE = "\033[35m"

COLOR_MAP = {
    "success": GREEN,
    "fail": RED,
    "pending": GRAY,
    "-": GRAY,
    "running": BLUE,
    "request_confirmation": YELLOW
}

TOKEN_FILE = '/run/ngcp-upgrade-api/client_session.token'
UPGRADE_API_SOCKET = '/run/ngcp-upgrade-api/ngcp-upgrade-api.sock'

url = "http://localhost:8000"
headers = {'Content-Type': 'application/json'}


def get_api_client():
    return httpx.Client(transport=httpx.HTTPTransport(uds=UPGRADE_API_SOCKET))


def load_token():
    if os.path.exists(TOKEN_FILE):
        return open(TOKEN_FILE, "r").read().strip()
    return None


def save_token(token):
    with open(TOKEN_FILE, "w") as f:
        f.write(token)
    os.chmod(TOKEN_FILE, 0o600)


def login():
    if os.getenv('NGCP_UPGRADE_API_AUTO'):
        password = os.getenv('NGCP_UPGRADE_API_PASSWORD')
    else:
        print("No active session. Please enter root password:")
        password = getpass.getpass(prompt='Password: ')
    client = get_api_client()
    try:
        response = client.post(f"{url}/login", json={
            "username": "root",
            "password": password
        })
    except Exception as exc:
        sys.stderr.write(f"Can't connect to ngcp-upgrade-api: {exc}\n")
        sys.stderr.write("Check if ngcp-upgrade-api systemd service "
                         "is running\n")
        exit(1)
    try:
        data = response.json()
    except Exception as exc:
        sys.stderr.write(f"Can't decode answer from ngcp-upgrade-api: {exc}\n")
        exit(1)
    if 'token' in data:
        token = data['token']
        save_token(token)
        return token
    elif 'detail' in data:
        sys.stderr.write(f"{data['detail']}\n")
        exit(1)
    else:
        sys.stderr.write("Can't login")
        exit(1)


def run_api_method(method, payload):
    token = load_token()
    if not token:
        token = login()
    headers['X-Session-Token'] = token

    m_url = f"{url}/{method}"
    client = get_api_client()

    while True:
        try:
            response = client.post(
                m_url,
                headers=headers,
                json=payload)
        except Exception as exc:
            sys.stderr.write(f"Can't connect to ngcp-upgrade-api: {exc}\n")
            sys.stderr.write("Check if ngcp-upgrade-api systemd service "
                             "is running\n")
            exit(1)

        try:
            data = response.json()
        except Exception as exc:
            sys.stderr.write(f"Can't decode answer from "
                             f"ngcp-upgrade-api: {exc}\n")
            exit(1)

        status_code = response.status_code
        if status_code == 401:
            if data['detail'] == 'Session expired':
                token = login()
                headers['X-Session-Token'] = token
        if status_code == 200:
            break
        else:
            sys.stderr.write("Error: Can't query ngcp-upgrade-api, "
                             "exiting...\n")
            if 'detail' in data:
                sys.stderr.write(f"{data['detail']}\n")
            exit(1)

    return (data)


def create_scenario(release_from, release_to):
    payload = {'release_from': release_from, 'release_to': release_to}
    answer = run_api_method('create_scenario', payload)
    print(answer['scenario_name'])


def run_scenario(name):
    payload = {'scenario_name': name}
    if os.getenv('NGCP_UPGRADE_API_AUTO'):
        payload['force'] = 'true'
    answer = run_api_method('run_scenario', payload)
    print(answer['status'])


def get_scenario_status(name):
    payload = {'scenario_name': name}
    answer = run_api_method('get_scenario_status', payload)
    print(answer['status'])
    if os.getenv('NGCP_UPGRADE_API_AUTO'):
        return

    steps = run_api_method('get_step_status', payload)
    if answer['status'] == 'running':
        for step in steps:
            for key in step:
                step_name = key
            step_status = step[step_name]['status']
            if step_status == 'running':
                nodes = []
                for key in step[step_name]['nodes']:
                    nodes.append(key)
                print(f"Running step '{step_name}' on nodes: {nodes}")
                return
    elif answer['status'] == 'fail':
        for step in steps:
            for key in step:
                step_name = key
            step_status = step[step_name]['status']
            if step_status == 'fail':
                print(f"Step '{step_name}':")
                for key in step[step_name]['nodes']:
                    print(f"{key}: {step[step_name]['nodes'][key]['status']}")
                return
    elif answer['status'] == 'pending':
        found_progress = False
        for step in reversed(steps):
            for key in step:
                step_name = key
            step_status = step[step_name]['status']
            if step_status == 'success':
                found_progress = True
                break
        if step_name == 'report_check_stop':
            print("The check stage is complete. "
                  "Use 'run' command again to continue to stage 1")
        elif step_name == 'stop':
            print("Stage 1 is complete. Enable maintanance mode, "
                  "run /usr/share/ngcp-upgrade/ngcp-switch-root-partition, "
                  "and reboot the node.")
            print("After it - use 'run' command to proceed with stage 2")
        elif step_name == 'report_upgrade_stop':
            print("Stage 2 is complete on a node. Use 'run' command again to "
                  "proceed with upgrade of next node(s)")

        if not found_progress:
            print(f"It looks like scenario {name} wasn't run yet.")
            print(f"Use 'run {name}' command to proceed with upgrade")


def get_running_step(name):
    payload = {'scenario_name': name}
    answer = run_api_method('get_step_status', payload)

    running_step_found = False
    for step in answer:
        for key in step:
            step_name = key
        step_status = step[step_name]['status']
        if step_status == 'running':
            print(f"Running step '{step_name}'")
            running_step_found = True
            break
    if not running_step_found:
        print("Running step not found")


def stop_scenario():
    answer = run_api_method('stop', {})
    print(answer['detail'])


def show_matrix(name):
    payload = {'scenario_name': name}
    answer = run_api_method('get_step_status', payload)

    table = PrettyTable()
    reset = RESET

    first_step_name = list(answer[0].keys())[0]
    node_names = list(answer[0][first_step_name]['nodes'].keys())
    # Carrier special order
    if 'db01a' in node_names:
        ordered_list = ['web01a', 'db01a']
        for node in node_names:
            if node.endswith('a') and node not in ('web01a', 'db01a'):
                ordered_list.append(node)
        for node in node_names:
            if node.endswith('b'):
                ordered_list.append(node)
        node_names = ordered_list

    table.field_names = ["Step Name"] + node_names

    for step in answer:
        for key in step:
            step_name = key
        row = [step_name]
        for node in node_names:
            try:
                status = step[step_name]['nodes'][node]['status']
            except KeyError as e:
                status = '-'
            color = COLOR_MAP.get(status)
            if os.getenv('NGCP_UPGRADE_API_AUTO'):
                color = ''
                reset = ''
            row.append(f"{color}{status}{reset}")
        table.add_row(row)

    table.align = "l"
    print(f"\nScenario: {name}")
    print(table)


def get_scenarios_list():
    answer = run_api_method('get_scenarios_list', {})

    table = PrettyTable()
    table.field_names = ["Scenario", "Status"]

    reset = RESET
    for scenario in answer['scenarios']:
        for key in scenario:
            scenario_name = key
        status = scenario[scenario_name]
        color = COLOR_MAP.get(status)
        if os.getenv('NGCP_UPGRADE_API_AUTO'):
            color = ''
            reset = ''
        table.add_row([f"{color}{scenario_name}{reset}",
                       f"{color}{status}{reset}"])

    table.align = "l"
    print(table)


def delete_scenatio(name):
    payload = {'scenario_name': name}
    answer = run_api_method('delete_scenario', payload)
    print(answer['status'])


def get_log(scenario_name, step_name, node_name):
    payload = {'scenario_name': scenario_name,
               'step_name': step_name,
               'node_name': node_name}
    answer = run_api_method('get_step_log', payload)
    print(answer['detail'])


def get_failed_log(scenario_name):
    payload = {'scenario_name': scenario_name}
    answer = run_api_method('get_step_status', payload)
    failed_step_found = False
    for step in answer:
        step_name = list(step.keys())[0]
        if not step[step_name]['status'] == 'fail':
            continue
        print(f"Failed step: {step_name}")
        failed_step_found = True
        failed_nodes = []
        for node_name in step[step_name]['nodes']:
            if step[step_name]['nodes'][node_name]['status'] == 'fail':
                failed_nodes.append(node_name)
        for node_name in failed_nodes:
            payload = {'scenario_name': scenario_name,
                       'step_name': step_name,
                       'node_name': node_name}
            answer = run_api_method('get_step_log', payload)
            print(f"  Node: {node_name} output:")
            print(answer['detail'])

    if not failed_step_found:
        print(f"No failed step found in scenario {scenario_name}")


def main():
    if os.getuid() != 0:
        print("This script must be run as root.")
        exit(1)

    parser = argparse.ArgumentParser(description="CLI for ngcp-upgrade-api")
    subparsers = parser.add_subparsers(
        dest="command", help="Available commands")
    subparsers.required = True

    parser_create = subparsers.add_parser(
        "create", help="Create new upgrade scenario")
    parser_create.add_argument(
        "from_rel", metavar="FROM", help="The from release")
    parser_create.add_argument(
        "to_rel", metavar="TO", help="The to release")

    parser_run = subparsers.add_parser(
        "run", help="Run a specific scenario")
    parser_run.add_argument(
        "name", metavar="NAME", help="The name of the scenario")

    parser_status = subparsers.add_parser(
        "status", help="Get the status of scenario")
    parser_status.add_argument(
        "name", metavar="NAME", help="The name of the scenario")

    parser_status = subparsers.add_parser(
        "get-running-step", help="Get the current running step")
    parser_status.add_argument(
        "name", metavar="NAME", help="The name of the scenario")

    parser_stop = subparsers.add_parser(
        "stop", help="Stop running scenario")

    parser_matrix = subparsers.add_parser(
        "show-matrix", help="Show the matrix of scenario")
    parser_matrix.add_argument(
        "name", metavar="NAME", help="The name of the scenario")

    parser_list = subparsers.add_parser(
        "scenarios-list", help="Show existing scenarios")

    parser_delete = subparsers.add_parser(
        "delete", help="Delete a specific scenario")
    parser_delete.add_argument(
        "name", metavar="NAME", help="The name of the scenario")

    parser_get_log = subparsers.add_parser(
        "get-log", help="Return the log output of the step")
    parser_get_log.add_argument(
        "scenario_name", metavar="scenario_name",
        help="The name of the scenario")
    parser_get_log.add_argument(
        "step_name", metavar="step_name", help="The name of the step")
    parser_get_log.add_argument(
        "node_name", metavar="node_name", default='localhost',
        help="The name of the node (localhost for CE)")

    parser_get_failed_log = subparsers.add_parser(
        "get-failed-log", help="Return the log output of the failed step")
    parser_get_failed_log.add_argument(
        "scenario_name", metavar="scenario_name",
        help="The name of the scenario")

    args = parser.parse_args()

    if args.command == "create":
        create_scenario(args.from_rel, args.to_rel)
    elif args.command == "run":
        run_scenario(args.name)
    elif args.command == "status":
        get_scenario_status(args.name)
    elif args.command == "get-running-step":
        get_running_step(args.name)
    elif args.command == "stop":
        stop_scenario()
    elif args.command == "show-matrix":
        show_matrix(args.name)
    elif args.command == "scenarios-list":
        get_scenarios_list()
    elif args.command == "delete":
        delete_scenatio(args.name)
    elif args.command == "get-log":
        get_log(args.scenario_name, args.step_name, args.node_name)
    elif args.command == "get-failed-log":
        get_failed_log(args.scenario_name)


if __name__ == "__main__":
    main()
