from __future__ import annotations

import sys
from datetime import datetime, timedelta
from typing import Union, Any

import numpy as np
import pandas as pd
from loguru import logger

DatetimeLike = Union[str, datetime, np.datetime64]


class Clock:
    def __init__(self, sim_start: str | datetime):
        self.sim_start = pd.to_datetime(sim_start)

    def to_datetime(self, simtime: int) -> datetime:
        return self.sim_start + timedelta(seconds=simtime)

    def to_simtime(self, dt: datetime) -> int:
        return int((dt - self.sim_start).total_seconds())


def disable_rt_warnings(behind_threshold: float):
    """Disables Mosaik's rt_check warnings.

    Mosaik's current implementation of the real-time checks are faulty and a new
    implementation is already outlined in
    https://gitlab.com/mosaik/mosaik/-/issues/172. As this issue was opened in
    June 2023, it is unlikely that the issue will be resolved soon. Therefore,
    we disable the warnings with a threshold after which these offsets are
    logged.

    Args:
        behind_threshold: Time the simulation is allowed to be behind schedule.
    """

    def filter_record(record):
        is_warning = record["level"].name == "WARNING"
        is_mosaik_log = record["name"].startswith("mosaik")
        is_user_warning = record["function"] == "user_warning"
        try:
            message_parts = record["message"].split(" - ")[1].split("s")[0]
            is_below_threshold = float(message_parts) < behind_threshold
        except IndexError:
            return False
        return not (is_warning and is_mosaik_log and is_user_warning and is_below_threshold)

    # Add the filter to the logger
    logger.remove()
    logger.add(sys.stdout, filter=filter_record)


def flatten_dict(d: dict, parent_key: str = "") -> dict:
    items: list[tuple[str, Any]] = []
    for k, v in d.items():
        new_key = parent_key + "." + k if parent_key else k
        if isinstance(v, dict):
            items.extend(flatten_dict(v, str(new_key)).items())
        else:
            items.append((new_key, v))
    return dict(items)
