import collections.abc
import typing as ta

from ... import check
from ... import dataclasses as dc
from ... import reflect as rfl
from ..api.contexts import UnmarshalContext
from ..api.contexts import UnmarshalFactoryContext
from ..api.types import Unmarshaler
from ..api.types import UnmarshalerFactory
from ..api.values import Value
from .api import ObjectSpecials
from .infos import FieldInfo
from .infos import FieldInfos


##


@dc.dataclass(frozen=True)
class ObjectUnmarshaler(Unmarshaler):
    factory: ta.Callable
    fields_by_unmarshal_name: ta.Mapping[str, tuple[FieldInfo, Unmarshaler]]

    _: dc.KW_ONLY

    specials: ObjectSpecials = ObjectSpecials()

    defaults: ta.Mapping[str, ta.Any] | None = None

    embeds: ta.Mapping[str, type] | None = None
    embeds_by_unmarshal_name: ta.Mapping[str, tuple[str, str]] | None = None

    ignore_unknown: bool = False

    unwrap_if_single_field: FieldInfo | None = None

    @classmethod
    def make(
            cls,
            ctx: UnmarshalFactoryContext,
            factory: ta.Callable,
            fis: FieldInfos,
            **kwargs: ta.Any,
    ) -> Unmarshaler:
        fields_by_unmarshal_name = {
            n: (fi, ctx.make_unmarshaler(fi.type))
            for fi in fis
            for n in fi.unmarshal_names
        }

        defaults = {
            fi.name: dfl.must()
            for fi in fis
            if (dfl := fi.options.default) is not None and dfl.present
        }

        return cls(
            factory,
            fields_by_unmarshal_name,
            defaults=defaults,
            **kwargs,
        )

    #

    def unmarshal(self, ctx: UnmarshalContext, v: Value) -> ta.Any:
        ma: collections.abc.Mapping
        if isinstance(v, collections.abc.Mapping):
            ma = v

        elif (usf := self.unwrap_if_single_field) is not None:
            ma = {usf.unmarshal_names[0]: v}

        else:
            raise TypeError(v)

        u: ta.Any
        kw: dict[str, ta.Any] = {}
        ukf: dict[str, ta.Any] | None = None

        ekws: dict[str, dict[str, ta.Any]] = {en: {} for en in self.embeds or ()}

        if self.specials.source is not None:
            kw[self.specials.source] = v

        if self.specials.unknown is not None:
            kw[self.specials.unknown] = ukf = {}

        for k, mv in ma.items():
            ks = check.isinstance(k, str)

            try:
                fi, u = self.fields_by_unmarshal_name[ks]

            except KeyError:
                if ukf is not None:
                    ukf[ks] = mv  # FIXME: unmarshal?
                    continue

                if self.ignore_unknown:
                    continue

                raise

            if self.embeds_by_unmarshal_name and (en := self.embeds_by_unmarshal_name.get(ks)):
                tkw, tk = ekws[en[0]], en[1]
            else:
                tkw, tk = kw, fi.name

            if tk in tkw:
                raise KeyError(f'Duplicate keys for field {tk!r}: {ks!r}')

            tkw[tk] = u.unmarshal(ctx, mv)

        for em, ecls in self.embeds.items() if self.embeds else ():
            ekw = ekws[em]
            ev = ecls(**ekw)
            kw[em] = ev

        if self.defaults:
            for dk, dv in self.defaults.items():
                kw.setdefault(dk, dv)

        return self.factory(**kw)


##


@dc.dataclass(frozen=True)
class SimpleObjectUnmarshalerFactory(UnmarshalerFactory):
    dct: ta.Mapping[type, ta.Sequence[FieldInfo]]

    _: dc.KW_ONLY

    specials: ObjectSpecials = ObjectSpecials()

    def make_unmarshaler(self, ctx: UnmarshalFactoryContext, rty: rfl.Type) -> ta.Callable[[], Unmarshaler] | None:
        if not (isinstance(rty, type) and rty in self.dct):
            return None

        def inner() -> Unmarshaler:
            ty = check.isinstance(rty, type)

            fis = FieldInfos(self.dct[ty])

            return ObjectUnmarshaler.make(
                ctx,
                ty,
                fis,
                specials=self.specials,
            )

        return inner
