import http.client
import urllib.request
import urllib.parse
import hashlib
import hmac
import base64
import json
import os
import time

API_KEY = "q1QjOkt6S0oOmmknkuh2MSyBqKE1suEGfNb6DlRm9mUpAhZsz7GFpnx7"
API_SECRET = "ZhIrO4Sh5e6khajbXmfGEYvnPE91EerdYXdDiPcJlBAUYNq5G4qr+tF6dwIH8vN/wTlBXCJyAJVZ3RkH4HaVfA=="
ASSET_MAP = {
        "XETH": "ETH",
        "XXBT": "BTC"
        }
ANSI_RESET = "\033[0m"
ANSI_BOLD = "\033[1m"
ANSI_GREEN = "\033[32m"
ANSI_RED = "\033[31m"
ANSI_CYAN = "\033[36m"
INTERVAL = 10

def main():
    previous_trade_balance = None
    previous_portfolio = None
    while True:
        trade_balance, portfolio = get_account_overview()
        os.system("clear")
        print(f"{ANSI_RESET}{ANSI_BOLD}Kraken Pro Portfolio{ANSI_RESET}")
        print("Currency\tBalance\t\tRate\t\tValue")
        for investment in portfolio:
            if previous_portfolio and previous_portfolio[investment] and previous_portfolio[investment]['value'] <= portfolio[investment]['value']:
                print(f"{ANSI_GREEN}{investment}\t\t{portfolio[investment]['balance']:010.8f}\t${portfolio[investment]['exchange_rate']:010.4f}\t${portfolio[investment]['value']:010.4f}{ANSI_RESET}")
            elif previous_portfolio and previous_portfolio[investment] and previous_portfolio[investment]['value'] > portfolio[investment]['value']:
                print(f"{ANSI_RED}{investment}\t\t{portfolio[investment]['balance']:010.8f}\t${portfolio[investment]['exchange_rate']:010.4f}\t${portfolio[investment]['value']:010.4f}{ANSI_RESET}")
            else:
                print(f"{ANSI_RESET}{investment}\t\t{portfolio[investment]['balance']:010.8f}\t${portfolio[investment]['exchange_rate']:010.4f}\t${portfolio[investment]['value']:010.4f}{ANSI_RESET}")
        print(f"{ANSI_RESET}------------------------------------------------------------")
        if previous_trade_balance and previous_trade_balance <= trade_balance:
            print(f"{ANSI_GREEN}Total\t\t\t\t\t\t${trade_balance:010.4f}{ANSI_RESET}")
        elif previous_trade_balance and previous_trade_balance > trade_balance:
            print(f"{ANSI_RED}Total\t\t\t\t\t\t${trade_balance:010.4f}{ANSI_RESET}")
        else:
            print(f"{ANSI_RESET}Total\t\t\t\t\t\t${trade_balance:010.4f}{ANSI_RESET}")
        time.sleep(INTERVAL)
        previous_trade_balance = trade_balance
        previous_portfolio = portfolio

def get_account_overview():
    account_balance = get_account_balance()
    if account_balance['error'] != []:
        for error in account_balance['error']:
            raise Exception(error) 
    else:
        account_balance = account_balance['result']
    investments = [each for each in account_balance if float(account_balance[each]) > 0]
    portfolio = {}
    for currency in investments:
        if currency in ASSET_MAP:
            pair = f"{ASSET_MAP[currency]}/USD"
        else:
            pair = f"{currency}/USD"
        ohlc = get_ohlc_data(pair)
        if ohlc['error'] != []:
            for error in ohlc['error']:
                print(f"{error}: {pair}")
                continue
        else:
            ohlc = ohlc['result']
            for pair in ohlc:
                portfolio[currency] = {
                        'balance': float(account_balance[currency]),
                        'exchange_rate': float(ohlc[pair][-2][4]),
                        'value': float(account_balance[currency]) * float(ohlc[pair][-2][4])
                        }
                if portfolio[currency]['value'] < 0.01:
                    portfolio.pop(currency)
                break
    trade_balance = get_trade_balance()
    if trade_balance['error'] != []:
        for error in trade_balance['error']:
            raise Exception(error) 
    else:
        trade_balance = float(trade_balance['result']['eb'])
    return trade_balance, portfolio
    

def get_ohlc_data(pair):
    response = request(
            method="GET",
            path="/0/public/OHLC",
            query={
                "pair": pair,
                "interval": 5,
            },
            environment="https://api.kraken.com",
        )
    return json.loads(response.read().decode())

def get_account_balance():
    response = request(
       method="POST",
       path="/0/private/Balance",
       public_key=API_KEY,
       private_key=API_SECRET,
       environment="https://api.kraken.com",
    )
    return json.loads(response.read().decode())

def get_trade_balance():
    response = request(
       method="POST",
       path="/0/private/TradeBalance",
       public_key=API_KEY,
       private_key=API_SECRET,
       environment="https://api.kraken.com",
    )
    return json.loads(response.read().decode())

def request(method: str = "GET", path: str = "", query: dict | None = None, body: dict | None = None, public_key: str = "", private_key: str = "", environment: str = "") -> http.client.HTTPResponse:
    url = environment + path
    query_str = ""
    if query is not None and len(query) > 0:
       query_str = urllib.parse.urlencode(query)
       url += "?" + query_str
    nonce = ""
    if len(public_key) > 0:
       if body is None:
          body = {}
       nonce = body.get("nonce")
       if nonce is None:
          nonce = get_nonce()
          body["nonce"] = nonce
    headers = {}
    body_str = ""
    if body is not None and len(body) > 0:
       body_str = json.dumps(body)
       headers["Content-Type"] = "application/json"
    if len(public_key) > 0:
       headers["API-Key"] = public_key
       headers["API-Sign"] = get_signature(private_key, query_str+body_str, nonce, path)
    req = urllib.request.Request(
       method=method,
       url=url,
       data=body_str.encode(),
       headers=headers,
    )
    return urllib.request.urlopen(req)

def get_nonce() -> str:
    return str(int(time.time() * 1000))

def get_signature(private_key: str, data: str, nonce: str, path: str) -> str:
    return sign(
       private_key=private_key,
       message=path.encode() + hashlib.sha256(
             (nonce + data)
          .encode()
       ).digest()
    )

def sign(private_key: str, message: bytes) -> str:
    return base64.b64encode(
       hmac.new(
          key=base64.b64decode(private_key),
          msg=message,
          digestmod=hashlib.sha512,
       ).digest()
    ).decode()


if __name__ == "__main__":
    main()