from __future__ import annotations

import json
from typing import Any
from typing import List

import pandas as pd
import pyarrow as pa


class MetadataJSONEncoder(json.JSONEncoder):
    def default(self, obj: Any) -> Any:
        # This will correctly encode DFs and other objects w/ a "to_json" method
        if hasattr(obj, 'to_json'):
            return json.loads(obj.to_json())
        else:
            try:
                return super().default(obj)
            except TypeError:
                return str(obj)


class SegmentMetadata(dict):
    @classmethod
    def from_flight_metadata(cls, metadata: pa.Buffer) -> SegmentMetadata:
        metadata_dict = json.loads(metadata.to_pybytes().decode())

        # Convert segment onset/offset values of type str to true timestamps
        if 'segment_params' in metadata_dict:
            segment_params = metadata_dict['segment_params']
            for key in ('onset', 'offset', 'series_onset', 'series_offset'):
                if key in segment_params and isinstance(segment_params[key], str):
                    segment_params[key] = pd.Timestamp(segment_params[key])

        return cls(**metadata_dict)

    def to_json(self) -> str:
        return json.dumps(self, cls=MetadataJSONEncoder)


BatchMetadata = List[SegmentMetadata]
