#!/usr/bin/env python
import argparse
from datetime import datetime, timezone
from io import StringIO
import os
from subprocess import PIPE, Popen

import skyway
from skyway.cloud.aws import *
from skyway.cloud.gcp import *
from skyway.cloud.azure import *
from skyway.cloud.slurm import *
from skyway.cloud.oci import *

import colorama
from colorama import Fore

# export SKYWAYROOT=/project/rcc/trung/skyway-github
#./skyway_list --account=rcc-aws

class InstanceDescriptor:
    def __init__(self, jobname: str, account_name: str, node_type: str, walltime: str, vendor_name: str):
        self.jobname = jobname
        self.account_name = account_name
        self.node_type = node_type
        self.walltime = walltime
        self.vendor_name = vendor_name

        self.account = None
        if 'aws' in vendor_name:
            self.account = AWS(account_name)
        elif 'gcp' in vendor_name:
            self.account = GCP(account_name)
        elif 'azure' in vendor_name:
            self.account = AZURE(account_name)
        elif 'oci' in vendor_name:
            self.account = OCI(account_name)
        elif 'midway3' in vendor_name:
            self.account = SLURMCluster(account_name)

        self.user = os.environ['USER']

    def getUsage(self, user_name):
        # retrieve from database for the given account
        user_budget = self.account.get_budget(user_name=user_name, verbose=False)
        accumulating_cost, remaining_balance = self.account.get_cost_and_usage_from_db(user_name=user_name)
        return user_budget, accumulating_cost, remaining_balance

    def getUsageHistory(self, user_name):
        df = self.account.get_usage_history_from_db(user_name=user_name)
        return df

if __name__ == "__main__":

    colorama.init(autoreset=True)

    default_user = os.environ['USER']

    msg = "Skyway usage of a user"
    
    parser = argparse.ArgumentParser(description=msg)
    parser.add_argument('-u', '--user',  dest='username', default=default_user, help="User name")
    parser.add_argument('-A', '--account', dest='account', default="", help="Account name")
    parser.add_argument('--byjob', dest='byjob', action='store_true', default=False, help="Show jobs if specified")
    parser.add_argument('--provider', dest='provider', default="", help="Vendor name: AWS, GCP, Azure, or RCC Midway")
    
    args = parser.parse_args()

    user_name = args.username
    account_name = args.account
    provider = args.provider.lower()
    byjob = args.byjob

    if provider == "":
        # try to infer the vendor name from account
        if 'aws' in account_name:
            vendor_name = "aws"
        elif 'gcp' in account_name:
            vendor_name = "gcp"
        elif 'azure' in account_name:
            vendor_name = "azure"
        elif 'oci' in account_name:
            vendor_name = "oci"
        elif 'midway3' in account_name or 'rcc-staff' in account_name:
            vendor_name = "rcc-midway3"
    else:
        vendor_name = provider

    # get the env variable SKYWAYROOT
    skywayroot = os.environ['SKYWAYROOT']
    if skywayroot == "":
        raise Exception("SKYWAYROOT is not defined.")

    # create an instance descriptor (like with the dashboard)
    instanceDescriptor = InstanceDescriptor("", account_name, "", "", vendor_name)

    # usage
    headers=["User", 'Allocation', 'Usage', 'Balance']
    user_budget, usage, balance = instanceDescriptor.getUsage(user_name)

    data = [[user_name, user_budget, usage, balance]]
    print(tabulate(data, headers=headers))

    if byjob == True:
        df = instanceDescriptor.getUsageHistory(user_name)

        if df.empty == True:
            headers = ['User','InstanceID','InstanceType','Start','End']
            print(tabulate([], headers=headers))
            print("")
        else:
            print(df)
