import os
from configparser import ConfigParser


class Lsc:
    def __init__(self, modfactory, ls_uri, ls_ioname='lsio', ls_devname='ls', ls_model='336', **kwds):
        self.modfactory = Mod = modfactory
        self.model = ls_model
        self.ioname = ls_ioname
        self.devname = ls_devname
        self.io = Mod(self.ioname, cls=f'frappy_psi.lakeshore.IO{self.model}',
                      description='comm. to lakeshore in cc rack',
                      uri=ls_uri)
        self.dev = Mod(self.devname, cls=f'frappy_psi.lakeshore.Device{self.model}',
                       description='lakeshore in cc rack', io=self.ioname, curve_handling=True)
        self.loops = {}
        self.outputs = {}

    def sensor(self, name, channel, calcurve, **kwds):
        Mod = self.modfactory
        kwds.setdefault('cls', f'frappy_psi.lakeshore.Sensor{self.model}')
        kwds.setdefault('description', f'T sensor {name}')
        return Mod(name, channel=channel, calcurve=calcurve,
                   io=self.ioname, device=self.devname, **kwds)

    def loop(self, name, channel, calcurve, **kwds):
        Mod = self.modfactory
        kwds.setdefault('cls', f'frappy_psi.lakeshore.Loop{self.model}')
        kwds.setdefault('description', f'T loop {name}')
        mod = Mod(name, channel=channel, calcurve=calcurve,
                  io=self.ioname, device=self.devname, **kwds)
        self.loops[name] = mod
        return mod

    def heater(self, name, max_heater, resistance, output_no=1, **kwds):
        Mod = self.modfactory
        if output_no == 1:
            kwds.setdefault('cls', f'frappy_psi.lakeshore.MainOutput{self.model}')
        elif output_no == 2:
            kwds.setdefault('cls', f'frappy_psi.lakeshore.SecondaryOutput{self.model}')
        else:
            return
        kwds.setdefault('description', '')
        mod = Mod(name, max_heater=max_heater, resistance=resistance,
                  io=self.ioname, device=self.devname, **kwds)
        self.outputs[name] = mod
        return mod

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        outmodules = dict(self.outputs)
        for name, loop in self.loops.items():
            outname = loop.get('output_module')
            if outname:
                out = outmodules.pop(outname, None)
                if not out:
                    raise KeyError(f'{outname} is not a output module in this lakeshore')
            else:
                if not outmodules:
                    raise KeyError(f'{name} needs an output module on this lakeshore')
                outname = list(outmodules)[0]
                out = outmodules.pop(outname)
                loop['output_module'] = outname
                if not out['description']:
                    out['description'] = f'heater for {outname}'


class CCU:
    def __init__(self, Mod, ccu_uri, ccu_ioname='ccuio', ccu_devname='ccu', he=None, n2=None, **kwds):
        self.ioname = ccu_ioname
        self.devname = ccu_devname
        Mod(self.ioname, 'frappy_psi.ccu4.CCU4IO',
            'comm. to CCU4', uri=ccu_uri)
        if he:
            if not isinstance(he, str):  # e.g. True
                he = 'He_lev'
            Mod(he, cls='frappy_psi.ccu4.HeLevel',
                description='the He Level', io=self.ioname)
        if n2:
            if isinstance(n2, str):
                n2 = n2.split(',')
            else:  # e.g. True
                n2 = []
            n2, valve, upper, lower = n2 + ['N2_lev', 'N2_valve', 'N2_upper', 'N2_lower'][len(n2):]
            print(n2, valve, upper, lower)
            Mod(n2, cls='frappy_psi.ccu4.N2Level',
                description='the N2 Level', io=self.ioname,
                valve=valve, upper=upper, lower=lower)
            Mod(valve, cls='frappy_psi.ccu4.N2FillValve',
                description='LN2 fill valve', io=self.ioname)
            Mod(upper, cls='frappy_psi.ccu4.N2TempSensor',
                description='upper LN2 sensor')
            Mod(lower, cls='frappy_psi.ccu4.N2TempSensor',
                description='lower LN2 sensor')


class HePump:
    def __init__(self, Mod, hepump_uri, hepump_io='hepump_io', hemotname='hepump_mot', **kwds):
        Mod(hepump_io, 'frappy_psi.trinamic.BytesIO', 'He pump connection', uri=hepump_uri)
        Mod(hemotname, 'frappy_psi.trinamic.Motor', 'He pump valve motor', io=hepump_io)


class Rack:
    rackfile = '/home/l_samenv/.config/racks.ini'

    def __init__(self, modfactory, **kwds):
        self.modfactory = modfactory
        parser = ConfigParser()
        parser.optionxform = str
        parser.read([self.rackfile])
        kwds.update(parser.items(os.environ['Instrument']))
        self.kwds = kwds

    def lakeshore(self):
        return Lsc(self.modfactory, **self.kwds)

    def ccu(self, **kwds):
        kwds.update(self.kwds)
        return CCU(self.modfactory, **kwds)

    def hepump(self):
        return HePump(self.modfactory, **self.kwds)