import json

from braket.device_schema.iqm import IqmDeviceCapabilities
from networkx import DiGraph

from qhub.quantum.sdk.braket.braket_provider import HubBraketProvider
from qhub.quantum.sdk.braket.gate_based_device import HubAwsGateBasedDevice


@HubBraketProvider.register_device("aws.iqm.garnet")
class HubAwsIqmGarnetDevice(HubAwsGateBasedDevice):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @property
    def name(self) -> str:
        return "Garnet"

    @property
    def provider_name(self) -> str:
        return "IQM"

    @property
    def properties(self) -> IqmDeviceCapabilities:
        """IqmDeviceCapabilities: Return the device properties"""
        config = self._get_backend_config()
        return IqmDeviceCapabilities.parse_raw(json.dumps(config))

    @property
    def topology_graph(self) -> DiGraph:
        """DiGraph: topology of device with all qubits as nodes.

        Ensures all qubits from oneQubitProperties are present in the graph,
        even if they have no connections (e.g., defective/unavailable qubits).

        This fixes an issue where the live AWS Braket connectivityGraph may
        omit qubits that have no connections due to errors, but oneQubitProperties still
        lists them, causing KeyError in qiskit-braket-provider's adapter.py.
        """
        graph = super()._construct_topology_graph()

        # Ensure all qubits from oneQubitProperties are in the graph
        props = self.properties
        if hasattr(props, 'standardized') and props.standardized:
            one_qubit_props = props.standardized.oneQubitProperties
            if one_qubit_props:
                for qubit_key in one_qubit_props.keys():
                    qubit_id = int(qubit_key)
                    if qubit_id not in graph.nodes:
                        graph.add_node(qubit_id)

        return graph
