import logging
import time
from pathlib import Path
from threading import Thread
from typing import Callable

from matplotlib import animation, pyplot as plt


def animate(
    fig, frames: list[int] | None, **kwargs
) -> Callable[..., animation.FuncAnimation]:
    """Decorator-style wrapper for FuncAnimation."""
    # Why isn't this how the function is exposed by default?
    return lambda func: animation.FuncAnimation(fig, func, frames, **kwargs)

# writer = animation.writers["ffmpeg"](fps=15, metadata={"artist": "Me"})

def bindfig(fig):
    """
    Create a function which is compatible with the signature of `pyplot.figure`,
    but alters the already-existing figure `fig` instead.
    """
    def ret(**kwargs):
        for i, j in kwargs.items():
            if i == "figsize":
                if j is not None:
                    fig.set_figwidth(j[0])
                    fig.set_figheight(j[1])
                continue
            fig.__dict__["set_" + i](j)
        return fig
    return ret

class SympyAnimationWrapper:
    """
    Context manager for binding a figure for animation.
    """
    def __init__(self, filename: Path | str):
        self._filename = filename if isinstance(filename, Path) else Path(filename)
        self._current_figure = plt.gcf()

        self._figure_temp = plt.figure

    def __enter__(self):
        plt.figure = bindfig(self._current_figure)
        return self.animate

    def __exit__(self, *_):
        plt.figure = self._figure_temp

    def animate(self, *args, **kwargs):
        def wrapper(func):
            ret = animate(self._current_figure, *args, **kwargs)(func)
            current_save = ret.save
            def save(*args, **kwargs):
                if self._filename.exists():
                    logging.error("Ignoring saving animation '%s': file already exists", self._filename)
                    return
                thread = Thread(target=lambda: time.sleep(1) or plt.close())
                thread.run()
                current_save(self._filename, *args, **kwargs)
            ret.save = save

            return ret

        return wrapper