import gc
import numpy as np
from typing import Any


class GatePlugin:
    name   = "UNNAMED_GATE"
    matrix = None

    def apply(self, circuit, *qubits):
        raise NotImplementedError("Override apply() in your GatePlugin subclass.")

    def validate(self):
        if self.matrix is None:
            raise ValueError(f"Gate '{self.name}' has no matrix defined.")
        m = np.array(self.matrix)
        if m.ndim != 2 or m.shape[0] != m.shape[1]:
            raise ValueError(f"Gate '{self.name}' matrix must be square.")
        eye  = np.eye(m.shape[0])
        prod = m.conj().T @ m
        if not np.allclose(prod, eye, atol=1e-6):
            raise ValueError(f"Gate '{self.name}' matrix is not unitary.")

    def __repr__(self):
        return f"GatePlugin(name='{self.name}')"


class BackendPlugin:
    name = "UNNAMED_BACKEND"

    def run(self, circuit, shots=1024):
        raise NotImplementedError("Override run() in your BackendPlugin subclass.")

    def validate(self):
        if not hasattr(self, "run"):
            raise ValueError(f"Backend '{self.name}' must implement run().")

    def __repr__(self):
        return f"BackendPlugin(name='{self.name}')"


class AlgorithmPlugin:
    name = "UNNAMED_ALGORITHM"

    def execute(self, circuit, **kwargs):
        raise NotImplementedError("Override execute() in your AlgorithmPlugin subclass.")

    def validate(self):
        if not hasattr(self, "execute"):
            raise ValueError(f"Algorithm '{self.name}' must implement execute().")

    def __repr__(self):
        return f"AlgorithmPlugin(name='{self.name}')"


class PluginManager:
    def __init__(self):
        self._gates      = {}
        self._backends   = {}
        self._algorithms = {}

    def register_gate(self, plugin, overwrite=False):
        if not isinstance(plugin, GatePlugin):
            raise TypeError("plugin must be a GatePlugin instance.")
        plugin.validate()
        if plugin.name in self._gates and not overwrite:
            raise ValueError(f"Gate '{plugin.name}' already registered. Use overwrite=True.")
        self._gates[plugin.name] = plugin
        print(f"  [PluginManager] ✅ Gate registered    : {plugin.name}")

    def register_backend(self, plugin, overwrite=False):
        if not isinstance(plugin, BackendPlugin):
            raise TypeError("plugin must be a BackendPlugin instance.")
        plugin.validate()
        if plugin.name in self._backends and not overwrite:
            raise ValueError(f"Backend '{plugin.name}' already registered. Use overwrite=True.")
        self._backends[plugin.name] = plugin
        print(f"  [PluginManager] ✅ Backend registered  : {plugin.name}")

    def register_algorithm(self, plugin, overwrite=False):
        if not isinstance(plugin, AlgorithmPlugin):
            raise TypeError("plugin must be an AlgorithmPlugin instance.")
        plugin.validate()
        if plugin.name in self._algorithms and not overwrite:
            raise ValueError(f"Algorithm '{plugin.name}' already registered. Use overwrite=True.")
        self._algorithms[plugin.name] = plugin
        print(f"  [PluginManager] ✅ Algorithm registered: {plugin.name}")

    def get_gate(self, name):
        if name not in self._gates:
            raise KeyError(f"Gate '{name}' not found. Available: {list(self._gates.keys())}")
        return self._gates[name]

    def get_backend(self, name):
        if name not in self._backends:
            raise KeyError(f"Backend '{name}' not found. Available: {list(self._backends.keys())}")
        return self._backends[name]

    def get_algorithm(self, name):
        if name not in self._algorithms:
            raise KeyError(f"Algorithm '{name}' not found. Available: {list(self._algorithms.keys())}")
        return self._algorithms[name]

    def unregister(self, plugin_type, name):
        registry = self._registry_for(plugin_type)
        if name not in registry:
            raise KeyError(f"{plugin_type.capitalize()} '{name}' not found.")
        del registry[name]
        gc.collect()
        print(f"  [PluginManager] Unregistered {plugin_type}: {name}")

    def list_plugins(self):
        print("=" * 56)
        print("  PKTron Plugin Registry")
        print("=" * 56)
        for title, registry in [("Gates", self._gates), ("Backends", self._backends), ("Algorithms", self._algorithms)]:
            print(f"  -- {title} ({len(registry)}) --")
            for name, plugin in registry.items():
                print(f"     * {name:<22} {type(plugin).__name__}")
            if not registry:
                print("     (none)")
        print("=" * 56)

    def counts(self):
        return {
            "gates"     : len(self._gates),
            "backends"  : len(self._backends),
            "algorithms": len(self._algorithms),
        }

    def clear(self, plugin_type="all"):
        if plugin_type == "all":
            self._gates.clear()
            self._backends.clear()
            self._algorithms.clear()
        else:
            self._registry_for(plugin_type).clear()
        gc.collect()
        print(f"  [PluginManager] Cleared: {plugin_type}")

    def _registry_for(self, plugin_type):
        mapping = {"gate": self._gates, "backend": self._backends, "algorithm": self._algorithms}
        if plugin_type not in mapping:
            raise ValueError(f"Unknown plugin type '{plugin_type}'. Choose: gate | backend | algorithm")
        return mapping[plugin_type]

    def __repr__(self):
        c = self.counts()
        return f"PluginManager(gates={c['gates']}, backends={c['backends']}, algorithms={c['algorithms']})"


plugin_manager = PluginManager()
