import etcd3
import json
from dataclasses import dataclass

@dataclass(init=False)
class EtcdEntry:
    def __init__(self, meta, value, value_in_json=False):
        self.key = meta.key.decode("utf-8")

        self.value = value.decode("utf-8")
        
        if value_in_json:
            self.value = json.loads(self.value)

    key: str
    value: str

class Etcd3Wrapper(object):
    def __init__(self, *args, **kwargs):
        self.client = etcd3.client(*args, **kwargs)
    
    def get(self, *args, value_in_json=False, **kwargs):   
        _value, _key = self.client.get(*args, **kwargs)
        if _key is None:
            return None
        return EtcdEntry(_key, _value, value_in_json=value_in_json)
    
    def put(self, *args, value_in_json=False, **kwargs):
        _key, _value = args
        if value_in_json:
            _value = json.dumps(_value)
        return self.client.put(_key.decode("utf-8"), _value)


    def get_prefix(self, *args, value_in_json=False, **kwargs):
        r = self.client.get_prefix(*args, **kwargs)
        for entry in r:
            yield EtcdEntry(*entry[::-1], value_in_json=value_in_json)

    def watch_prefix(self, *args, value_in_json=False, **kwargs):
        r, _ = self.client.watch_prefix(*args, **kwargs)
        for event in r:
            yield EtcdEntry(event, event.value, value_in_json=value_in_json)