import os
import re
import subprocess
import tempfile
from importlib.metadata import version

import fitz

from tikzpics.color import Color
from tikzpics.coordinate import TikzCoordinate
from tikzpics.layer import LayerCollection
from tikzpics.loop import Loop
from tikzpics.node import Node
from tikzpics.path import Path
from tikzpics.plot import Plot3D
from tikzpics.variable import Variable

line_length = 50
version_string = f"% Tikzfigure generated by tikzpics v{version('tikzpics')}"
version_string += " " * (line_length - len(version_string) - 2) + "%\n"
link_string = "% https://github.com/max-models/tikzpics"
link_string += " " * (line_length - len(link_string) - 2) + "%\n"
line_separator = "% " + "-" * (line_length - 5) + " %\n"
TIKZFIGURE_HEADER = line_separator + version_string + link_string + line_separator


class TikzFigure:
    def __init__(
        self,
        ndim=2,
        label: str | None = None,
        grid: bool = False,
        tikz_code: str | None = None,
        figure_setup: str | None = None,
        figsize: tuple = (10, 6),
        caption: str | None = None,
        description: str | None = None,
    ):
        """Initialize the TikzFigure class for creating TikZ figures."""

        self._figsize = figsize
        self._caption = caption
        self._description = description
        self._label = label
        self._grid = grid
        self._tikz_code = tikz_code
        self._figure_setup = figure_setup

        # Initialize lists to hold Node and Path objects
        # TODO: nodes, paths, layers should have @property and @setter methods
        self._layers = LayerCollection()
        self._variables = []
        self._colors = []
        self._ndim = ndim

        # Counter for unnamed nodes
        self._node_counter = 0

        if self._tikz_code:
            lines = self._tikz_code.split("\n")
            lines = [line.lstrip().rstrip() for line in lines]
            lines = [line for line in lines if line not in ["", "\n"]]
            lines = [line for line in lines if not line[0] == "%"]

            assert lines[0] == "\\begin{tikzpicture}"
            assert lines[-1] == "\\end{tikzpicture}"

            current_layer = 0
            for line in lines[1:-1]:
                # print(line)

                # Match \begin{pgfonlayer}{layer}
                match_pgfdeclarelayer = re.search(r"\\pgfdeclarelayer\{(\d+)\}", line)
                if match_pgfdeclarelayer:
                    layer = match_pgfdeclarelayer.group(1)
                    self.layers.add_layer(layer)
                    # print(f"Layer number: {layer}")

                # Match \begin{pgfonlayer}{layer}
                match_begin_pgfonlayer = re.search(
                    r"\\begin\{pgfonlayer\}\{(\d+)\}",
                    line,
                )
                if match_begin_pgfonlayer:
                    current_layer = match_begin_pgfonlayer.group(1)
                    # print(f'Current layer: {current_layer}')

                # Match \end{pgfonlayer}{layer}
                match_end_pgfonlayer = re.search(r"\\end\{pgfonlayer\}\{(\d+)\}", line)
                if match_end_pgfonlayer:
                    current_layer = 0
                    # print(f'Current layer: {current_layer}')

                # Match \node[attributes] at (x,y) {content};
                match_node = re.search(
                    r"\\node(?:\[([^\]]+)\])? \((\w+)\) at \(([^,]+, [^)]+)\) \{(.*)\};",
                    line,
                )
                if match_node:
                    attributes = match_node.group(1)
                    attributes = [
                        attribute.lstrip().rstrip()
                        for attribute in attributes.split(",")
                    ]
                    attributes = [
                        attribute for attribute in attributes if not attribute == ""
                    ]
                    attributes_dict = dict(attr.split("=") for attr in attributes)
                    node_name = match_node.group(2)
                    coordinates = match_node.group(3)
                    coordinates = [c for c in coordinates.split(",")]
                    coordinates = [c for c in coordinates]
                    coordinates = [c.lstrip().rstrip() for c in coordinates]
                    content = match_node.group(4)
                    # print(f"Attributes: {attributes_dict}")
                    # print(f"Node name: {node_name}")
                    # print(f"Coordinates: {coordinates}")
                    # print(f"Node text: '{node_text}'")
                    # print(coordinates.split(','))

                    self.add_node(
                        x=coordinates[0],
                        y=coordinates[1],
                        label=node_name,
                        layer=current_layer,
                        content=content,
                        **attributes_dict,
                    )

                # Match \draw[attributes] (node1.center) to (node2.center) to (node2.center)...;
                match_draw = re.search(
                    r"\\draw(?:\[([^\]]+)\])? ((?:\(\w+\.*\) to )+\(\w+\.*\));",
                    line,
                )
                if match_draw:
                    attributes = match_draw.group(1)
                    attributes = [
                        attribute.lstrip().rstrip()
                        for attribute in attributes.split(",")
                    ]
                    attributes = [
                        attribute for attribute in attributes if not attribute == ""
                    ]
                    path = match_draw.group(2)
                    nodes = re.findall(r"\((\w+)\.*\)", path)

                    # print(f"Attributes: {attributes}")
                    # print(f"Path: {path}")
                    # print(f"Nodes: {nodes}")
                    self._add_path(nodes, options=attributes, layer=current_layer)

    # ------------------------------------------------------------- #
    # Public methods

    def add(self, items: list | tuple | Node, layer=0, verbose: bool = False):
        if not isinstance(items, list | tuple):
            items = [items]

        for item in items:
            if isinstance(item, Node):
                if item.label == "":
                    item._label = f"node{self._node_counter}"
                self._node_counter += 1
                self.layers.add_item(item=item, layer=layer, verbose=verbose)

    def colorlet(
        self,
        name,
        color_spec: str,
        layer: int = 0,
    ) -> Color:
        color = Color(color_spec=color_spec)

        self._colors.append((name, color))
        return color

    def add_node(
        self,
        x: float | int | str | None = None,
        y: float | int | str | None = None,
        z: float | int | str | None = None,
        label: str | None = None,
        content: str = "",
        layer: int = 0,
        comment: str | None = None,
        options: list | str | None = None,
        verbose: bool = False,
        **kwargs,
    ):
        """
        Add a node to the TikZ figure.

        Parameters:
        - x (float): X-coordinate of the node.
        - y (float): Y-coordinate of the node.
        - label (str, optional): Label of the node. If None, a default label will be assigned.
        - content: (str, optional): Content of the node
        - comment: (str, optional): Comment to be added before the node
        - **kwargs: Additional TikZ node options (e.g., shape, color).

        Returns:
        - node (Node): The Node object that was added.
        """
        if options is None:
            options = []

        if isinstance(options, str):
            options = [options]

        if label is None:
            label = f"node{self._node_counter}"
        node = Node(
            x=x,
            y=y,
            z=z,
            label=label,
            layer=layer,
            content=content,
            comment=comment,
            options=options,
            **kwargs,
        )
        self._node_counter += 1
        self.layers.add_item(item=node, layer=layer, verbose=verbose)
        return node

    def add_variable(
        self,
        label: str,
        value: int | float,
        layer: int | None = 0,
        comment: str | None = None,
        verbose: bool = False,
    ):
        variable = Variable(
            label=label,
            value=value,
            layer=layer,
            comment=comment,
        )
        self._variables.append(variable)

        # TODO: Allow for special variables in each layer
        # self.layers.add_item(item=variable, layer=layer, verbose=verbose)
        return variable

    def filldraw(
        self,
        nodes: list,
        layer: int = 0,
        comment: str | None = None,
        center=False,
        **kwargs,
    ):
        path = self._add_path(
            nodes=nodes,
            layer=layer,
            comment=comment,
            center=center,
            tikz_command="filldraw",
            **kwargs,
        )
        return path

    def draw(
        self,
        nodes: list,
        layer: int = 0,
        comment: str | None = None,
        center=False,
        **kwargs,
    ):
        path = self._add_path(
            nodes=nodes,
            layer=layer,
            comment=comment,
            center=center,
            tikz_command="draw",
            **kwargs,
        )
        return path

    def plot3d(
        self,
        x: list,
        y: list,
        z: list,
        layer: int = 0,
        comment: str | None = None,
        center=False,
        verbose: bool = False,
        **kwargs,
    ):
        plot = Plot3D(
            x=x,
            y=y,
            z=z,
            comment=comment,
            center=center,
            **kwargs,
        )

        self.layers.add_item(item=plot, layer=layer, verbose=verbose)
        return plot

    def add_loop(
        self,
        variable,
        values,
        layer=0,
        comment=None,
        verbose: bool = False,
    ):
        loop_obj = Loop(
            variable=variable,
            values=values,
            layer=layer,
            comment=comment,
        )

        self.layers.add_item(item=loop_obj, layer=layer, verbose=verbose)
        return loop_obj

    def generate_tikz(self):
        """
        Generate the TikZ script for the figure.

        Returns:
        - tikz_script (str): The TikZ script as a string.
        """
        tikz_script = TIKZFIGURE_HEADER
        tikz_script += "\\begin{tikzpicture}\n"
        if self._figure_setup:
            tikz_script += f"[{self._figure_setup}]"

        # Add variables
        if len(self.variables) > 0:
            for variable in self.variables:
                tikz_script += (
                    f"\\pgfmathsetmacro{{\\{variable.label}}}{{{variable.value}}}\n"
                )

        # Add colors
        if len(self.colors) > 0:
            for name, color in self.colors:
                tikz_script += f"\\colorlet{{{name}}}{{{color.color_spec}}}\n"

        # Add 3D axis
        if self.ndim == 3:
            tikz_script += "\\begin{axis}[\n"
            tikz_script += "view={20}{30},\n"
            tikz_script += "axis lines=center,\n"
            tikz_script += "xlabel={$x$},\n"
            tikz_script += "ylabel={$y$},\n"
            tikz_script += "zlabel={$z$},\n"
            tikz_script += "grid=major\n"
            tikz_script += "    ]\n"

        tikz_script += "\n"

        tikz_script += "% Define the layers library\n"
        layers = sorted([str(layer) for layer in self.layers.layers.keys()])
        for layer in layers:
            tikz_script += f"\\pgfdeclarelayer{{{layer}}}\n"
        if len(layers) > 0:
            tikz_script += f"\\pgfsetlayers{{{','.join(layers)}}}\n"

        # Add grid if enabled
        # TODO: Create a Grid class
        if self._grid:
            tikz_script += (
                "    \\draw[step=1cm, gray, very thin] (-10,-10) grid (10,10);\n"
            )
        ordered_layers = []
        buffered_layers = set()

        for key, layer in self.layers.layers.items():
            # layer_order, buffered_layers = update_layer_order(layer, layer_order, buffered_layers)
            reqs = layer.get_reqs()
            if all([r == layer.label for r in reqs]):
                ordered_layers.append(layer)
            elif all([r in [layer.label for layer in ordered_layers] for r in reqs]):
                ordered_layers.append(layer)
            else:
                buffered_layers.add(layer)

            for buffered_layer in buffered_layers:
                buff_reqs = buffered_layer.get_reqs()
                if all(
                    [r in [layer.label for layer in ordered_layers] for r in buff_reqs],
                ):
                    print("Move layer from buffer")
                    ordered_layers.append(key)
                    buffered_layers.remove(key)
        assert (
            len(buffered_layers) == 0
        ), f"Layer order is impossible for layer {[layer.label for layer in buffered_layers]}"
        for layer in ordered_layers:
            tikz_script += layer.generate_tikz()

        if self.ndim == 3:
            tikz_script += "\\end{axis}\n"
        tikz_script += "\\end{tikzpicture}"

        # Wrap in figure environment if necessary
        if self._caption or self._description or self._label:
            figure_env = "\\begin{figure}\n" + tikz_script + "\n"
            if self._caption:
                figure_env += f"    \\caption{{{self._caption}}}\n"
            if self._label:
                figure_env += f"    \\label{{{self._label}}}\n"
            figure_env += "\\end{figure}"
            tikz_script = figure_env
        tikz_script = self._add_tabs(tikz_script)
        return tikz_script

    def generate_standalone(self):
        tikz_code = self.generate_tikz()

        # Create a minimal LaTeX document
        latex_document = (
            "\\documentclass[border=10pt]{standalone}\n"
            "\\usepackage{tikz}\n"
            "\\usepackage{pgfplots}\n"
            "\\pgfplotsset{compat=newest}\n"
            "\\usetikzlibrary{arrows.meta}\n"
            "\\begin{document}\n"
            f"{tikz_code}\n"
            "\\end{document}"
        )
        return latex_document

    def compile_pdf(self, filename="output.pdf", verbose=False):
        """
        Compile the TikZ script into a PDF using pdflatex.

        Parameters:
        - filename (str): The name of the output PDF file (default is 'output.pdf').

        Notes:
        - Requires 'pdflatex' to be installed and accessible from the command line.
        """
        latex_document = self.generate_standalone()
        if verbose:
            print(latex_document)
        # Use a temporary directory to store the LaTeX files
        with tempfile.TemporaryDirectory() as tempdir:
            tex_file = os.path.join(tempdir, "figure.tex")
            with open(tex_file, "w") as f:
                f.write(latex_document)
            # Run pdflatex
            try:
                # Split the path
                head_tail = os.path.split(os.path.abspath(filename))
                output_directory = head_tail[0]
                jobname = head_tail[1].replace(".pdf", "")
                cmd = [
                    "pdflatex",
                    "-interaction=nonstopmode",
                    "-jobname",
                    f"{jobname}",
                    "-output-directory",
                    f"{output_directory}",
                    tex_file,
                ]
                if verbose:
                    print(f"{cmd =}")
                subprocess.run(
                    cmd,
                    cwd=tempdir,
                    check=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
                # Remove .aux and .log files
                os.remove(os.path.abspath(filename).replace(".pdf", ".aux"))
                os.remove(os.path.abspath(filename).replace(".pdf", ".log"))
            except subprocess.CalledProcessError as e:
                print("An error occurred while compiling the LaTeX document:")
                print(e.stderr.decode())
                return

    def savefig(self, filename, verbose=False):
        """
        Save the TikZ figure to a file (PDF, PNG, JPG) using pure Python tools.
        """
        ext = os.path.splitext(filename)[1].lower()

        if ext == ".pdf":
            # Direct compile
            self.compile_pdf(filename=filename, verbose=verbose)
        elif ext == ".tikz":
            tikz_code = self.generate_tikz()
            with open(filename, "w") as f:
                f.write(tikz_code)
        elif ext in [".png", ".jpg", ".jpeg"]:
            # Compile to a temporary PDF first
            with tempfile.TemporaryDirectory() as tempdir:
                temp_pdf = os.path.join(tempdir, "temp.pdf")
                # temp_pdf = "_temp.pdf"
                self.compile_pdf(filename=temp_pdf, verbose=verbose)

                if verbose:
                    print(f"Converting {temp_pdf} → {filename}")

                # Convert PDF → image
                doc = fitz.open(temp_pdf)
                page = doc[0]
                pix = page.get_pixmap(dpi=300)
                pix.save(filename)
                doc.close()
        else:
            raise ValueError(f"Unsupported file format: {ext}")

    def show(self, width=None, height=None, verbose=False):
        from IPython.display import Image, display

        with tempfile.TemporaryDirectory() as tempdir:
            temp_pdf = os.path.join(tempdir, "temp.png")
            self.savefig(filename=temp_pdf, verbose=verbose)
            display(Image(filename=temp_pdf, width=width, height=height))

    # ------------------------------------------------------------- #
    # Private methods

    def _add_path(
        self,
        nodes: list,
        layer: int = 0,
        comment: str | None = None,
        center=False,
        tikz_command="draw",
        verbose=False,
        **kwargs,
    ):
        """
        Add a line or path connecting multiple nodes.

        Parameters:
        - nodes (list of str): List of node names to connect OR list of coordinates
        - **kwargs: Additional TikZ path options (e.g., style, color).

        Examples:
        - add_path(['A', 'B', 'C'], color='blue')
          Connects nodes A -> B -> C with a blue line.
        """
        if not isinstance(nodes, list):
            raise ValueError("nodes parameter must be a list of node names.")

        nodes_cleaned = []

        for node in nodes:
            if isinstance(node, Node):
                nodes_cleaned.append(node)
            elif isinstance(node, str):
                # Find the node by its label
                nodes_cleaned.append(self.layers.get_node(node))
            elif isinstance(node, tuple) or isinstance(node, list):
                node = tuple(node)
                nodes_cleaned.append(TikzCoordinate(*node, layer=layer))
            else:
                raise NotImplementedError(
                    f"{node =}, {type(node) =} is not a valid node type!",
                )

        if verbose:
            print(f"Creating a path with the following nodes {nodes_cleaned}")

        path = Path(
            nodes=nodes_cleaned,
            comment=comment,
            center=center,
            tikz_command=tikz_command,
            **kwargs,
        )
        self.layers.add_item(item=path, layer=layer, verbose=verbose)
        return path

    def _add_tabs(self, tikz_script):
        tikz_script_new = ""
        tab_str = "    "
        num_tabs = 0
        for line in tikz_script.split("\n"):
            if "\\end" in line or "end \\foreach" in line:
                num_tabs = max(num_tabs - 1, 0)
            tikz_script_new += f"{tab_str*num_tabs}{line}\n"
            if "\\begin" in line or "start \\foreach" in line:
                num_tabs += 1
        return tikz_script_new

    # ---------------------------------------------------------------- #
    # Properties

    @property
    def ndim(self):
        return self._ndim

    @property
    def layers(self) -> LayerCollection:
        return self._layers

    @property
    def colors(self):
        return self._colors

    @property
    def variables(self) -> list:
        return self._variables
