# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# define error causes as documented in https://docs.nvidia.com/deploy/xid-errors/index.html
from typing import Optional


class ErrorCause:
    HW_ERROR = "HW Error"
    DRIVER_ERROR = "Driver Error"
    USER_APP_ERROR = "User App Error"
    SYSTEM_MEMORY_CORRUPTION = "System Memory Corruption"
    BUS_ERROR = "Bus Error"
    THERMAL_ISSUE = "Thermal Issue"
    FB_CORRUPTION = "FB Corruption"
    CUSTOM_FAIR = "Custom XID Error List from Meta Corp"

    XID_TO_CAUSES = {
        1: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        2: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        3: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        4: {
            DRIVER_ERROR,
            USER_APP_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            FB_CORRUPTION,
        },
        5: {},
        6: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        7: {DRIVER_ERROR, BUS_ERROR, FB_CORRUPTION},
        8: {DRIVER_ERROR, USER_APP_ERROR, BUS_ERROR, THERMAL_ISSUE},
        9: {DRIVER_ERROR},
        10: {},
        11: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        12: {DRIVER_ERROR},
        13: {
            HW_ERROR,
            DRIVER_ERROR,
            USER_APP_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            THERMAL_ISSUE,
            FB_CORRUPTION,
        },
        14: {},
        15: {},
        16: {DRIVER_ERROR},
        17: {},
        18: {DRIVER_ERROR},
        19: {DRIVER_ERROR},
        20: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        21: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        22: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR, FB_CORRUPTION},
        23: {},
        24: {
            DRIVER_ERROR,
            USER_APP_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            THERMAL_ISSUE,
            FB_CORRUPTION,
        },
        25: {
            DRIVER_ERROR,
            USER_APP_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            FB_CORRUPTION,
        },
        26: {DRIVER_ERROR},
        27: {DRIVER_ERROR},
        28: {DRIVER_ERROR},
        29: {DRIVER_ERROR},
        30: {DRIVER_ERROR},
        31: {HW_ERROR, DRIVER_ERROR, USER_APP_ERROR},
        32: {
            DRIVER_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            THERMAL_ISSUE,
            FB_CORRUPTION,
        },
        33: {DRIVER_ERROR},
        34: {DRIVER_ERROR},
        35: {DRIVER_ERROR},
        36: {DRIVER_ERROR},
        37: {DRIVER_ERROR, SYSTEM_MEMORY_CORRUPTION, BUS_ERROR},
        38: {DRIVER_ERROR},
        39: {},
        40: {},
        41: {},
        42: {DRIVER_ERROR},
        43: {DRIVER_ERROR, USER_APP_ERROR},
        44: {DRIVER_ERROR},
        45: {DRIVER_ERROR},
        46: {DRIVER_ERROR},
        47: {DRIVER_ERROR},
        48: {HW_ERROR},
        49: {},
        50: {},
        51: {},
        52: {},
        53: {},
        54: {},
        55: {},
        56: {HW_ERROR, DRIVER_ERROR},
        57: {HW_ERROR, DRIVER_ERROR, FB_CORRUPTION},
        58: {HW_ERROR, DRIVER_ERROR},
        59: {DRIVER_ERROR},
        60: {DRIVER_ERROR},
        61: {},
        62: {HW_ERROR, DRIVER_ERROR, THERMAL_ISSUE},
        63: {HW_ERROR, DRIVER_ERROR, FB_CORRUPTION},
        64: {HW_ERROR, DRIVER_ERROR},
        65: {HW_ERROR, DRIVER_ERROR},
        66: {DRIVER_ERROR, USER_APP_ERROR},
        67: {DRIVER_ERROR, USER_APP_ERROR},
        68: {HW_ERROR, DRIVER_ERROR},
        69: {HW_ERROR, DRIVER_ERROR},
        70: {HW_ERROR, DRIVER_ERROR},
        71: {HW_ERROR, DRIVER_ERROR},
        72: {HW_ERROR, DRIVER_ERROR},
        73: {HW_ERROR, DRIVER_ERROR},
        74: {HW_ERROR, DRIVER_ERROR, BUS_ERROR},
        75: {HW_ERROR, DRIVER_ERROR},
        76: {HW_ERROR, DRIVER_ERROR},
        77: {HW_ERROR, DRIVER_ERROR},
        78: {DRIVER_ERROR},
        79: {
            HW_ERROR,
            DRIVER_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            THERMAL_ISSUE,
        },
        80: {
            HW_ERROR,
            DRIVER_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            FB_CORRUPTION,
        },
        81: {HW_ERROR},
        82: {HW_ERROR, DRIVER_ERROR},
        83: {HW_ERROR, DRIVER_ERROR},
        84: {HW_ERROR, DRIVER_ERROR},
        85: {HW_ERROR, DRIVER_ERROR},
        86: {HW_ERROR, DRIVER_ERROR},
        87: {},
        88: {HW_ERROR, DRIVER_ERROR},
        89: {HW_ERROR, DRIVER_ERROR},
        90: {HW_ERROR, DRIVER_ERROR},
        91: {HW_ERROR, DRIVER_ERROR},
        92: {HW_ERROR, DRIVER_ERROR},
        93: {DRIVER_ERROR, USER_APP_ERROR},
        94: {HW_ERROR, DRIVER_ERROR, FB_CORRUPTION},
        95: {HW_ERROR, DRIVER_ERROR, FB_CORRUPTION},
        96: {HW_ERROR, DRIVER_ERROR},
        97: {HW_ERROR, DRIVER_ERROR},
        98: {HW_ERROR, DRIVER_ERROR},
        99: {HW_ERROR, DRIVER_ERROR},
        100: {HW_ERROR, DRIVER_ERROR},
        101: {HW_ERROR, DRIVER_ERROR},
        102: {HW_ERROR, DRIVER_ERROR},
        103: {HW_ERROR, DRIVER_ERROR},
        104: {HW_ERROR, DRIVER_ERROR},
        105: {HW_ERROR, DRIVER_ERROR},
        106: {USER_APP_ERROR},
        107: {USER_APP_ERROR},
        108: {},
        109: {
            HW_ERROR,
            DRIVER_ERROR,
            USER_APP_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            THERMAL_ISSUE,
            FB_CORRUPTION,
        },
        110: {HW_ERROR},
        111: {HW_ERROR, DRIVER_ERROR, BUS_ERROR},
        112: {HW_ERROR, DRIVER_ERROR},
        113: {HW_ERROR, DRIVER_ERROR},
        114: {HW_ERROR, DRIVER_ERROR, FB_CORRUPTION},
        115: {HW_ERROR, DRIVER_ERROR},
        116: {HW_ERROR, DRIVER_ERROR},
        117: {HW_ERROR, DRIVER_ERROR},
        118: {HW_ERROR, DRIVER_ERROR},
        119: {
            HW_ERROR,
            DRIVER_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            THERMAL_ISSUE,
            FB_CORRUPTION,
        },
        120: {
            HW_ERROR,
            DRIVER_ERROR,
            SYSTEM_MEMORY_CORRUPTION,
            BUS_ERROR,
            THERMAL_ISSUE,
            FB_CORRUPTION,
        },
        121: {HW_ERROR, BUS_ERROR},
        122: {HW_ERROR, DRIVER_ERROR},
        123: {HW_ERROR, DRIVER_ERROR},
        124: {HW_ERROR, DRIVER_ERROR},
        125: {HW_ERROR, DRIVER_ERROR},
        126: {},
        127: {},
        128: {},
        129: {},
        130: {},
        131: {},
        132: {},
        133: {},
        134: {},
        135: {},
        136: {},
        137: {HW_ERROR, DRIVER_ERROR, USER_APP_ERROR},
        138: {},
        139: {},
        140: {HW_ERROR, DRIVER_ERROR, FB_CORRUPTION},
        141: {},
        142: {},
        143: {HW_ERROR, DRIVER_ERROR, FB_CORRUPTION},
    }

    @classmethod
    def get_causes_for_xid(cls, xid: Optional[int]) -> list[str]:
        if xid is None:
            return []
        return sorted(cls.XID_TO_CAUSES.get(xid, []))

    SOFT_ERROR_TYPE = {DRIVER_ERROR, USER_APP_ERROR}
    HARD_ERROR_TYPE = {
        HW_ERROR,
        SYSTEM_MEMORY_CORRUPTION,
        BUS_ERROR,
        THERMAL_ISSUE,
        FB_CORRUPTION,
    }

    # non-critical errors: https://fburl.com/jtkc8d6m
    # 63, and 92 comes from ad-hoc conversation with NVIDIA engineers: https://fburl.com/3ornsec3
    # Non-critical XIDs by NVIDIA: https://github.com/NVIDIA/k8s-device-plugin/blob/main/internal/rm/health.go#L65
    # XIDs docs: http://docs.nvidia.com/deploy/xid-errors/index.html#topic_4
    NON_CRITICAL_ERRORS = {13, 31, 43, 45, 63, 68, 92, 109}
