#!/usr/bin/env python3

import os
import subprocess
import time
    
#-----------------
# Settings

MONGO_BASE_PATH = '/opt/my-mongo'
REPLICASET_NAME = 'techtransthai-db'
NETWORK_NAME = 'mongonet'
NUMBER_OF_CONTAINERS = '5'
MAX_MEM_ALLOC_PERCENT = '95'
MONGO_INITDB_ROOT_USERNAME = 'sasha'
MONGO_INITDB_ROOT_PASSWORD = '12345'
#-----------------

if os.getuid() != 0:
    print("You must be root to run this script")
    exit()

if int(NUMBER_OF_CONTAINERS) < 1 and int(NUMBER_OF_CONTAINERS) > 7:
    print("Number of containers must be between 1 and 7")
    exit()

if int(MAX_MEM_ALLOC_PERCENT) < 50 and int(NMAX_MEM_ALLOC_PERCENT) > 100:
    print("Memory allocation must be between 50% and 100%")
    exit()

print(f"\n\033[1;32mInstalling Python dependencies and OpenSSL from Ubuntu repo\033[0m\n")

subprocess.run("apt install -y python3-pymongo python3-psutil openssl", shell=True, check=True)

print(f"\n\033[1;32mGenerating keys\033[0m\n")

keyfile_location = MONGO_BASE_PATH + "/mongodb-keyfile"
subprocess.run("openssl rand -base64 741 > " + keyfile_location, shell=True, check=True)
subprocess.run("chmod 600 " + keyfile_location, shell=True, check=True)
subprocess.run("chown 999 " + keyfile_location, shell=True, check=True)

print(f"\n\033[1;32mInstalling Docker CE from official repo\033[0m\n")

docker_install="""
apt update
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
chmod a+r /etc/apt/keyrings/docker.asc

echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \
  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
  tee /etc/apt/sources.list.d/docker.list > /dev/null
apt update 

apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
"""

subprocess.run(docker_install, shell=True, check=True)

print(f"\n\033[1;32mCreating Docker network {NETWORK_NAME} for MongoDB Cluster\033[0m\n")

network_command = "docker network create " + NETWORK_NAME

subprocess.run(network_command, shell=True, check=True)

print(f"\n\033[1;32mPulling MongoDB Image\033[0m\n")

subprocess.run("docker pull mongo", shell=True, check=True)

import psutil

mem_per_container = f'{(psutil.virtual_memory().total/1073741824)*(int(MAX_MEM_ALLOC_PERCENT)/(int(NUMBER_OF_CONTAINERS)*100)):.2f}'

for container_number in range(int(NUMBER_OF_CONTAINERS)):
  print(f"\n\033[1;32mStarting MongoDB Container #{container_number} with {mem_per_container} GB of memory\033[0m\n")
  docker_cmd = f"docker run -d -p 3000{container_number}:27017 --name mongo{container_number} --net {NETWORK_NAME} -e MONGO_INITDB_ROOT_USERNAME={MONGO_INITDB_ROOT_USERNAME} -e MONGO_INITDB_ROOT_PASSWORD={MONGO_INITDB_ROOT_PASSWORD} -v {MONGO_BASE_PATH}/mongo{container_number}:/data/db -v {MONGO_BASE_PATH}:/opt/keyfile  --restart unless-stopped mongo mongod --replSet {REPLICASET_NAME} --wiredTigerCacheSizeGB {mem_per_container} --keyFile /opt/keyfile/mongodb-keyfile"

  subprocess.run(docker_cmd, shell=True, check=True)

print(f"\n\033[1;32mWaiting for all containers to start\033[0m\n")

time.sleep(10)
#----------------------------------------------------------------------------------

print(f"\n\033[1;32mInitializing replica set {REPLICASET_NAME} of {NUMBER_OF_CONTAINERS} containers\033[0m\n")

from pymongo import MongoClient
from pymongo.errors import PyMongoError

# Connect to the MongoDB instances

client = MongoClient('mongodb://sasha:12345@localhost:30000')

# Initialize a replica set
try:
    # Build the replica set configuration
    config = {
    "_id": REPLICASET_NAME,
    # "keyFile": "/opt/keyfile/mongodb-keyfile",
    "members": []
    }

    for i in range(int(NUMBER_OF_CONTAINERS)):
        member = {
            "_id": i,
            "host": "mongo"+str(i)
        }
        config["members"].append(member)

    # Initialize the replica set
    client.admin.command('replSetInitiate', config)
    print("Replica set "+ REPLICASET_NAME + " initialized successfully.")

except PyMongoError as e:
    print("Error initializing replica set:", e)