from datetime import date, datetime, time, timedelta, timezone
from typing import Any
import json
import sys

from dateutil.parser import parse
import numpy as np
import pandas as pd


class SarusJSONEncoder(json.JSONEncoder):
    def default(self, obj: Any) -> Any:
        if isinstance(obj, np.ndarray):
            return {'_type': 'numpy.ndarray', 'data': obj.tolist()}
        elif isinstance(obj, pd.DataFrame):
            return {
                '_type': 'pandas.DataFrame',
                'data': obj.to_json(date_format='iso'),
            }
        elif isinstance(obj, pd.Series):
            return {
                '_type': 'pandas.Series',
                'data': obj.to_json(date_format='iso'),
            }
        elif isinstance(obj, pd.Timestamp):
            return {'_type': 'pandas.Timestamp', 'data': obj.isoformat()}
        elif isinstance(obj, datetime):
            return {'_type': 'datetime', 'data': obj.isoformat()}
        elif isinstance(obj, timedelta):
            return {'_type': 'timedelta', 'data': obj.total_seconds()}
        elif isinstance(obj, timezone):
            utcoffset_result = obj.utcoffset(None)
            if utcoffset_result is not None:
                return {
                    '_type': 'timezone',
                    'data': utcoffset_result.total_seconds(),
                }
            else:
                raise ValueError("Invalid timezone object")
        elif isinstance(obj, time):
            return {'_type': 'time', 'data': obj.isoformat()}
        elif isinstance(obj, date):
            return {'_type': 'date', 'data': obj.isoformat()}
        elif isinstance(obj, np.generic):
            if np.issubdtype(obj, np.complexfloating):
                complex_obj = obj.astype(np.complex128)
                return {
                    '_type': 'numpy.complex',
                    'data': {
                        '"real"': complex_obj.real,
                        '"imag"': complex_obj.imag,
                    },
                }
            else:
                return {'_type': 'numpy.generic', 'data': obj.item()}
        elif isinstance(obj, pd.MultiIndex):
            return {
                '_type': 'pandas.MultiIndex',
                'data': obj.tolist(),
                'names': obj.names,
                'levels': [level.tolist() for level in obj.levels],
                'codes': list(obj.codes),
            }
        elif isinstance(obj, pd.Index):
            return {
                '_type': 'pandas.Index',
                'class': type(obj).__name__,
                'data': obj.tolist(),
                'dtype': str(obj.dtype),
            }
        elif isinstance(obj, pd.Period):
            return {
                '_type': 'pandas.Period',
                'data': str(obj),
                'freq': obj.freqstr,
            }
        elif isinstance(obj, pd.Timedelta):
            return {'_type': 'pandas.Timedelta', 'data': obj.value}
        elif isinstance(obj, pd.Interval):
            return {'_type': 'pandas.Interval', 'data': (obj.left, obj.right)}
        elif isinstance(obj, pd.Categorical):
            return {
                '_type': 'pandas.Categorical',
                'data': obj.tolist(),
                'categories': obj.categories.tolist(),
                'ordered': obj.ordered,
            }
        elif isinstance(obj, type):
            return {
                '_type': 'class',
                'data': {'"name"': obj.__name__, '"module"': obj.__module__},
            }
        elif isinstance(obj, pd.api.extensions.ExtensionDtype):
            return {'_type': 'dtype', 'data': str(obj)}
        elif isinstance(obj, slice):
            return {'_type': 'slice', 'data': (obj.start, obj.stop, obj.step)}
        elif isinstance(obj, range):
            return {
                '_type': 'range',
                'data': {
                    '"start"': obj.start,
                    '"stop"': obj.stop,
                    '"step"': obj.step,
                },
            }
        return super().default(obj)

    def encode_obj(self, obj: Any) -> Any:
        if isinstance(obj, tuple):
            return {
                '_type': 'tuple',
                'data': [self.encode_obj(v) for v in obj],
            }
        elif isinstance(obj, list):
            return [self.encode_obj(v) for v in obj]
        elif isinstance(obj, dict):
            return {self.encode(k): self.encode_obj(v) for k, v in obj.items()}
        return obj

    def encode(self, obj: Any) -> str:
        obj_transformed = self.encode_obj(obj)
        return super().encode(obj_transformed)

    @classmethod
    def encode_bytes(cls, obj: Any) -> bytes:
        encoder = cls()
        return (encoder.encode(obj)).encode('utf-8')


class SarusJSONDecoder(json.JSONDecoder):
    def decode(self, s: str, *args: Any, **kwargs: Any) -> Any:
        obj = super().decode(s, *args, **kwargs)
        return self.decode_obj(obj)

    def decode_obj(self, obj: Any) -> Any:
        if isinstance(obj, dict):
            if '_type' in obj:
                data = self.decode_obj(obj['data'])
                if obj['_type'] == 'tuple':
                    return tuple(self.decode_obj(v) for v in data)
                elif obj['_type'] == 'numpy.ndarray':
                    return np.array(data)
                elif obj['_type'] == 'pandas.DataFrame':
                    return pd.read_json(data, convert_dates=True)
                elif obj['_type'] == 'pandas.Series':
                    return pd.read_json(data, typ='series', convert_dates=True)
                elif obj['_type'] == 'pandas.Timestamp':
                    return pd.Timestamp(data)
                elif obj['_type'] == 'datetime':
                    return parse(data)
                elif obj['_type'] == 'timedelta':
                    return timedelta(seconds=data)
                elif obj['_type'] == 'timezone':
                    return timezone(timedelta(seconds=data))
                elif obj['_type'] == 'time':
                    return parse(data).time()
                elif obj['_type'] == 'date':
                    return parse(data).date()
                elif obj['_type'] == 'numpy.generic':
                    return np.array(data).item()
                elif obj['_type'] == 'numpy.complex':
                    return np.complex128(complex(data['real'], data['imag']))
                elif obj['_type'] == 'pandas.Index':
                    cls = getattr(pd, obj['class'])
                    return cls(data, dtype=obj['dtype'])
                elif obj['_type'] == 'pandas.MultiIndex':
                    return pd.MultiIndex(
                        levels=[pd.Index(level) for level in obj['levels']],
                        codes=self.decode_obj(obj['codes']),
                        names=obj['names'],
                    )
                elif obj['_type'] == 'pandas.Period':
                    return pd.Period(data, freq=obj['freq'])
                elif obj['_type'] == 'pandas.Timedelta':
                    return pd.to_timedelta(data)
                elif obj['_type'] == 'pandas.Interval':
                    return pd.Interval(*data)
                elif obj['_type'] == 'pandas.Categorical':
                    return pd.Categorical(
                        data,
                        categories=obj['categories'],
                        ordered=obj['ordered'],
                    )
                elif obj['_type'] == 'class':
                    if data['module'] in ['builtins', 'numpy', 'pandas']:
                        cls = getattr(
                            sys.modules[data['module']], data['name']
                        )
                        if isinstance(cls, type):
                            return cls
                        else:
                            raise ValueError("Decoded object is not a type")
                    else:
                        raise ValueError("Invalid module name")
                elif obj['_type'] == 'dtype':
                    return pd.api.types.pandas_dtype(data)
                elif obj['_type'] == 'slice':
                    return slice(*data)
                elif obj['_type'] == 'range':
                    return range(data['start'], data['stop'], data['step'])
            return {self.decode(k): self.decode_obj(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [self.decode_obj(v) for v in obj]
        return obj

    @classmethod
    def decode_bytes(cls, b: bytes, *args: Any, **kwargs: Any) -> Any:
        decoder = cls()
        return decoder.decode(b.decode('utf-8'), *args, **kwargs)
