from __future__ import annotations
from datetime import datetime, timedelta, tzinfo
import dataclasses
import importlib
import importlib.util
import json
import typing as t
import os
import time

from lazy_loader import load as _lazy_load
from loguru import logger
from pydantic import BaseModel
import numpy as np

# import numpy.typing as npt
# import pyarrow as pa
import polars as pl
import datetime


# -----------------------------------------------------------
# utility routines for JSON serialization - parts picked off
# from the openai-evals codebase
# -----------------------------------------------------------


def to_py_types(obj: t.Any) -> t.Any:
    import inspect

    # for nested dataclasses/pydantic models/operators
    if isinstance(obj, datetime.datetime):
        return str(obj)
    if isinstance(obj, dict):
        return {k: to_py_types(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [to_py_types(v) for v in obj]
    elif dataclasses.is_dataclass(obj):
        return to_py_types(dataclasses.asdict(obj))
    elif hasattr(obj, "_uptrain_op_name"):
        if hasattr(obj, "_uptrain_op_custom"):
            if hasattr(obj, "_uptrain_op_custom_source"):
                source = obj._uptrain_op_custom_source
            else:
                source = inspect.getsource(obj.__class__)
            return {
                "op_name": getattr(obj, "_uptrain_op_name"),
                "params": obj.dict(include=set(obj.__fields__)),
                "source": source,
            }
        else:
            return {
                "op_name": getattr(obj, "_uptrain_op_name"),
                "params": obj.dict(include=set(obj.__fields__)),
            }
    elif isinstance(obj, BaseModel):
        return obj.model_dump()

    # for numpy types
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()

    return obj


class UpTrainEncoder(json.JSONEncoder):
    """Special json encoder for numpy types"""

    def default(self, obj: t.Any) -> str:
        return to_py_types(obj)


def jsondumps(obj: t.Any, **kwargs) -> str:
    return json.dumps(obj, cls=UpTrainEncoder, **kwargs)


def jsondump(obj: t.Any, fp: t.Any, **kwargs) -> None:
    json.dump(obj, fp, cls=UpTrainEncoder, **kwargs)


def jsonloads(s: str, **kwargs) -> t.Any:
    return json.loads(s, **kwargs)


def jsonload(fp: t.Any, **kwargs) -> t.Any:
    return json.load(fp, **kwargs)


# -----------------------------------------------------------
# utility routines for converting between Arrow and Numpy
# -----------------------------------------------------------


# def array_np_to_arrow(arr: npt.NDArray) -> pa.Array:
#     assert arr.ndim in (1, 2), "Only 1D and 2D arrays are supported."
#     if arr.ndim == 1:
#         return pa.array(arr)
#     else:
#         dim1, dim2 = arr.shape
#         return pa.ListArray.from_arrays(
#             np.arange(0, (dim1 + 1) * dim2, dim2), arr.ravel()
#         )


# def array_arrow_to_np(arr: pa.Array) -> npt.NDArray:
#     if isinstance(arr, pa.ChunkedArray):
#         arr = arr.combine_chunks()

#     if not pa.types.is_list(arr.type):
#         return arr.to_numpy(zero_copy_only=False)  # assume a 1D array
#     else:
#         dim1 = len(arr)  # assume a 2D array
#         return np.asarray(arr.values.to_numpy(zero_copy_only=False)).reshape(dim1, -1)


# def arrow_batch_to_table(batch_or_tbl: t.Union[pa.Table, pa.RecordBatch]) -> pa.Table:
#     if not isinstance(batch_or_tbl, pa.Table):
#         return pa.Table.from_batches([batch_or_tbl])
#     else:
#         return batch_or_tbl


# def table_arrow_to_np_arrays(
#     tbl: pa.Table, cols: list[t.Union[str, int]]
# ) -> list[np.ndarray]:
#     return [array_arrow_to_np(tbl.column(c)) for c in cols]


# def np_arrays_to_arrow_table(arrays: list[npt.NDArray], cols: list[str]) -> pa.Table:
#     return pa.Table.from_pydict(
#         {c: array_np_to_arrow(arr) for c, arr in zip(cols, arrays)}
#     )


def polars_to_json_serializable_dict(data: pl.DataFrame):
    data_dictn = data.to_dicts()

    # Remove cases which are not json serializable
    for key in list(data_dictn[0].keys()):
        delete_key = False
        if "datetime" in str(type(data_dictn[0][key])):
            delete_key = True

        if delete_key:
            for row in data_dictn:
                del row[key]

    try:
        json.dumps(data_dictn)
    except Exception:
        for key in list(data_dictn[0].keys()):
            try:
                json.dumps([x[key] for x in data_dictn])
            except Exception:
                for row in data_dictn:
                    del row[key]

    return data_dictn


def polars_to_pandas(data: pl.DataFrame):
    """Convert a polars dataframe to a pandas dataframe"""
    # FIXME: obscure error during pandas conversion through pyarrow for string columns, though
    # they already are utf-8 encoded by arrow/polars?? I tried pandas 1.5.3 and pyarrow>12,
    # as rec-d on github.
    import pandas as pd

    try:
        pd_data = data.to_pandas()
    except Exception:
        # convert to python native types first and then to pandas
        logger.warning(
            "Error converting polars to pandas. Trying to convert to python native types first."
        )

        lazy_load_dep("pyarrow", "pyarrow>=10.0.0")

        # We can't iterate over polars as such, since rust panics can't be caught in python
        # convert to pyarrow first
        all_rows = []
        arrow_data = data.to_arrow()
        for row in range(len(data)):
            try:
                row_dict = {}
                for col in data.columns:
                    row_dict[col] = arrow_data[col][row].as_py()
                all_rows.append(row_dict)
            except Exception as e:
                logger.warning(f"Error converting row {row}: {e}")

        pd_data = pd.DataFrame(all_rows)
    return pd_data


# -----------------------------------------------------------
# routines for working with local files and directories
# -----------------------------------------------------------


def clear_directory(dir_path: str):
    """Clears the directory at dir_path but without deleting the directory itself. `shutil.rmtree` will
    have difficulties with mounted volumes or network drives.
    """
    import shutil

    for filename in os.listdir(dir_path):
        file_path = os.path.join(dir_path, filename)
        if os.path.isfile(file_path) or os.path.islink(file_path):
            os.unlink(file_path)
        elif os.path.isdir(file_path):
            shutil.rmtree(file_path)


# -----------------------------------------------------------
# routines to deal with time
# -----------------------------------------------------------


class Clock:
    """Makes testing easier by anchoring to an older time. This is helpful when
    testing with older datasets. While if initialized with no arguments, the clock
    starts at the current time.
    """

    behind_by: timedelta
    tzone: t.Optional[tzinfo]

    def __init__(self, init_at: datetime):
        tz = init_at.tzinfo
        if tz is None or tz.utcoffset(init_at) is None:
            self.tzone = None
            self.behind_by = datetime.now() - init_at
        else:
            self.tzone = tz
            self.behind_by = datetime.now(tz=tz) - init_at

    def now(self) -> datetime:
        """Return the current time, adjusted by the amount of time the clock is behind."""
        return datetime.now(tz=self.tzone) - self.behind_by

    def sleep(self, seconds: float):
        """If the clock is behind, catch up. Else, sleep for the given duration."""
        seconds_behind = self.behind_by.total_seconds()
        if seconds_behind > 0:
            print(f"advancing the clock by {seconds} seconds")
        else:
            print(f"sleeping for {seconds} seconds")
            time.sleep(seconds)
        self.behind_by = timedelta(seconds=max(seconds_behind - seconds, 0))


class Timer:
    """Context manager for timing code blocks"""

    time: float

    def __enter__(self):
        self.time = time.perf_counter()
        return self

    def __exit__(self, type, value, traceback):
        self.time = round(time.perf_counter() - self.time, 3)


# -----------------------------------------------------------
# routines to deal with optional dependencies
# -----------------------------------------------------------


def lazy_load_dep(import_name: str, package_name: str):
    """Helper function to lazily load optional dependencies. If the dependency is not
    present, the function will raise an error _when used_.

    NOTE: This wrapper adds a warning message at import time.
    """
    try:
        spec = importlib.util.find_spec(import_name)
    except Exception:
        spec = None
    if spec is None:
        logger.warning(
            f"Optional feature dependent on missing package: {import_name} was initialized.\n"
            f"Use `pip install {package_name}` to install the package if running locally."
        )

    return _lazy_load(import_name)
