import argparse
from enum import Enum
import glob
import platform
import re
import sys
import os
from typing import List

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../mtcr_py")))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../reg_access")))

from mtcr import MstDevice, MtcrException
from regaccess import RegAccess, RegAccException

try:
    import mst_vmdl_int
    from mst_vmdl_int import CPO_DEVICE
except ImportError:
    mst_vmdl_int = None
    class CPO_DEVICE(Enum):
        V_MODULE = 0

V_MODULE_MASK = 0xFF
DEV_TYPE_REG_ADDR = 0xf0014
DEVICE_ID_MASK = 0xFFFF
SWITCH_GLOB = "/dev/mst/mt5????_pciconf*"

CPO_SUPPORTED_DEVS = [
    0x270,  # Spectrum 5
]

V_MODULE_DEV_ID_ADDRESS = 0x0
V_MODULE_DEV_IDS = [0x0]  # TODO: change when we have a devid

ADDRESS_SIZE = 1

VERBOSE = False


class CpoAsic:
    def __init__(self, device):
        self.device = device
        self.mst_device = MstDevice(device)
        self.reg_access = RegAccess(dev=self.mst_device)

        self.mgir = self.reg_access.getMGIR()
        self.ga = self.mgir.hw_info.ga

    def close(self):
        self.reg_access.close()
        self.mst_device.close()

    def is_cpo_eligible_device(self):
        could_read = True
        try:
            output = self.mst_device.read4(DEV_TYPE_REG_ADDR)
            output = output & DEVICE_ID_MASK
        except BaseException as e:
            could_read = False

        return could_read and output in CPO_SUPPORTED_DEVS

    def is_cpo_device(self):
        try:
            return self.mgir.hw_info.cpo_indication == 1
        except RegAccException:
            return False

    def num_of_modules_per_system(self):
        try:
            num_of_modules = self.reg_access.num_of_modules_per_system()
        except RegAccException:
            num_of_modules = 0

        return num_of_modules

    def mmam(self, index):
        try:
            mmam = self.reg_access.get_mmam(index)
        except RegAccException:
            mmam = None

        return mmam


def verbose_print(*args, **kwargs):
    if VERBOSE:
        print(*args, **kwargs)


def get_dev_id(dev_name, address):
    dev = None
    try:
        # Trim /dev/mst/ from the device name
        dev_name = dev_name[9:]
        # Run mopen on newly created device
        # If device is non-responsive or if mread fails, exception will be thrown
        dev = MstDevice(dev_name)
        id = dev.read4(address) & V_MODULE_MASK
        return id
    except MtcrException:
        pass
    finally:
        if dev:
            dev.close()


def create_device(device_type: CPO_DEVICE, address, dev_name, dev_ids, discover_err_message):
    try:
        with open(dev_name, "w"):
            pass
        # For SPC5, CMIS of vmdl won't be ready
        # For SPC6 remove this line
        if device_type == CPO_DEVICE.V_MODULE:
            return True
        dev_id = get_dev_id(dev_name, address)
        if dev_id is None:
            verbose_print(f"Failed to read device id for {dev_name}")
            if os.path.isfile(dev_name):
                os.remove(dev_name)
            return False

        if dev_id not in dev_ids:
            verbose_print(discover_err_message)
            if os.path.isfile(dev_name):
                os.remove(dev_name)
            return False
        return True

    except BaseException as e:
        verbose_print(f"Failed to open device {dev_name}")
        if os.path.isfile(dev_name):
            os.remove(dev_name)
        return False


def parse_args(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--verbose', '-v', action='store_true', help='Enable verbose output')
    if mst_vmdl_int is not None:
        mst_vmdl_int.add_args_int(parser)
    return parser.parse_args(args=args)


def main(args=None):
    global VERBOSE
    # Only check for root on linux
    if platform.system() != "Windows" and os.name != "nt" and os.geteuid() != 0:
        print("This script must be run as root.")
        sys.exit(1)

    args = parse_args(args=args)
    VERBOSE = args.verbose if args else False

    all_cpo_asics: List[CpoAsic] = []
    for asic in glob.glob(SWITCH_GLOB):
        # Filter out cables
        if not re.search(r"mt5\d{4}_pciconf\d+$", asic):
            continue
        a = CpoAsic(asic)
        if a.is_cpo_eligible_device() and a.is_cpo_device():
            all_cpo_asics.append(a)

    if len(all_cpo_asics) == 0:
        verbose_print("-I- No virtual module ASICs found")
        return

    num_devs_added = 0
    asic_ga_dict = {s.ga: s for s in all_cpo_asics}

    try:
        asic = all_cpo_asics[0]
        total_modules = asic.num_of_modules_per_system()
        for i in range(total_modules):
            mmam = asic.mmam(i)
            if mmam.module_type != 25:  # CPO
                continue

            primary_asic = asic_ga_dict[mmam.ga]

            local_module_num = mmam.local_module
            dev_name = f'{primary_asic.device}_vmdl_{local_module_num}'

            if not create_device(
                CPO_DEVICE.V_MODULE,
                V_MODULE_DEV_ID_ADDRESS,
                dev_name,
                V_MODULE_DEV_IDS,
                f"Module {local_module_num} on {primary_asic.device} is not a virtual module"
            ):
                continue
            num_devs_added += 1

            if mst_vmdl_int is not None:
                mst_vmdl_int.create_devices_int(args, create_device, dev_name)

    finally:
        for asic in all_cpo_asics:
            asic.close()

    print(f"-I- Added {num_devs_added} virtual modules")
    if mst_vmdl_int is not None:
        mst_vmdl_int.print_summary(args)


if __name__ == "__main__":
    main()
