import csv
import paramiko
import time
from cube_activate_ssh import activate_ssh
from dotenv import load_dotenv
import os
import re
from azure.iot.hub import IoTHubRegistryManager
from azure.iot.hub.models import Twin, TwinProperties

load_dotenv(override=True)

ip_address_prefix = "10.81.56."
ssh_command = "hostname"

csv_filename = "DK2_01.csv"
SITE_NAME = "DK2"

ssh_username = os.getenv("DEFAULT_CUBE_LINUX_ADMIN_USER")
ssh_password = os.getenv("DEFAULT_CUBE_LINUX_ADMIN_PASSWORD")
CONNECTION_STRING = str(os.getenv("CONNECTION_STRING_INOX_PROD"))

def execute_ssh_command(ip, command):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(ip, port=11022, username=ssh_username, password=ssh_password, allow_agent=False, look_for_keys=False)
        stdin, stdout, stderr = client.exec_command(command)
        result = stdout.read().decode().lower().strip()
        return result
    except Exception as e:
        print(f"SSH Error: {str(e)}")
        raise
    finally:
        client.close()

def update_cloud_config(ip, new_content):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(ip, port=11022, username=ssh_username, password=ssh_password, allow_agent=False, look_for_keys=False)
        stdin, stdout, stderr = client.exec_command(f'sudo -S bash -c \'cat > /etc/cube/config-azure.properties << EOF\n{new_content}\nEOF\'\n')
        stdin.write(ssh_password + "\n")
        stdin.flush()
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        for output in stdoutput:
            print(output.strip())
    except Exception as e:
        print(f"SSH Error: {str(e)}")
        raise
    finally:
        client.close()

def restart_cloudagent(ip):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(ip, port=11022, username=ssh_username, password=ssh_password, allow_agent=False, look_for_keys=False)
        stdin, stdout, stderr = client.exec_command(f'sudo -S bash -c \'systemctl restart cube-web-cloudagent << EOF\n\nEOF\'\n')
        stdin.write(ssh_password + "\n")
        stdin.flush()
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        for output in stdoutput:
            print(output.strip())
    except Exception as e:
        print(f"SSH Error: {str(e)}")
        raise
    finally:
        client.close()

def main():

    print("Starting...")

    with open(csv_filename, mode="w", newline="") as file:
        writer = csv.writer(file)
        writer.writerow(["Number", "IP address", "Cube ID", "Environment", "Correct configuration"])

        registry_manager = IoTHubRegistryManager.from_connection_string(CONNECTION_STRING)

        results = []

        numbers = [193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214]
        for i in numbers:
            ip_address = f"{ip_address_prefix}{i}"
            print(f"Activating SSH for {ip_address}:", end=" ")
            
            try:
                activate_ssh(ip_address)
            except Exception as e:
                print("Failed!")
                writer.writerow([i, ip_address, "UNREACHABLE", "NA", "NA"])
                file.flush()
                continue

            print("Activated!")

            print(f"Executing {ssh_command} for {ip_address}:", end=" ")
            try:
                cube_id = execute_ssh_command(ip_address, ssh_command)
            except Exception as e:
                print("Failed!")
                writer.writerow([i, ip_address, "UNREACHABLE", "NA", "NA"])
                file.flush()
                continue
            print(cube_id)

            print(f"Getting configured Connection String")
            try:
                connection_string = execute_ssh_command(ip_address, "grep \"connection-string\" /etc/cube/config-azure.properties")
                if connection_string == "":
                    raise Exception("No Connection String extracted!")

                iothub_match = re.search(r"hostname\\=(.*?);", connection_string, re.IGNORECASE)
                iothub = iothub_match.group(1) if iothub_match else None
                if iothub.lower() == "IotHub-CUBE-PROD.azure-devices.net".lower():
                    migration = "SAFT"
                elif iothub.lower() == "iot-ingest-ess-prod.azure-devices.net".lower():
                    migration = "INOX"
                else:
                    migration = "NONE"

                device_id_match =  re.search(r"deviceid\\=(.*?);", connection_string, re.IGNORECASE)
                cloud_cube_id = device_id_match.group(1) if device_id_match else None
                if cloud_cube_id.lower() == cube_id.lower():
                    status = "CORRECT"
                else:
                    status = "INCORRECT"

            except Exception as e:
                print(e)
                migration = "NONE"
                status = "INCORRECT"

            if migration == "SAFT" or migration == "NONE":
                print("SAFT device, migrating to INOX...")

                print("Creating device on INOX...")
                try:
                    registry_manager.create_device_with_sas(
                        cube_id,
                        primary_key="", secondary_key="",
                        status="enabled",
                        iot_edge=False
                    )
                except Exception as iot_e:
                    print("Error creating new device!")
                    print(iot_e)
                    continue

                print("Adding tags to new device...")
                try:
                    twin = registry_manager.get_twin(cube_id)
                    twin_patch = Twin(properties=TwinProperties(desired={}), tags={
                        "site": SITE_NAME,
                        "number": i
                    })
                    registry_manager.update_twin(cube_id, twin_patch, twin.etag)
                except Exception as iot_e:
                    print("Error assigning tags to new device!")
                    print(iot_e)

                print("Requesting primary key...")
                try:
                    device_info = registry_manager.get_device(cube_id)
                    primary_key = device_info.authentication.symmetric_key.primary_key
                    new_connection_string = f"HostName\\={CONNECTION_STRING.split(';')[0].split('=')[1]};DeviceId\\={cube_id};SharedAccessKey\\={primary_key}"
                    new_content = f'light-telemetry=false\ncompression-enabled=true\ntelemetry-on=true\nremote-update-on=true\nconnection-string={new_connection_string}'
                except Exception as iot_e:
                    print("Error getting new Connection String!")
                    print(iot_e)
                    continue
                print("Setting new Connection String...")
                try:
                    update_cloud_config(ip_address, new_content)
                except Exception as ssh_e:
                    print("Error when setting the new Connection String!")
                    print(ssh_e)
                    continue
            print("Done!")

            restart_cloudagent(ip_address)

            writer.writerow([i, ip_address, cube_id, migration, status])
            file.flush()

if __name__ == "__main__":
    main()