import json
import os
import requests
from dotenv import load_dotenv
from pprint import pprint
from constants import URL, DOMAIN, IPV4, IPV6

# Would be nice to define some constants, f.e. for the various API urls, the headers, etc.


def environment():
    load_dotenv()
    api_key = os.getenv("VT_API")
    return api_key


# Unfortunately this works for actual URLs, not domains. See: https://docs.virustotal.com/reference/domain-info
# This also doesn't work for IPv6 addresses, where the response_dict does not have a 'data' key. So I would have to revamp this module and make separate functions called based on host type (URL, IPv4 and -6, domain).


def analysis_object(api_key, host):
    url = "https://www.virustotal.com/api/v3/urls"
    payload = {"url": ""}
    payload.update({"url": host})
    headers = {
        "accept": "application/json",
        "content-type": "application/x-www-form-urlencoded",
        "x-apikey": api_key,
    }
    response = requests.post(url, data=payload, headers=headers)
    response_dict = json.loads(response.text)
    response_id = response_dict["data"]["id"]
    return response_id


def analyse_domain(api_key, host):
    url = "https://www.virustotal.com/api/v3/domains/" + host
    headers = {
        "accept": "application/json",
        "content-type": "application/x-www-form-urlencoded",
        "x-apikey": api_key,
    }
    analysis_response = requests.get(url, headers=headers)
    response_dict = json.loads(analysis_response.text)
    # Probably still need to turn the requests.get into a json like below
    return response_dict


def analyse_URL(api_key, response_id):
    analysis_url = "https://www.virustotal.com/api/v3/analyses/{}".format(response_id)
    headers = {"accept": "application/json", "x-apikey": api_key}
    analysis_response = requests.get(analysis_url, headers=headers)
    analysis_dict = json.loads(analysis_response.text)
    # return analysis_response.text
    return analysis_dict


def analyse_IP(api_key, host):
    analysis_url = "https://www.virustotal.com/api/v3/ip_addresses/{}".format(host)
    headers = {"accept": "application/json", "x-apikey": api_key}
    analysis_response = requests.get(analysis_url, headers=headers)
    analysis_dict = json.loads(analysis_response.text)
    # Implement this: https://docs.virustotal.com/reference/ip-info
    return analysis_dict


def analyse(host, host_type):
    api_key = environment()
    if host_type == URL:
        response_id = analysis_object(api_key, host)
        result = analyse_URL(api_key, response_id)
    elif host_type == DOMAIN:
        result = analyse_domain(api_key, host)
    # elif for IPv4 and IPv6.
    elif host_type == IPV4 or IPV6:
        result = analyse_IP(api_key, host)
    return result