import dataclasses
import json
from typing import Any, List, Type, TypeVar, overload

import pydantic
from pydantic.main import BaseModel

from chalk.features import DataFrame, Feature, Filter, serialize_dtype
from chalk.features.feature_field import HasOnePathObj
from chalk.features.resolver import Cron, OfflineResolver, OnlineResolver, SinkResolver, StreamResolver
from chalk.parsed.duplicate_input_gql import (
    UpsertDataFrameGQL,
    UpsertFeatureGQL,
    UpsertFeatureIdGQL,
    UpsertFeatureReferenceGQL,
    UpsertFeatureTimeKindGQL,
    UpsertFilterGQL,
    UpsertHasManyKindGQL,
    UpsertHasOneKindGQL,
    UpsertReferencePathComponentGQL,
    UpsertResolverGQL,
    UpsertResolverInputUnionGQL,
    UpsertResolverOutputGQL,
    UpsertScalarKindGQL,
    UpsertSinkResolverGQL,
    UpsertStreamResolverGQL,
    UpsertStreamResolverParamGQL,
    UpsertStreamResolverParamKeyedStateGQL,
    UpsertStreamResolverParamMessageGQL,
    VersionInfoGQL,
)
from chalk.streams import Windowed
from chalk.streams.types import (
    StreamResolverParam,
    StreamResolverParamKeyedState,
    StreamResolverParamMessage,
    StreamResolverParamMessageWindow,
)
from chalk.utils import paths
from chalk.utils.duration import timedelta_to_duration
from chalk.utils.string import to_snake_case

T = TypeVar("T")

try:
    import attrs
except ImportError:
    attrs = None


@dataclasses.dataclass
class _ConvertedType:
    name: str
    bases: List[str]


def _get_qualified_class_name(cls: Type[Any]):
    mod = cls.__module__
    return cls.__qualname__ if mod == "builtins" else f"{mod}.{cls.__qualname__}"


_primitive_types = {bool, str, int, float}


def _convert_type(underlying_type: Type) -> _ConvertedType:
    if isinstance(underlying_type, Windowed):
        if underlying_type._kind in _primitive_types:
            return _ConvertedType(name=f"Windowed[{str(underlying_type._kind.__name__)}]", bases=["CHALK_WINDOWED"])
        else:
            kind = str(underlying_type._kind).replace("typing.", "")
            return _ConvertedType(name=f"Windowed[{kind}]", bases=["CHALK_WINDOWED"])
    assert isinstance(underlying_type, type), f"Underlying type is not a type: '{str(underlying_type)}'"
    base_classes = [x.__name__ for x in type.mro(underlying_type)]
    if "Enum" in base_classes:
        # add type(s) of Enum to base_classes
        # Enums are already confirmed to be of one type by here
        value_type = [type(t).__name__ for t in underlying_type._value2member_map_.keys()][0]
        base_classes.insert(base_classes.index("Enum") + 1, value_type)

    # Attrs and Dataclasses don't technically have base classes
    # Pydantic calls their base class BaseModel which is way too generic for string comparison
    # For simplicity on the server-side validation, we'll come up with our own "base class" names
    if dataclasses.is_dataclass(underlying_type):
        base_classes.append("__dataclass__")
    if attrs is not None and isinstance(underlying_type, type) and attrs.has(underlying_type):
        base_classes.append("__attrs__")
    if issubclass(underlying_type, pydantic.BaseModel):
        base_classes.append("__pydantic__")

    return _ConvertedType(
        name=underlying_type.__name__,
        bases=base_classes,
    )


def _get_feature_id(s: Feature):
    assert s.features_cls is not None
    assert s.name is not None, "Could not find feature name"
    assert s.namespace is not None, "Could not find feature namespace"
    return UpsertFeatureIdGQL(
        fqn=s.fqn,
        name=s.name,
        namespace=s.namespace,
        className=s.features_cls.__name__,
        attributeName=s.attribute_name,
        explicitNamespace=s.namespace != to_snake_case(s.features_cls.__name__),
    )


def _convert_df(df: Type[DataFrame]) -> UpsertDataFrameGQL:
    return UpsertDataFrameGQL(
        columns=[_get_feature_id(f) for f in df.columns if not f.is_autogenerated],
        filters=[convert_type_to_gql(f) for f in df.filters],
    )


def _get_path_component(pc: HasOnePathObj) -> UpsertReferencePathComponentGQL:
    assert isinstance(pc.parent, Feature), f"Parent in relationship path not a feature, but {type(pc).__name__}"
    return UpsertReferencePathComponentGQL(
        parent=_get_feature_id(pc.parent),
        child=_get_feature_id(pc.child),
        parentToChildAttributeName=pc.parent_to_child_attribute_name,
    )


@overload
def convert_type_to_gql(t: Filter) -> UpsertFilterGQL:
    ...


@overload
def convert_type_to_gql(t: StreamResolver) -> UpsertStreamResolverGQL:
    ...


@overload
def convert_type_to_gql(t: SinkResolver) -> UpsertSinkResolverGQL:
    ...


@overload
def convert_type_to_gql(t: OnlineResolver) -> UpsertResolverGQL:
    ...


@overload
def convert_type_to_gql(t: OfflineResolver) -> UpsertResolverGQL:
    ...


@overload
def convert_type_to_gql(t: StreamResolverParam) -> UpsertStreamResolverParamGQL:
    ...


@overload
def convert_type_to_gql(t: Feature) -> UpsertFeatureGQL:
    ...


def convert_type_to_gql(t: Any):
    if isinstance(t, StreamResolver):
        return UpsertStreamResolverGQL(
            fqn=t.fqn,
            kind="stream",
            sourceClassName=_get_qualified_class_name(t.source.__class__),
            sourceConfig=t.source._config_to_json(),
            functionDefinition=t.function_definition,
            environment=[t.environment] if isinstance(t.environment, str) else t.environment,
            doc=t.fn.__doc__,
            machineType=t.machine_type,
            output=[
                _get_feature_id(f)
                for f in t.output_features
                if (not isinstance(f, type) or not issubclass(f, DataFrame)) and not f.is_autogenerated
            ],
            inputs=[
                convert_type_to_gql(i)
                for i in t.signature.params
                if not isinstance(i, StreamResolverParamMessageWindow)
            ],
            owner=t.owner,
        )

    if isinstance(t, SinkResolver):
        return UpsertSinkResolverGQL(
            fqn=t.fqn,
            functionDefinition=t.function_definition,
            inputs=[
                UpsertFeatureReferenceGQL(
                    underlying=_get_feature_id(f),
                    path=[_get_path_component(p) for p in f.path or []],
                )
                for f in t.inputs
                if not f.is_autogenerated
            ],
            environment=t.environment,
            tags=t.tags,
            doc=t.doc,
            machineType=t.machine_type,
            bufferSize=t.buffer_size,
            debounce=None if t.debounce is None else timedelta_to_duration(t.debounce),
            maxDelay=None if t.max_delay is None else timedelta_to_duration(t.max_delay),
            upsert=t.upsert,
            owner=t.owner,
        )

    if isinstance(t, (OnlineResolver, OfflineResolver)):
        cron = t.cron
        if isinstance(cron, Cron):
            assert cron.schedule is not None, "`Cron` object must be constructed with a `schedule` property."
            cron = cron.schedule
        inputs = []
        all_inputs = []
        for f in t.inputs:
            if (not isinstance(f, type) or not issubclass(f, DataFrame)) and not f.is_autogenerated:
                feature_ref = UpsertFeatureReferenceGQL(
                    underlying=_get_feature_id(f),
                    path=[_get_path_component(p) for p in f.path or []],
                )
                all_inputs.append(UpsertResolverInputUnionGQL(feature=feature_ref))
                inputs.append(feature_ref)
            elif isinstance(f, type) and issubclass(f, DataFrame):
                all_inputs.append(UpsertResolverInputUnionGQL(dataframe=_convert_df(f)))

        assert t.output is not None
        return UpsertResolverGQL(
            fqn=t.fqn,
            kind="offline" if isinstance(t, OfflineResolver) else "online",
            functionDefinition=t.function_definition,
            inputs=inputs,
            all_inputs=all_inputs,
            output=UpsertResolverOutputGQL(
                features=[
                    _get_feature_id(f)
                    for f in t.output.features
                    if (not isinstance(f, type) or not issubclass(f, DataFrame)) and not f.is_autogenerated
                ],
                dataframes=[
                    _convert_df(f) for f in t.output.features if isinstance(f, type) and issubclass(f, DataFrame)
                ],
            ),
            environment=t.environment,
            tags=t.tags,
            doc=t.doc,
            cron=cron,
            machineType=t.machine_type,
            owner=t.owner,
            timeout=None if t.timeout is None else timedelta_to_duration(t.timeout),
        )

    if isinstance(t, StreamResolverParam):
        if isinstance(t, StreamResolverParamMessage):
            converted_type = _convert_type(t.typ)
            schema = None
            if issubclass(t.typ, BaseModel):
                schema = json.loads(t.typ.schema_json())
            return UpsertStreamResolverParamGQL(
                message=UpsertStreamResolverParamMessageGQL(
                    name=t.name,
                    typeName=converted_type.name,
                    bases=converted_type.bases,
                    schema=schema,
                ),
                state=None,
            )
        elif isinstance(t, StreamResolverParamKeyedState):
            converted_type = _convert_type(t.typ)
            default_value = None
            schema = None
            if dataclasses.is_dataclass(t.default_value):
                default_value = json.loads(json.dumps(dataclasses.asdict(t.default_value)))
                schema = pydantic.dataclasses.dataclass(t.typ).__pydantic_model__.schema()
            elif isinstance(t.default_value, BaseModel):
                assert issubclass(t.typ, BaseModel)
                default_value = json.loads(t.default_value.json())
                schema = json.loads(t.typ.schema_json())
            return UpsertStreamResolverParamGQL(
                state=UpsertStreamResolverParamKeyedStateGQL(
                    name=t.name,
                    defaultValue=default_value,
                    typeName=converted_type.name,
                    bases=converted_type.bases,
                    schema=schema,
                ),
                message=None,
            )

    if isinstance(t, Feature):
        assert t.name is not None, "Feature has no name"
        assert t.namespace is not None, "Feature has no namespace"
        assert t.features_cls is not None
        scalar_kind_gql = None
        has_one_kind_gql = None
        has_many_kind_gql = None
        feature_time_kind_gql = None
        if t.is_autogenerated:
            raise RuntimeError("Autogenerated features should not be converted")
        if t.is_has_one:
            assert t.join is not None
            has_one_kind_gql = UpsertHasOneKindGQL(join=convert_type_to_gql(t.join))
        elif t.is_has_many:
            assert t.join is not None
            has_many_kind_gql = UpsertHasManyKindGQL(
                join=convert_type_to_gql(t.join),
                columns=None,
                filters=None,
            )
        elif t.is_feature_time:
            feature_time_kind_gql = UpsertFeatureTimeKindGQL()
        else:
            assert t.typ is not None
            converted_type = _convert_type(t.typ.underlying)
            scalar_kind_gql = UpsertScalarKindGQL(
                scalarKind=converted_type.name,
                primary=t.primary,
                baseClasses=converted_type.bases,
                version=t.version and t.version.version,
                versionInfo=t.version
                and VersionInfoGQL(
                    version=t.version.version,
                    maximum=t.version.maximum,
                    default=t.version.default,
                    versions=[f.fqn for f in t.version.reference.values()],
                ),
                dtype=serialize_dtype(t.converter.pyarrow_dtype),
            )
        return UpsertFeatureGQL(
            id=UpsertFeatureIdGQL(
                fqn=t.fqn,
                name=t.name,
                namespace=t.namespace,
                className=t.features_cls.__name__,
                attributeName=t.attribute_name,
                explicitNamespace=t.namespace != to_snake_case(t.features_cls.__name__),
            ),
            maxStaleness=timedelta_to_duration(t.max_staleness),
            description=t.description,
            owner=t.owner,
            windowBuckets=None if t.window_buckets is None else sorted(list(t.window_buckets)),
            etlOfflineToOnline=t.etl_offline_to_online,
            tags=t.tags,
            hasOneKind=has_one_kind_gql,
            hasManyKind=has_many_kind_gql,
            scalarKind=scalar_kind_gql,
            featureTimeKind=feature_time_kind_gql,
            namespacePath=str(paths.get_classpath(t.features_cls)),
        )

    if isinstance(t, Filter):
        if isinstance(t.lhs, Feature):
            assert (
                t.lhs.is_scalar or t.lhs.is_feature_time
            ), f"Filters must be on scalar features, but {t.lhs.fqn} is not a scalar."

        if isinstance(t.rhs, Feature):
            assert (
                t.rhs.is_scalar or t.rhs.is_feature_time
            ), f"Filters must be on scalar features, but {t.rhs.fqn} is not a scalar."

        return UpsertFilterGQL(
            lhs=_get_feature_id(t.lhs),
            op=t.operation,
            rhs=_get_feature_id(t.rhs),
        )

    raise ValueError(f"Unable to convert {t} to GQL")
