from __future__ import annotations

import io
from enum import Enum
from typing import Any, Iterable, cast

import cairosvg  # pyright: ignore
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.artist import Artist
from matplotlib.axes import Axes
from matplotlib.figure import Figure

from newclid.animation.artists_in_animation import init_point_name_to_artist
from newclid.animation.clause import ClauseArtists, init_clauses_to_draw
from newclid.animation.goal import init_goals_to_draw
from newclid.animation.justifications import init_justifications_to_draw
from newclid.animation.text_animation import TextAnimation
from newclid.draw.figure import init_figure
from newclid.draw.theme import DrawTheme
from newclid.jgex.clause import JGEXClause
from newclid.jgex.formulation import JGEXFormulation
from newclid.justifications._index import JustificationType
from newclid.predicate_types import PredicateArgument
from newclid.proof_data import ProofData
from newclid.symbols.symbols_registry import SymbolsRegistry


class ProofAnimationParts(Enum):
    SHOWCASE = "Showcase of the final figure"
    CONSTRUCTION = "Construction of the initial figure"
    PROOF = "Proof"
    GOAL = "Summary of proven goals"


class ProofAnimation:
    def __init__(
        self,
        proof_data: ProofData,
        symbols: SymbolsRegistry,
        theme: DrawTheme,
        jgex_problem: JGEXFormulation | None = None,
        figure_kwargs: dict[str, Any] | None = None,
        image_path: str | None = None,
    ):
        self.jgex_problem = jgex_problem
        self.symbols = symbols
        self.theme = theme

        figure_kwargs = figure_kwargs or {"figsize": (24, 12), "dpi": 80}
        self.fig, axes = cast(
            tuple[Figure, dict[str, Axes]],
            plt.subplot_mosaic(  # pyright: ignore
                [
                    ["figure", "figure", "text", "text"],
                    ["figure", "figure", "text", "text"],
                ],
                **figure_kwargs,
            ),
        )

        self.figure_ax = axes["figure"]
        self.text_ax = axes["text"]

        init_figure(self.figure_ax, list(self.symbols.points))

        self.aux_point_names: set[PredicateArgument] = set()
        if self.jgex_problem is not None:
            self.aux_point_names = set(
                pname
                for clause in self.jgex_problem.auxiliary_clauses
                for pname in clause.points
            )

        self.point_name_to_artist = init_point_name_to_artist(
            ax=self.figure_ax,
            color_theme=self.theme,
            aux_point_names=self.aux_point_names,
            symbols=self.symbols,
        )

        self.goals_artists = init_goals_to_draw(
            ax=self.figure_ax,
            proven_goals=proof_data.proven_goals,
            point_name_to_artist=self.point_name_to_artist,
            symbols=self.symbols,
            theme=self.theme,
        )

        self.artists_by_clause: dict[JGEXClause, ClauseArtists] = {}
        self.clauses_to_draw: list[JGEXClause] = []
        if self.jgex_problem is not None:
            self.artists_by_clause = init_clauses_to_draw(
                ax=self.figure_ax,
                jgex_problem=self.jgex_problem,
                point_name_to_artist=self.point_name_to_artist,
                symbols=self.symbols,
                theme=self.theme,
            )
            self.clauses_to_draw = list(self.artists_by_clause.keys())

        (self.justifications_to_draw, self.justification_artists) = (
            init_justifications_to_draw(
                ax=self.figure_ax,
                point_name_to_artist=self.point_name_to_artist,
                symbols=self.symbols,
                proof_data=proof_data,
                theme=self.theme,
            )
        )

        text_lines: list[str] = [ProofAnimationParts.SHOWCASE.value]
        if self.jgex_problem is not None:
            text_lines.extend(
                [ProofAnimationParts.CONSTRUCTION.value]
                + [str(clause) for clause in self.clauses_to_draw]
            )

        text_lines.extend(
            [ProofAnimationParts.PROOF.value]
            + [
                justification_to_draw.text
                for justification_to_draw in self.justifications_to_draw
            ]
        )
        text_lines.extend([ProofAnimationParts.GOAL.value])

        self.text_animation = TextAnimation(
            self.text_ax,
            text_lines,
            theme=self.theme,
            n_lines_shown=21,
            fontsize=18,
        )
        self.text_animation.init()

        self.current_part = ProofAnimationParts.SHOWCASE
        self._has_done_showcase_transition: bool = False
        self._has_done_construction_transition: bool = False
        self._has_done_proof_transition: bool = False

        problem_name = (
            self.jgex_problem.name if self.jgex_problem is not None else "Problem"
        )
        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
        self.fig.suptitle(  # pyright: ignore
            f"Newclid solution for {problem_name} | Goals: {', '.join(str(goal_in_proof.predicate) for goal_in_proof in proof_data.proven_goals)}",
            fontsize=30,
            # fontweight="bold",
            color=self.theme.title_color,
            family=["DejaVu Sans", "STIXGeneral"],
        )
        self.title = self.figure_ax.text(  # pyright: ignore
            0.01,
            0.01,
            self.current_part.value,
            transform=self.figure_ax.transAxes,
            fontsize=22,
            fontweight="bold",
            ha="left",
            va="bottom",
            color=self.theme.title_color,
            family=["DejaVu Sans", "STIXGeneral"],
        )

        if image_path:
            try:
                # Convert SVG to PNG in memory
                png_data = cairosvg.svg2png(url=image_path, scale=0.5)  # pyright: ignore

                # Read the PNG data from memory
                img = mpimg.imread(io.BytesIO(png_data))  # pyright: ignore

                # Place the image on the figure
                self.fig.figimage(  # pyright: ignore
                    img,  # pyright: ignore
                    15,
                    int(self.fig.get_figheight() * self.fig.dpi) - img.shape[0] - 15,
                )

            except FileNotFoundError:
                print(f"Image file not found at: {image_path}")
            except Exception as e:
                print(f"An error occurred while loading the SVG: {e}")

    def animate(self) -> FuncAnimation:
        return FuncAnimation(
            self.fig,
            self._next_step,
            init_func=self._init,
            frames=len(self.text_animation.proof_lines) + 1,
            interval=400,
            blit=True,
        )

    def _init(self) -> Iterable[Artist]:
        self.current_part = ProofAnimationParts.SHOWCASE
        return set()

    def _next_step(self, frame: int) -> Iterable[Artist]:
        match self.current_part:
            case ProofAnimationParts.SHOWCASE:
                return self._showcase(frame)
            case ProofAnimationParts.CONSTRUCTION:
                return self._construction(frame)
            case ProofAnimationParts.PROOF:
                return self._proof(frame)
            case ProofAnimationParts.GOAL:
                return self._goals(frame)

    def _showcase(self, frame: int) -> set[Artist]:
        if self._has_done_showcase_transition:
            return set()

        updated_artists: set[Artist] = set()
        updated_artists.update(self.text_animation.update(frame))
        updated_artists.update(self._hide_all_rule_applications())
        updated_artists.update(self._hide_all_extra_applications())
        updated_artists.update(self._hide_goals())
        self._has_done_showcase_transition = True
        if self.jgex_problem is not None:
            self.current_part = ProofAnimationParts.CONSTRUCTION
        else:
            self.current_part = ProofAnimationParts.PROOF
        return updated_artists

    def _construction(self, frame: int) -> set[Artist]:
        updated_artists: set[Artist] = set()
        if not self._has_done_construction_transition:
            self.title.set_text(ProofAnimationParts.CONSTRUCTION.value)
            updated_artists.add(self.title)
            updated_artists.update(self.text_animation.update(frame))
            for _pname, point_artists in self.point_name_to_artist.items():
                updated_artists.update(point_artists.hide())
            for clause, clause_artists in self.artists_by_clause.items():
                updated_artists.update(clause_artists.hide())

            self._has_done_construction_transition = True
            return updated_artists

        if not self.clauses_to_draw:
            raise ValueError("Should have done transition to proof")
        clause = self.clauses_to_draw.pop(0)
        if not self.clauses_to_draw:
            self.current_part = ProofAnimationParts.PROOF

        for other_clause, other_clause_artists in self.artists_by_clause.items():
            if other_clause == clause:
                continue
            updated_artists.update(other_clause_artists.gray_out())

        updated_artists.update(self.text_animation.update(frame))
        clause_artists = self.artists_by_clause[clause]
        if self.jgex_problem is None:
            raise ValueError("Should have a jgex problem to draw construction")

        is_aux = clause in self.jgex_problem.auxiliary_clauses
        highlight_color = (
            self.theme.aux_point_color if is_aux else self.theme.construction_color
        )
        updated_artists.update(
            clause_artists.manifest(
                aux_point_names=self.aux_point_names, highlight_color=highlight_color
            )
        )
        return updated_artists

    def _proof(self, frame: int) -> set[Artist]:
        updated_artists: set[Artist] = set()
        if not self._has_done_proof_transition:
            self.title.set_text(ProofAnimationParts.PROOF.value)
            updated_artists.add(self.title)
            updated_artists.update(self.text_animation.update(frame))
            updated_artists.update(self._gray_out_all_constructions())
            updated_artists.update(self._highlight_goals())
            self._has_done_proof_transition = True
            return updated_artists

        if not self.justifications_to_draw:
            raise ValueError("Should have done transition to goals")

        updated_artists.update(self._gray_out_all_constructions())
        updated_artists.update(self._gray_out_all_rule_applications())
        updated_artists.update(self._hide_goals())
        updated_artists.update(self._hide_all_extra_applications())

        next_justification = self.justifications_to_draw.pop(0)
        if not self.justifications_to_draw:
            self.current_part = ProofAnimationParts.GOAL

        self.text_animation.update(frame)
        justification_artists = self.justification_artists[
            next_justification.justification
        ]
        updated_artists.update(
            justification_artists.highlight(aux_point_names=self.aux_point_names)
        )
        return updated_artists

    def _goals(self, frame: int) -> set[Artist]:
        updated_artists: set[Artist] = set()
        updated_artists.update(self._gray_out_all_constructions())
        updated_artists.update(self._gray_out_all_rule_applications())
        updated_artists.update(self._hide_all_extra_applications())
        updated_artists.update(self.text_animation.update(frame))
        updated_artists.update(self._highlight_goals())
        return updated_artists

    def _gray_out_all_constructions(self) -> set[Artist]:
        updated_artists: set[Artist] = set()
        for clause_artists in self.artists_by_clause.values():
            updated_artists.update(clause_artists.gray_out())
        return updated_artists

    def _gray_out_all_rule_applications(self) -> set[Artist]:
        updated_artists: set[Artist] = set()
        for _, rule_application in self.justification_artists.items():
            if (
                rule_application.justification.dependency_type
                == JustificationType.RULE_APPLICATION
            ):
                updated_artists.update(rule_application.gray_out())
        return updated_artists

    def _hide_goals(self) -> set[Artist]:
        updated_artists: set[Artist] = set()
        for _, goal_artists in self.goals_artists.items():
            updated_artists.update(goal_artists.hide())
        return updated_artists

    def _highlight_goals(self) -> set[Artist]:
        updated_artists: set[Artist] = set()
        for _, goal_artists in self.goals_artists.items():
            updated_artists.update(
                goal_artists.highlight(
                    aux_point_names=self.aux_point_names,
                    highlight_color=self.theme.goal_color,
                )
            )
        return updated_artists

    def _hide_all_rule_applications(self) -> set[Artist]:
        updated_artists: set[Artist] = set()
        for _, rule_application in self.justification_artists.items():
            if (
                rule_application.justification.dependency_type
                == JustificationType.RULE_APPLICATION
            ):
                updated_artists.update(rule_application.hide())
        return updated_artists

    def _hide_all_extra_applications(self) -> set[Artist]:
        updated_artists: set[Artist] = set()
        for _dep, extra_application in self.justification_artists.items():
            if (
                extra_application.justification.dependency_type
                != JustificationType.RULE_APPLICATION
            ):
                updated_artists.update(extra_application.hide())
        return updated_artists
