import dataclasses
import datetime as dt
import enum
import sys
import traceback
from typing import Any, List, Optional, Set, Tuple, Type, Union

import dataclasses_json

from chalk import Resolver
from chalk.features import Feature, Features, is_features_cls
from chalk.parsed.branch_state_rich import FeatureFieldRichSummary, ResolverRichSummary
from chalk.parsed.duplicate_input_gql import FailedImport, UpsertFeatureGQL, UpsertGraphGQL, UpsertResolverGQL
from chalk.parsed.json_conversions import convert_type_to_gql
from chalk.utils.collections import get_unique_item
from chalk.utils.log_with_context import get_logger

_logger = get_logger(__name__)


class EntityUpdateSource(enum.Enum):
    live_update = "live_update"
    deployment_source = "deployment_source"


class EntityUpdateType(enum.Enum):
    resolver = "resolver"
    feature_class = "feature_class"


# Need this to convert datetimes via ISO instead of numberic timestamp -- deser wasn't working otherwise.
DATACLASSES_JSON_TIMESTAMP_FIELD = dataclasses.field(
    metadata=dataclasses_json.config(
        encoder=dt.datetime.isoformat,
        decoder=dt.datetime.fromisoformat,
    )
)


@dataclasses_json.dataclass_json
@dataclasses.dataclass
class BlobLoadError:
    traceback: str
    message: str


@dataclasses_json.dataclass_json
@dataclasses.dataclass
class EntityUpdate:
    """
    Represents a record of loading a pickled resolver/feature update on the branch server.
    Either produces a resolver/feature that gets added to the graph, or an error message.
    """

    blob_id: str
    source: EntityUpdateSource
    success: bool  # True if this entity was unpickled & added to the graph w/o errors
    # TODO (rkargon) This currently shows the time the blob was loaded by the current active worker - not the time it was created/sent.
    timestamp: dt.datetime = DATACLASSES_JSON_TIMESTAMP_FIELD

    # These might be None depending on whether an error was raised when updating the graph
    name: Optional[str] = None  # either resolver.fqn or feature_class.name
    entity_type: Optional[EntityUpdateType] = None
    resolver_gql: Optional[UpsertResolverGQL] = None
    feature_fields_gql: Optional[List[UpsertFeatureGQL]] = None
    error: Optional[BlobLoadError] = None

    @classmethod
    def from_obj(cls, blob_id: str, obj: Union[Resolver, Type[Features]]):
        kwargs: dict[str, Any] = {
            "blob_id": blob_id,
            "timestamp": dt.datetime.now(tz=dt.timezone.utc),
            "source": EntityUpdateSource.live_update,
            "success": True,
        }
        gql = None

        if isinstance(obj, Resolver):
            try:
                gql = convert_type_to_gql(obj)
            except Exception:
                _logger.error(f"Failed to convert {obj=} with {blob_id=} to GQL.", exc_info=True)
            kwargs.update(
                {
                    "name": obj.fqn,
                    "entity_type": EntityUpdateType.resolver,
                    "resolver_gql": gql,
                    "feature_fields_gql": None,
                }
            )
        elif is_features_cls(obj):
            try:
                gql = [convert_type_to_gql(f) for f in obj.features if not f.is_autogenerated]
            except Exception:
                _logger.error(f"Failed to convert {obj=} with {blob_id=} to GQL.", exc_info=True)
            kwargs.update(
                {
                    "name": obj.namespace,
                    "entity_type": EntityUpdateType.feature_class,
                    "resolver_gql": None,
                    "feature_fields_gql": gql,
                }
            )
        else:
            raise ValueError(f"Unrecognized object type: {obj}")
        return cls(**kwargs)

    @classmethod
    def from_exc(cls, blob_id: str, exc: Exception):
        # need to build these explicitly in 3.8
        exc_type, exc_value, exc_traceback = sys.exc_info()
        trace_lines = traceback.format_exception(exc_type, value=exc_value, tb=exc_traceback)

        err = BlobLoadError(
            traceback="".join(trace_lines),
            message=str(exc),
        )
        return cls(
            blob_id=blob_id,
            timestamp=dt.datetime.now(tz=dt.timezone.utc),
            source=EntityUpdateSource.live_update,
            error=err,
            success=False,
        )


@dataclasses_json.dataclass_json
@dataclasses.dataclass
class BranchGraphSummary:
    """
    Summary of the state of the feature/resolver graph on a branch deployment.
    This includes both data from the deployment source and from live updates sent
    from e.g. Jupyter notebooks.
    """

    # includes all features/resolvers currently used to run queries, as well as any validation errors/warnings
    current_graph_gql: Optional[UpsertGraphGQL]

    # Errors caused when importing the deployment source
    import_errors: List[FailedImport]

    # Result or error of each pickled entity update, in chronological order.
    blob_history: List[EntityUpdate]

    # Name's of feature classes currently in the graph that came from pickled blobs
    # i.e. these names are features that _differ_ from the deployment source
    blob_feature_classes: Set[str]

    # FQN's of resolvers currently in the graph that came from pickled blobs
    # i.e. these names are resolvers that _differ_ from the deployment source
    blob_resolvers: Set[str]

    # Need the time the deployment was created in order to display updated_at for the graph
    deployment_created: Optional[str] = None

    def resolver(self, ref: Tuple[Resolver, str]) -> Optional[ResolverRichSummary]:
        if isinstance(ref, Resolver):
            fqn = ref.fqn
        elif isinstance(ref, str):
            fqn = ref
        else:
            raise ValueError(f"Unrecognized type: {ref}")
        short_name = _short_name(fqn)

        if not self.current_graph_gql:
            return None
        assert self.current_graph_gql.resolvers is not None
        resolver = get_unique_item(x for x in self.current_graph_gql.resolvers if _short_name(x.fqn) == short_name)
        update_time = self._get_last_updated_time(EntityUpdateType.resolver, name=resolver.fqn)
        return ResolverRichSummary(resolver, update_time)

    def feature(self, ref: Tuple[Feature, str]) -> Optional[FeatureFieldRichSummary]:
        if isinstance(ref, Feature):
            fqn = ref.fqn
        elif isinstance(ref, str):
            fqn = ref
        else:
            raise ValueError(f"Unrecognized type: {ref}")
        if not self.current_graph_gql:
            return None
        assert self.current_graph_gql.features is not None
        feature_field = get_unique_item(x for x in self.current_graph_gql.features if x.id.fqn == fqn)
        update_time = self._get_last_updated_time(EntityUpdateType.resolver, name=feature_field.id.namespace)
        return FeatureFieldRichSummary(feature_field, update_time)

    # def features_class(self, ref: Tuple[Features, str]) -> Optional["FeaturesClassRichSummary"]:
    #     pass

    # def entity(self, ref: Union[str, Feature, Resolver, Features]):
    #     return self.get(ref)
    #
    # def get(self, ref: Union[str, Feature, Resolver, Features]):
    #     pass

    def _get_last_updated_time(self, entity_type: EntityUpdateType, name: str) -> Optional[dt.datetime]:
        for update in reversed(self.blob_history):
            if not update.success:
                continue
            if update.name == name and update.entity_type == entity_type:
                return _blob_id_to_timestamp(update.blob_id)
        return None


def _short_name(fqn: str):
    return fqn.split(".")[-1]


def _blob_id_to_timestamp(blob_id: str) -> dt.datetime:
    """
    TODO (rkargon) - this relies on the implementation of cuid(), instead we should store the 'blob created timestamp' in the EntityUpdate
    """
    encoding = "0123456789abcdefghijklmnopqrstuvwxyz"

    def _str_to_timestamp(s: str):
        """Assumes `s` is a base-len(encoding) encoded integer."""
        v = 0
        for x in s:
            v *= len(encoding)
            v += encoding.index(x)
        return v

    time_str = blob_id[1:9]
    timestamp = _str_to_timestamp(time_str) // 1000
    return dt.datetime.fromtimestamp(timestamp, tz=dt.timezone.utc)
