from __future__ import annotations

from ._feature import Feature


class FeatureDict(dict[str, Feature]):
    """A dictionary mapping keys to Features, with special tracking for time/position.

    Inherits from dict[str, Feature], so can be used directly as a dictionary.
    Provides convenient access to time and position features through properties.

    Attributes:
        time_key (str): The key used for the time feature
        position_key (str | list[str] | None): The key(s) used for position feature(s)
    """

    def __init__(
        self,
        features: dict[str, Feature],
        time_key: str,
        position_key: str | list[str] | None,
        tracklet_key: str | None,
    ):
        """
        Args:
            features (dict[str, Feature]): Mapping from feature keys to Features
            time_key (str): The key for the time feature (must be in features)
            position_key (str | list[str] | None): The key(s) for position feature(s)
            tracklet_key (str | None): The key for the tracklet feature
        """
        super().__init__(features)
        self.time_key = time_key
        self.position_key = position_key
        self.tracklet_key = tracklet_key

        # Validate that time and position keys exist
        if time_key not in self:
            raise KeyError(f"time_key '{time_key}' not found in features")
        # Validate position_key - now that we use register_position_feature(),
        # position_key passed at init should always be in features
        if isinstance(position_key, list):
            for key in position_key:
                if key not in self:
                    raise KeyError(f"position_key '{key}' not found in features")
        else:
            if position_key is not None and position_key not in self:
                raise KeyError(f"position_key '{position_key}' not found in features")

    @property
    def node_features(self) -> dict[str, Feature]:
        """A dict of all node features"""
        return {k: v for k, v in self.items() if v["feature_type"] == "node"}

    @property
    def edge_features(self) -> dict[str, Feature]:
        """A dict of all edge features"""
        return {k: v for k, v in self.items() if v["feature_type"] == "edge"}

    def dump_json(self) -> dict:
        """Dump this FeatureDict to a json compatible dictionary

        Returns:
            dict: A map from the key "FeatureDict" containing features, time_key,
                position_key, and tracklet_key
        """
        return {
            "FeatureDict": {
                "features": {k: dict(v) for k, v in self.items()},
                "time_key": self.time_key,
                "position_key": self.position_key,
                "tracklet_key": self.tracklet_key,
            }
        }

    @classmethod
    def from_json(cls, json_dict: dict) -> FeatureDict:
        """Generate a FeatureDict from a json dict such as one generated by dump_json

        Args:
            json_dict (dict): A dictionary with the key "FeatureDict" containing
                features, time_key, position_key, and tracklet_key

        Returns:
            FeatureDict: A FeatureDict object containing the features from the dictionary
        """
        data = json_dict["FeatureDict"]
        return cls(
            features=data["features"],
            time_key=data["time_key"],
            position_key=data["position_key"],
            # Use get() for backwards compatibility with old JSON files
            tracklet_key=data.get("tracklet_key"),
        )

    def register_position_feature(self, key: str, feature: Feature) -> None:
        """Register the position feature and set the position_key.

        Args:
            key: The key to use for the position feature
            feature: The Feature to register
        """
        self.position_key = key
        self[key] = feature

    def register_tracklet_feature(self, key: str, feature: Feature) -> None:
        """Register the tracklet/track_id feature and set the tracklet_key.

        Args:
            key: The key to use for the tracklet feature
            feature: The Feature to register
        """
        self.tracklet_key = key
        self[key] = feature
