from functools import wraps
from types import SimpleNamespace
from epics import PV


class Source(PV):
    """
    PV, but waits for connection on creation
    """

    def __init__(self, pvname, *args, **kwargs):
        self.name = pvname
        super().__init__(pvname, *args, **kwargs)
        self.wait_for_connection()



class PVCollection:
    """
    Creates all PVs, then waits for connection
    Disconnects all PVs
    """

    def __init__(self, name, **kwargs):
        self.name = name
#        self.pvnames = SimpleNamespace(**kwargs)
        self.pv_dict = pv_dict = {k: PV(v) for k, v in kwargs.items()}
        for pv in pv_dict.values():
            pv.wait_for_connection()
        self.pvs = SimpleNamespace(**pv_dict)

    def __str__(self):
        return self.name

    def disconnect(self):
        for pv in self.pv_dict.values(): #TODO is it better to use: pv_dict -> pvs.__dict__
            pv.disconnect()



class Camera(PVCollection):
    """
    Expects: pvname = "NAME:FPICTURE"
    Also creates NAME:WIDTH and NAME:HEIGHT
    """

    #TODO: could also accept base as argument, and also append ":FPICTURE"

    def __init__(self, pvname):
        assert pvname.endswith(":FPICTURE")
        base = pvname.rsplit(":", 1)[0]
        super().__init__(
            base,
            image = pvname,
            width = base + ":WIDTH",
            height = base + ":HEIGHT"
        )
        self.pvs.image.auto_monitor = True

    def get(self):
        data = self.pvs.image.get()
        data.shape = self.get_shape()
        return data

    def get_shape(self):
        width = self.pvs.width.get()
        height = self.pvs.height.get()
        width = int(round(width))
        height = int(round(height))
        shape = (height, width)
        return shape

    def add_callback(self, callback, **kwargs):
        @wraps(callback)
        def wrapper(*args, value=None, **kwargs):
            value.shape = self.get_shape()
            return callback(*args, value=value, **kwargs)
        return self.pvs.image.add_callback(callback=wrapper, with_ctrlvars=False, **kwargs)