from __future__ import annotations

import collections
import dataclasses
import json
from datetime import timedelta
from typing import TYPE_CHECKING, Any, List, Optional, Type, TypeVar, Union, overload

from chalk import Resolver, Windowed
from chalk.features import DataFrame, Feature, Features, Filter, serialize_dtype
from chalk.features._encoding.outputs import find_feature_or_resolver_by_fqn
from chalk.features.feature_field import HasOnePathObj
from chalk.features.feature_set import is_feature_set_class
from chalk.features.pseudofeatures import PSEUDONAMESPACE
from chalk.features.resolver import Cron, OfflineResolver, OnlineResolver, SinkResolver, StreamResolver
from chalk.features.scheduled_query import ScheduledQuery
from chalk.parsed.duplicate_input_gql import (
    FeatureClassGQL,
    RecomputeFeaturesGQL,
    UpsertCDCSourceGQL,
    UpsertCronQueryGQL,
    UpsertDataFrameGQL,
    UpsertFeatureGQL,
    UpsertFeatureIdGQL,
    UpsertFeatureReferenceGQL,
    UpsertFeatureTimeKindGQL,
    UpsertFilterGQL,
    UpsertHasManyKindGQL,
    UpsertHasOneKindGQL,
    UpsertReferencePathComponentGQL,
    UpsertResolverGQL,
    UpsertResolverInputUnionGQL,
    UpsertResolverOutputGQL,
    UpsertScalarKindGQL,
    UpsertSinkResolverGQL,
    UpsertStreamResolverGQL,
    UpsertStreamResolverParamGQL,
    UpsertStreamResolverParamKeyedStateGQL,
    UpsertStreamResolverParamMessageGQL,
    VersionInfoGQL,
)
from chalk.sql._internal.sql_source import BaseSQLSource, TableIngestMixIn
from chalk.streams.types import (
    StreamResolverParam,
    StreamResolverParamKeyedState,
    StreamResolverParamMessage,
    StreamResolverParamMessageWindow,
)
from chalk.utils import paths
from chalk.utils.collections import get_unique_item
from chalk.utils.duration import CronTab, Duration, timedelta_to_duration
from chalk.utils.pydanticutil.pydantic_compat import is_pydantic_basemodel, is_pydantic_basemodel_instance
from chalk.utils.string import removeprefix, to_snake_case

if TYPE_CHECKING:
    import pydantic
else:
    try:
        import pydantic.v1 as pydantic
    except ImportError:
        import pydantic

T = TypeVar("T")


@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__}"


def _convert_type(underlying_type: Type, is_windowed: bool) -> _ConvertedType:
    if isinstance(underlying_type, Windowed):
        is_windowed = True
        underlying_type = underlying_type.kind
    bases = []
    pretty_name = repr(underlying_type)
    if isinstance(underlying_type, type):
        bases.extend(x.__name__ for x in underlying_type.mro())
        pretty_name = underlying_type.__name__
    else:
        bases = ["object"]
    pretty_name = pretty_name.replace("typing.", "")
    if is_windowed:
        bases.append("CHALK_WINDOWED")
        return _ConvertedType(name=f"Windowed[{pretty_name}]", bases=bases)
    if isinstance(underlying_type, Feature):
        return _convert_type(underlying_type.converter.rich_type, underlying_type.is_windowed)

    bases = []
    if isinstance(underlying_type, type):
        bases.extend(x.__name__ for x in underlying_type.mro())
    else:
        bases = ["object"]
    return _ConvertedType(
        name=pretty_name,
        bases=bases,
    )


def _get_feature_id(s: Feature, root_fqn: bool):
    assert s.features_cls is not None
    return UpsertFeatureIdGQL(
        fqn=s.root_fqn if root_fqn else s.fqn,
        name=s.name,
        namespace=s.namespace,
        isPrimary=s.primary,
        className=s.features_cls.__name__,
        attributeName=s.attribute_name,
        explicitNamespace=s.namespace != to_snake_case(s.features_cls.__name__),
    )


def _get_pseudofeature_id(s: Feature):
    return UpsertFeatureIdGQL(
        fqn=s.fqn,
        name=s.name,
        namespace=s.namespace,
        isPrimary=s.primary,
    )


def _convert_df(df: Type[DataFrame]) -> UpsertDataFrameGQL:
    return UpsertDataFrameGQL(
        columns=[
            _get_feature_id(f, root_fqn=True)
            for f in df.columns
            if not f.is_autogenerated and not f.namespace == PSEUDONAMESPACE
        ],
        # filters=[convert_type_to_gql(f) for f in df.filters],
        filters=None,
    )


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, root_fqn=False),
        child=_get_feature_id(pc.child, root_fqn=False),
        parentToChildAttributeName=pc.parent_to_child_attribute_name,
    )


def _parse_generic_params_recursively(typ: Any) -> str:
    """
    Stringify a type with generic parameters like 'list[list[str]]',
    making sure parameters are included recursively.
    :return: the type name along with any parameters it may have
    """
    if not isinstance(typ, type):
        type_name = repr(typ)
    else:
        type_name = typ.__name__

    type_name = removeprefix(type_name, "typing.")
    if not hasattr(typ, "__args__"):
        return type_name
    params = [_parse_generic_params_recursively(x) for x in typ.__args__]
    param_string = ",".join(params)
    return f"{type_name}[{param_string}]"


@overload
def convert_type_to_gql(t: Filter, path_prefix: Optional[str] = None) -> UpsertFilterGQL:
    ...


@overload
def convert_type_to_gql(t: StreamResolver, path_prefix: Optional[str] = None) -> UpsertStreamResolverGQL:
    ...


@overload
def convert_type_to_gql(t: SinkResolver, path_prefix: Optional[str] = None) -> UpsertSinkResolverGQL:
    ...


@overload
def convert_type_to_gql(
    t: Union[OnlineResolver, OfflineResolver], path_prefix: Optional[str] = None
) -> UpsertResolverGQL:
    ...


@overload
def convert_type_to_gql(
    t: Resolver, path_prefix: Optional[str] = None
) -> UpsertStreamResolverGQL | UpsertSinkResolverGQL | UpsertResolverGQL:
    ...


@overload
def convert_type_to_gql(t: StreamResolverParam, path_prefix: Optional[str] = None) -> UpsertStreamResolverParamGQL:
    ...


@overload
def convert_type_to_gql(t: ScheduledQuery, path_prefix: Optional[str] = None) -> UpsertCronQueryGQL:
    ...


@overload
def convert_type_to_gql(t: Feature, path_prefix: Optional[str] = None) -> UpsertFeatureGQL:
    ...


@overload
def convert_type_to_gql(t: Type[Features], path_prefix: Optional[str] = None) -> FeatureClassGQL:
    ...


def convert_type_to_gql(
    t: Any,
    # Resolvers have their filenames attached to them. But for UI presentation, we don't
    # want to show the full path, e.g. /Users/customer-name/projects/chalk/....,
    # but we want to show just the path relative to the project root. To do this efficiently,
    # we'll pass the project root path here and delete it from those full paths
    path_prefix: Optional[str] = None,
):
    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, root_fqn=True)
                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,
            filename=t.filename if path_prefix is None else t.filename.replace(path_prefix, ""),
            sourceLine=t.source_line,
        )

    if isinstance(t, SinkResolver):
        return UpsertSinkResolverGQL(
            fqn=t.fqn,
            functionDefinition=t.function_definition,
            inputs=[
                UpsertFeatureReferenceGQL(
                    underlying=_get_feature_id(f, root_fqn=False),
                    path=[_get_path_component(p) for p in f.path or []],
                )
                for f in t.inputs
                if isinstance(f, Feature) and 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,
            filename=t.filename if path_prefix is None else t.filename.replace(path_prefix, ""),
            sourceLine=t.source_line,
        )

    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:
                if f.namespace == PSEUDONAMESPACE:
                    feature_ref = UpsertFeatureReferenceGQL(
                        underlying=_get_pseudofeature_id(f),
                        path=[_get_path_component(p) for p in f.path or []],
                    )
                    all_inputs.append(UpsertResolverInputUnionGQL(pseudoFeature=feature_ref))
                else:
                    feature_ref = UpsertFeatureReferenceGQL(
                        underlying=_get_feature_id(f, root_fqn=False),
                        path=[_get_path_component(p) for p in f.path or []],
                    )
                    inputs.append(feature_ref)
                    all_inputs.append(UpsertResolverInputUnionGQL(feature=feature_ref))
            elif isinstance(f, type) and issubclass(f, DataFrame):
                all_inputs.append(UpsertResolverInputUnionGQL(dataframe=_convert_df(f)))

        assert t.output is not None
        if t.data_sources is not None:
            data_sources = [source.to_json() for source in t.data_sources]
        else:
            data_sources = None
        return UpsertResolverGQL(
            fqn=t.fqn,
            kind="offline" if isinstance(t, OfflineResolver) else "online",
            functionDefinition=t.function_definition,
            inputs=inputs,
            allInputs=all_inputs,
            output=UpsertResolverOutputGQL(
                features=[
                    _get_feature_id(f, root_fqn=True)
                    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=convert_type_crontab(cron) if cron is not None else cron,
            machineType=t.machine_type,
            owner=t.owner,
            timeout=None if t.timeout is None else timedelta_to_duration(t.timeout),
            filename=t.filename if path_prefix is None else t.filename.replace(path_prefix, ""),
            sourceLine=t.source_line,
            dataSources=data_sources,
        )

    if isinstance(t, StreamResolverParam):
        if isinstance(t, StreamResolverParamMessage):
            converted_type = _convert_type(t.typ, is_windowed=False)
            schema = None
            if is_pydantic_basemodel(t.typ):
                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, is_windowed=False)
            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 is_pydantic_basemodel_instance(t.default_value):
                assert is_pydantic_basemodel(t.typ)
                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, ScheduledQuery):
        for f in t.output:
            validate_feature_reference(f)
        get_unique_item((f.split(".")[0] for f in t.output), " of root output feature namespaces")
        if not isinstance(t.recompute_features, bool):
            for f in t.recompute_features:
                validate_feature_reference(f)

        return UpsertCronQueryGQL(
            name=t.name,
            cron=convert_type_crontab(t.cron),
            filename=t.filename if path_prefix is None else t.filename.replace(path_prefix, ""),
            output=list(t.output),
            maxSamples=t.max_samples,
            recomputeFeatures=RecomputeFeaturesGQL(
                featureFqns=list(t.recompute_features) if not isinstance(t.recompute_features, bool) else None,
                all=t.recompute_features if isinstance(t.recompute_features, bool) else None,
            ),
            lowerBound=t.lower_bound,
            upperBound=t.upper_bound,
            tags=list(t.tags) if t.tags is not None else None,
            requiredResolverTags=list(t.required_resolver_tags) if t.required_resolver_tags is not None else None,
            storeOnline=t.store_online,
            storeOffline=t.store_offline,
            incrementalSources=None if t.incremental_resolvers is None else list(t.incremental_resolvers),
        )

    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:
            if t.join is None:
                builder = t.features_cls.__chalk_error_builder__
                builder.add_diagnostic(
                    range=builder.annotation_range(t.attribute_name),
                    message=(
                        f"The attribute '{t.attribute_name}' on the feature class '{t.features_cls.__name__}' "
                        f"references {t.typ}. However, the has-one relationship is not defined."
                    ),
                    raise_error=TypeError,
                    code_href="https://docs.chalk.ai/docs/has-one",
                    label="Missing has-one join",
                    code="35",
                )
            assert t.join is not None
            has_one_kind_gql = UpsertHasOneKindGQL(join=convert_type_to_gql(t.join))

        elif t.is_has_many:
            if t.join is None:
                builder = t.features_cls.__chalk_error_builder__
                builder.add_diagnostic(
                    range=builder.annotation_range(t.attribute_name),
                    message=(
                        f"The attribute '{t.attribute_name}' on the feature class '{t.features_cls.__name__}' "
                        f"references a DataFrame. However, the has-many relationship is not defined."
                    ),
                    label="Missing has-many join",
                    raise_error=TypeError,
                    code="34",
                    code_href="https://docs.chalk.ai/docs/has-many",
                )
            has_many_kind_gql = UpsertHasManyKindGQL(
                join=convert_type_to_gql(t.join),
                columns=None,
                filters=None,
            )

        elif t.join is not None:
            # If a has_one/has_many has an incorrect type annotation
            builder = t.features_cls.__chalk_error_builder__
            builder.add_diagnostic(
                range=builder.annotation_range(t.attribute_name),
                message=(
                    f"The attribute '{t.attribute_name}' on the feature class '{t.features_cls.__name__}' "
                    f"has a join filter ({t.join}) but its type annotation is not a feature class or "
                    f"DataFrame ({t.typ})."
                ),
                label="Incorrect join type annotation",
                raise_error=TypeError,
                code="34",
                code_href="https://docs.chalk.ai/docs/has-many",
            )

        elif t.is_feature_time:
            feature_time_kind_gql = UpsertFeatureTimeKindGQL()

        else:
            assert t.typ is not None
            try:
                converted_type = _convert_type(t.converter.rich_type, is_windowed=t.is_windowed)
            except Exception as e:
                builder = t.features_cls.__chalk_error_builder__
                if builder:
                    builder.add_diagnostic(
                        message=e.args[0] if len(e.args) > 0 else str(e),
                        range=builder.annotation_range(t.attribute_name) or builder.property_range(t.attribute_name),
                        code="39",
                        label="invalid type",
                        code_href="https://docs.chalk.ai/docs/feature-types",
                    )
                raise e
            pretty_name = repr(t.converter.rich_type)
            if isinstance(t.converter.rich_type, type):
                pretty_name = _parse_generic_params_recursively(t.converter.rich_type)
            pretty_name = pretty_name.replace("typing.", "")
            if t.is_windowed:
                pretty_name = f"Windowed[{pretty_name}]"

            scalar_kind_gql = UpsertScalarKindGQL(
                scalarKind=pretty_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,
                isPrimary=t.primary,
                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=None if t.tags is None else list(t.tags),
            hasOneKind=has_one_kind_gql,
            hasManyKind=has_many_kind_gql,
            scalarKind=scalar_kind_gql,
            featureTimeKind=feature_time_kind_gql,
            namespacePath=paths.get_classpath_or_name(t.features_cls),
            isSingleton=t.is_singleton,
        )

    if is_feature_set_class(t):
        return FeatureClassGQL(
            isSingleton=t.__chalk_is_singleton__,
            doc=t.__doc__,
            name=t.__chalk_namespace__,
            owner=t.__chalk_owner__,
            tags=t.__chalk_tags__,
        )

    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, root_fqn=True),
            op=t.operation,
            rhs=_get_feature_id(t.rhs, root_fqn=True),
        )

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


def convert_type_crontab(cron: Union[CronTab, Duration]) -> str:
    return f"{cron.total_seconds()}s" if isinstance(cron, timedelta) else cron


def validate_feature_reference(f: str):
    path = f.split(".")
    fqn = ".".join(path[: min(2, len(path))])
    if (feature := find_feature_or_resolver_by_fqn(fqn)) is None:
        raise ValueError(f"Feature or resolver '{f}' not found.")
    if isinstance(feature, Resolver):
        raise ValueError(f"'{f}' is a resolver, not a feature.")


def gather_cdc_sources() -> List[UpsertCDCSourceGQL]:
    integration_name_to_tables = collections.defaultdict(list)
    for source in BaseSQLSource.registry:
        if isinstance(source, TableIngestMixIn):
            for schema_dot_table, preferences in source.ingested_tables.items():
                if preferences.cdc is True:
                    assert isinstance(source, BaseSQLSource)
                    integration_name_to_tables[source.name].append(schema_dot_table)
    return [
        UpsertCDCSourceGQL(
            integrationName=integration_name,
            schemaDotTableList=tables,
        )
        for integration_name, tables in integration_name_to_tables.items()
    ]
