"""
Configuration Model.  This translates between the GRPC model and a higher-level
model for clients.
"""
import asyncio
from collections import namedtuple
from enum import Enum
import json
import logging

from google.protobuf.descriptor import FieldDescriptor
from preseem._apiutil import field_is_repeated
from google.protobuf.message_factory import GetMessageClass

from preseem.grpc_client import PreseemGrpcClient
from preseem_protobuf.config import company_configs_pb2
from preseem.pushid import generate_pushid

ConfigType = namedtuple('ConfigType', ())
ConfigOp = Enum('ConfigOp', 'ADD DEL')


def protobuf_to_dict(pb) -> dict:
    """Recursively convert a Protobuf message to a Python dict."""
    result = {}
    for field in pb.DESCRIPTOR.fields:
        value = getattr(pb, field.name)

        if field_is_repeated(field):
            # Handle repeated fields (including maps)
            if field.message_type and field.message_type.GetOptions().map_entry:
                # It's a map field
                result[field.name] = dict(value)
            else:
                # It's a repeated field (list)
                result[field.name] = [
                    protobuf_to_dict(v) if isinstance(v, Message) else v for v in value
                ]
        elif field.type == FieldDescriptor.TYPE_MESSAGE:
            # Handle nested messages
            if value is not None and value.ByteSize():
                result[field.name] = protobuf_to_dict(value)
        else:
            # Handle simple types
            result[field.name] = value

    return result


class ConfigClient(PreseemGrpcClient):

    def __init__(self, api_key, host=None, port=None, beta=None):
        super().__init__(api_key, host
                         or 'config{}.preseem.com'.format('-beta' if beta else ''),
                         port)


class ConfigModel(object):

    def __init__(self,
                 client,
                 company_id=None,
                 all_companies=None,
                 instances_only=None):
        """With a regular company-specific API key the model will act on the
           company of the API key.  With an all-companies API key, pass in
           a company_id to act as a specific company, or set all_companies=True
           to operate on all companies.  All companies operation does not
           subscribe to instances, it only processes add/delete operations
           and it calls back with (company_id, old, new) instead of the regular
           callback which just calls (old, new)."""
        self._client = client
        self.company_id = company_id  # used for multi-company keys
        self._cbk = None  # optional callback for updates
        self.all_companies = all_companies
        self.objs = {}  # public: current config objects
        self._sub_task = None  # Instance subscription
        self._sub_lock = asyncio.Lock()
        self._cfg_syn = None
        self._loaded = asyncio.Condition()
        self.instances_only = instances_only

        self.instances = {}  # current supported instances: {company -> {(type, id)}}
        self._sub_tasks = {}  # per-instance subscription state: (task, future)

        self._cfg_types = {}
        self._type_map = {}
        self._type_nt = {}
        self._type_field_is_list = {}
        self._map_types = {}
        self._nested_types = {}
        for ty, ct in self._client.config_types.items():
            for field, fd in ct.pb_type.DESCRIPTOR.fields_by_name.items():
                if fd.message_type:
                    # This is a nested message type.
                    if fd.message_type.GetOptions().map_entry:  # map type
                        value_field = fd.message_type.fields_by_name["value"]
                        self._map_types[(ty, field)] = GetMessageClass(
                            value_field.message_type
                        ) if value_field.message_type else None
                    self._nested_types[fd] = namedtuple(
                        field, [x.name for x in fd.message_type.fields])
            nt = namedtuple(ct.name, ct.pb_type.DESCRIPTOR.fields_by_name)
            self._cfg_types[ct.name.replace('_', '-')] = nt
            self._type_map[nt] = ty
            self._type_nt[ty] = nt
            self._type_field_is_list[ty] = {
                fd.name: field_is_repeated(fd)
                for fd in ct.pb_type.DESCRIPTOR.fields
            }

    def _pb(self, obj, id_prefix=None):
        """convert a config ojbect to a protobuf"""
        config_type = self._type_map.get(type(obj))
        if config_type is None:
            raise ValueError(f'Set was passed an unknown config class "{type(obj)}"')
        ct = self._client.config_types[config_type]
        d = obj._asdict()
        if not obj.id:
            d['id'] = generate_pushid()
            obj = obj._replace(id=d['id'])
        if id_prefix:
            d['id'] = f'{id_prefix}{d["id"]}'
            obj = obj._replace(id=d['id'])
        if obj.extra:
            d['extra'] = json.dumps(obj.extra)
        map_vars = {}
        for map_var in [v for v in obj._fields if (config_type, v) in self._map_types]:
            map_vars[map_var] = d.pop(map_var)
        pb = ct.pb_type(**d)
        for k, v in map_vars.items():
            if v:
                valtype = self._map_types[(config_type, k)]
                if valtype:
                    for kk, vv in v.items():
                        getattr(pb, k)[kk].CopyFrom(valtype(**vv))
                else:
                    setattr(pb, k, v)
        return pb

    def _pb_to_nt(self, config_type, pb):
        """Convert a config protobuf to a namedtuple object"""

        def pythonize_field(pb, f):
            if f == 'extra':
                return extra
            val = getattr(pb, f, None)
            if field_is_list.get(f) and (config_type, f) not in self._map_types:
                val = list(val)
            elif f in nested_fields:
                if (config_type, f) in self._map_types:
                    mapval = self._map_types.get((config_type, f))
                    if mapval:  # value is a protobuf type
                        val = {x: protobuf_to_dict(y) for x, y in val.items()}
                    else:  # value is a simple type
                        val = {x: y for x, y, in val.items()}
                else:
                    val = {
                        x.name: getattr(val, x.name)
                        for x in nested_fields[f].message_type.fields
                    }

            return val

        nt = self._type_nt.get(config_type)
        if nt is None:
            raise ValueError(f'Unsupported config type {config_type}')
        extra = None
        if pb.extra:
            try:
                extra = json.loads(pb.extra)
            except AttributeError:
                extra = {}
            except json.JSONDecodeError:
                logging.warning('Bad JSON in extra field for %s: "%s"', config_type,
                                pb.extra)
            if not isinstance(extra, dict):
                logging.warning('extra field is not a dict for %s: "%s"', config_type,
                                pb.extra)
        nested_fields = {
            x.name: x
            for x in pb.DESCRIPTOR.fields if x in self._nested_types
        }
        field_is_list = self._type_field_is_list[config_type]
        return nt(*[pythonize_field(pb, f) for f in nt._fields])

    def close(self):
        self._client.close_subscriptions()
        if self._sub_task:
            self._sub_task.cancel()
            self._sub_task = None
        for task, fut in self._sub_tasks.values():
            if fut:
                fut.cancel()
            task.cancel()
        self._sub_tasks.clear()

    def dict2cfg(self, ty, d):
        """Convert a dict to a config object."""
        nt = self._cfg_types.get(ty)
        if nt is None:
            raise RuntimeError('Unknown config type "%s"', ty)
        spec = {x: None for x in nt._fields}
        spec.update(d)
        return nt(**spec)

    def get_type(self, name):
        """Return a generated python type for the named config type."""
        return self._cfg_types.get(name)

    def get_default(self, name):
        """Return an "empty" config object for the named config type."""
        nt = self.get_type(name)
        return self._pb_to_nt(self._type_map[nt],
                              self._pb(nt(*[None for _ in nt._fields])))

    async def delete(self, obj, company_id=None):
        """Delete a configuration object."""
        config_type = self._type_map.get(type(obj))
        if config_type is None:
            raise ValueError(f'Set was passed an unknown config class "{type(obj)}"')
        ct = self._client.config_types[config_type]
        r = await self._client.del_config(config_type, obj.id, company_id
                                          or self.company_id)
        if not r:
            raise RuntimeError('Failed to delete config instance "{}"'.format(obj.id))

    async def list(self):
        """Return a list of config objects."""
        loop = asyncio.get_event_loop()
        fut = loop.create_future()
        result = []

        async def cbk(msg, syn=False):
            if msg.event_type == company_configs_pb2.CompanyConfigEvent.LOADED:
                fut.set_result(True)
            else:
                try:
                    r = await self._client.get_config(msg.config_type, msg.config_id,
                                                      msg.company_id)
                    result.append(self._pb_to_nt(msg.config_type, r))
                except ValueError:
                    pass  # we don't know about this type, skip it

        cor = self._client.sub_company_configs(cbk,
                                               company_id=self.company_id,
                                               subscribe_all=self.all_companies)
        task = loop.create_task(cor)
        try:
            await asyncio.wait_for(fut, timeout=10)
        finally:
            task.cancel()
        return result

    async def set(self, obj, company_id=None, id_prefix=None):
        """Set a configuration object."""
        config_type = self._type_map.get(type(obj))
        pb = self._pb(obj, id_prefix=id_prefix)
        r = await self._client.set_config(config_type, pb, company_id
                                          or self.company_id)
        if not r:
            raise RuntimeException('Failed to set config instance "{}"'.format(obj))
        return obj._replace(id=pb.id)

    async def handle_event(self, msg, syn=False):
        """
        Handle a config instance event from the event stream.
        The end result of this is an up to date set of config objects.
        The first time through, call the update() method to wait for everything
        to be fully loaded.  After that, "objs" always has a consistent state.
        This can also call a callback whenever something changes.
        """
        if syn:
            self._cfg_syn = {}  # Keep track of the initial stream of instances

        if msg.event_type == company_configs_pb2.CompanyConfigEvent.LOADED:
            self.instances = self._cfg_syn or {}
            self._cfg_syn = None
            # wait for any instance subscriptions to be initialized
            pend = [fut for _, fut in self._sub_tasks.values() if fut]
            if pend:
                done, pend = await asyncio.wait(pend)
                [x.result() for x in done]
            if self.instances_only and self._cbk:
                for company_id, instances in self.instances.items():
                    for config_type, config_id in instances:
                        try:
                            await self._cbk(ConfigOp.ADD, company_id, config_type,
                                            config_id)
                        except asyncio.CancelledError:
                            raise
                        except Exception as err:
                            logging.warning('Error with initial callback: %s', err)
            async with self._loaded:
                self._loaded.notify()
            return
        elif msg.event_type == company_configs_pb2.CompanyConfigEvent.UPDATE:
            return  # we don't care about modify events at the company level
        elif msg.config_type not in self._client.config_types:
            return  # we don't know/care about this config type

        loop = asyncio.get_event_loop()
        uid = (msg.company_id, msg.config_type, msg.config_id)
        if msg.event_type == company_configs_pb2.CompanyConfigEvent.SET:
            op = ConfigOp.ADD
            if self._cfg_syn is None:
                instances = self.instances.get(msg.company_id)
                if instances is None:
                    instances = self.instances[msg.company_id] = set()
                instances.add((msg.config_type, msg.config_id))
            else:
                instances = self._cfg_syn.get(msg.company_id)
                if instances is None:
                    instances = self._cfg_syn[msg.company_id] = set()
                instances.add((msg.config_type, msg.config_id))
            if not self.instances_only:
                sub_state = self._sub_tasks.get(uid)
                if sub_state:  # unexpected, race condition?
                    pass
                else:
                    # subscribe to the instance
                    task = loop.create_task(
                        self._client.sub_config(msg.config_type, msg.config_id,
                                                self.handle_instance_update,
                                                msg.company_id))
                    sub_state = self._sub_tasks[uid] = (task, loop.create_future())

        elif msg.event_type == company_configs_pb2.CompanyConfigEvent.DELETE:
            op = ConfigOp.DEL
            instances = self.instances.get(msg.company_id)
            if instances:
                instances.discard((msg.config_type, msg.config_id))
                if not instances:
                    del self.instances[msg.company_id]
            if not self.instances_only:
                # stop subscription and callback for instance removal
                sub_state = self._sub_tasks.get(uid)
                if sub_state:
                    task, fut = sub_state
                    if fut:
                        fut.cancel()
                    task.cancel()
                    del self._sub_tasks[uid]

                cur = self.objs.get(msg.config_id)
                if cur:
                    del self.objs[msg.config_id]
                    if self._cbk:  # send a deletion callback
                        try:
                            await self._cbk(cur, None)
                        except asyncio.CancelledError:
                            raise
                        except Exception as err:
                            logging.warning('Error deleting config "%s": %s',
                                            msg.config_id, err)

        if self.instances_only and self._cbk and self._cfg_syn is None:
            # Call a special callback for instance operations only
            try:
                await self._cbk(op, msg.company_id, msg.config_type, msg.config_id)
            except asyncio.CancelledError:
                raise
            except Exception as err:
                logging.warning('Error calling instances callback: %s', err)

    async def handle_instance_update(self, company_id, config_type, msg):
        """Handle a change to a config model object."""
        cur = self.objs.get(msg.id)
        obj = self._pb_to_nt(config_type, msg)
        if cur == obj:
            logging.info('Ignoring config update of "%s" with no change', msg.id)
            return
        self.objs[msg.id] = obj
        if self._cbk:
            try:
                await self._cbk(cur, obj)
            except Exception as err:
                logging.warning('Error updating "%s": %s', msg.id, err)
        # notify the initial company subscription if it is waiting
        sub_state = self._sub_tasks.get((company_id, config_type, msg.id))
        if sub_state and sub_state[1] and not sub_state[1].done():
            sub_state[1].set_result(True)

    async def listen(self, cbk):
        """Listen for events on config objects (in the background)."""
        if self._sub_task is None:
            async with self._sub_lock:
                if self._sub_task is None:
                    self._cbk = cbk
                    self._sub_task = asyncio.get_event_loop().create_task(
                        self._client.sub_company_configs(
                            self.handle_event,
                            company_id=self.company_id,
                            subscribe_all=self.all_companies))
                    async with self._loaded:
                        await self._loaded.wait()
