import psycopg2
# Manages operations on the database
import configparser
from random import randint

# Load Configurations
config = configparser.ConfigParser()
config.read('bgp.conf')
db_conf = config['bgp']


# Connect to database
try:
    conn = psycopg2.connect(dbname=db_conf['database'],
                            user=db_conf['user'],
                            password=db_conf['password'],
                            port=db_conf['port'],
                            host=db_conf['host'])
    conn.close()
except:
    print("[FAILED] Connect to database")



def load_mrt_announcements(victim_asn, attacker_asn):
    cur = conn.cursor()
    # Load Victim
    cur.execute(
        """
        INSERT INTO mrt_announcements (origin, as_path, prefix) VALUES (%s, %s, %s)
        """, (victim_asn, '{{0}}'.format(victim_asn), '1.2.0.0/16')
    )
    # Load Attacker
    cur.execute(
        """
        INSERT INTO mrt_announcements (origin, as_path, prefix) VALUES (%s, %s, %s)
        """, (attacker_asn, '{{0}}'.format(attacker_asn), '1.2.3.0/24')
    )


def select_rov_ases(percentROV):
    """

    """
    cur = conn.cursor()
    # Get Total Count of ASes
    cur.execute("SELECT COUNT(*) AS count FROM ases", ())
    num_ases = cur.fetchone()[0]
    # Calculate the number of ROV ases needed
    num_rov_ases = percentROV * num_ases
    rov_ases = set()
    for i in range(num_rov_ases):
        random_asn = randint(1, num_rov_ases)
        while random_asn in rov_ases:
            random_asn = randint(1, num_rov_ases)
        cur.execute("UPDATE ases SET as_type=%s WHERE asn=%s", ('rov', random_asn))
    raise NotImplementedError


def cleanup():
    raise NotImplementedError
