import re
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar

from sympy import Expr, lambdify, latex, simplify, symbols
from sympy.parsing.latex import parse_latex
from timeout_decorator import timeout

from latexvm.type_defs import EnvironmentVariables, Varname

T = TypeVar("T")
TimeoutFunction = Callable[[], T]


ASSIGNMENT_TOKEN: str = "="


class ExpressionType(Enum):
    ASSIGNMENT: int = 1
    FUNCTION: int = 2
    STATEMENT: int = 3


"""
An expression can either be a graphing statement
expression or an assignment expression.
"""


@dataclass
class ExpressionBuffer:
    expr_type: ExpressionType

    @staticmethod
    def new(expr: str) -> Optional["ExpressionBuffer"]:
        type: ExpressionType = Expression.get_expression_type(expr)

        match (type):
            case ExpressionType.FUNCTION:
                lhs, definition = Expression.break_expression(expr)

                fname = Expression.get_function_name(lhs)

                if lhs[-1] != ")":
                    raise Exception(f"Invalid function lhs: '{lhs}'")

                if len(re.findall(r"[a-zA-Z_][a-zA-Z0-9_]*", fname)) != 1:
                    raise Exception(f"Invalid function lhs: '{lhs}'")

                signature_str: str = Expression.get_parameters_str_from_function(lhs)
                signature: List[str] = Expression.get_parameters_from_function(
                    signature_str
                )
                return FunctionExpressionBuffer(
                    expr_type=type,
                    name=fname,
                    signature=signature,
                    signature_str=signature_str,
                    body=definition,
                )
            case ExpressionType.ASSIGNMENT:
                name, body = Expression.break_expression(expr)

                if s := re.search(r"(?<!\{)\b\d+\b(?!\}|\{)", name):
                    raise Exception(f"Invalid identifier: '{s.group()}'")

                if len(re.findall(r"\b[a-zA-Z_0-9{}]+\b", name)) > 1:
                    raise Exception(f"Invalid assignment lhs: '{name}'")

                return AssignmentExpressionBuffer(expr_type=type, name=name, body=body)
            case _:
                return StatementExpressionBuffer(expr_type=type, body=expr)

    def assemble(self) -> str:
        match (self.expr_type):
            case ExpressionType.FUNCTION:
                return f"{self.name}{self.signature_str} = {self.body}"
            case ExpressionType.ASSIGNMENT:
                return f"{self.name} = {self.body}"
            case _:
                return self.body

    def create_callable(self) -> Tuple[Callable, List[Varname]]:
        expr: Expr = Expression.get_expression(expr_str=self.body)
        variables = Expression.extract_variables(expr=expr)
        var_symbols = symbols(names=variables)
        return lambdify(args=var_symbols, expr=expr), var_symbols

    def has_unresolved_function(self) -> List[str]:
        unresolved: List[str] = [f[:-1] for f in re.findall(r"[a-zA-Z]+\(", self.body)]
        return unresolved


@dataclass
class Expression:
    @staticmethod
    def unpack(value: Varname) -> Optional[Varname]:
        if value[0] != "(":
            raise Exception("Value not packed")

        idx = 1
        resolution = 1
        size = len(value)

        while idx < size and resolution > 0:
            if value[idx] == "(":
                resolution += 1
            elif value[idx] == ")":
                resolution -= 1
            idx += 1

        if resolution != 0:
            raise Exception("Value pack missing closing parenthesis")

        unpacked_value = value[1:idx]
        return unpacked_value

    @staticmethod
    def pack(value: Varname) -> Varname:
        return "({})".format(value)

    @staticmethod
    def break_expression(raw_expr: str) -> Tuple[str, str]:
        try:
            # First index of "="
            asn_idx = raw_expr.index("=")
            lhs, rhs = raw_expr[:asn_idx], raw_expr[asn_idx + 1 :]
            return lhs.strip(), rhs.strip()
        except Exception:
            return raw_expr, ""

    @staticmethod
    def is_function(expr_str: str) -> bool:
        regex = r"\b.*\("
        capture = re.search(regex, expr_str)

        if capture is None:
            return False

        return expr_str.index(capture.group()) == 0

    @staticmethod
    def get_function_name(raw_equation: str) -> str:
        return raw_equation.split("(")[0]

    @staticmethod
    def is_function_expression(raw_equation: str) -> bool:
        lhs, _ = Expression.break_expression(raw_expr=raw_equation)
        return Expression.is_function(lhs)

    @staticmethod
    def get_parameters_str_from_function(function_equation: str) -> str:
        first_param_index: int = function_equation.index("(")
        resolution: int = 1
        idx: int = first_param_index + 1
        size = len(function_equation)

        while resolution > 0:
            if idx >= size:
                raise Exception("Function not closed")

            c = function_equation[idx]

            if c == ")":
                resolution -= 1
            elif c == "(":
                resolution += 1

            idx += 1

        parameters = function_equation[first_param_index:(idx)]  # noqa: E203
        return parameters

    @staticmethod
    def get_parameters_from_function(function_equation: str) -> List[Varname]:
        params = Expression.get_parameters_str_from_function(function_equation)[1:-1]
        return [
            "{}".format(param.strip()) for param in params.split(",") if len(param) > 0
        ]

    @staticmethod
    def get_expression_type(raw_equation: str) -> ExpressionType:
        is_assignment = False
        resolution = 0

        for c in raw_equation:
            if c == "{":
                resolution += 1
            elif c == "}":
                resolution -= 1
            elif c == "=" and resolution == 0:
                if is_assignment:
                    raise Exception("Chaining assignment is not allowed")
                is_assignment = True

        if is_assignment:
            if Expression.is_function_expression(raw_equation=raw_equation):
                return ExpressionType.FUNCTION
            return ExpressionType.ASSIGNMENT
        return ExpressionType.STATEMENT

    @staticmethod
    def get_expression(expr_str: str) -> Expr:
        expr: Expr = parse_latex(expr_str)
        return expr

    @staticmethod
    def extract_variables(expr: Expr) -> List[str]:
        variables = list(expr.free_symbols)
        return [str(var) for var in variables]

    @staticmethod
    def capture_function(input: str, func_name: str) -> str:
        fn_idx = input.index(func_name)

        search_str = input[fn_idx:]
        fname = Expression.get_function_name(search_str)
        fparams = Expression.get_parameters_str_from_function(search_str)

        return "{}{}".format(fname, fparams)

    @staticmethod
    def replace_variables(
        expression: str,
        variables: Dict[Varname, Any],
        force_ignore: List[Varname] = list(),
    ) -> str:
        sub_variables = {
            k: v
            for k, v in variables.items()
            if k in expression and k not in force_ignore
        }
        for variable, value in sub_variables.items():
            pat = r"(?<![a-zA-Z\\]){}(?![a-zA-Z])".format(variable)
            expression = re.sub(pat, value, expression)
        return expression

    @staticmethod
    def substitute_function(
        fn: str,
        filtered_variables: EnvironmentVariables,
    ) -> str:
        resolved_fn: str = fn

        for varname, value in filtered_variables.items():
            pos = 0
            pat = r"\b{}\b".format(re.escape(varname))

            while m := re.search(pat, resolved_fn[pos:]):
                start, end = m.span()
                replacement = Expression.pack(value)
                resolved_fn = (
                    resolved_fn[: pos + start] + replacement + resolved_fn[pos + end :]
                )
                pos += start + len(replacement)

        return resolved_fn

    @staticmethod
    def try_running(func: TimeoutFunction[T], timeout_value: float) -> (T | None):
        @timeout(timeout_value)
        def f() -> T:
            return func()

        try:
            result: T = f()
            return result if result else True
        except Exception:
            return None

    @staticmethod
    def try_simplify_expression(expr: ExpressionBuffer) -> None:

        snapshot = expr.body

        simplified_eq = Expression.try_running(
            lambda: Expression.simplify_expression(expr), 3.0
        )

        if simplified_eq is None:
            expr.body = snapshot

    @staticmethod
    def replace_params_with_temp(expr_str: str, params: List[str]) -> str:
        """
        Replaces function parameters in a given expression string with temporary strings.
        Args:
            expr_str (str): The expression string.
            params (List[str]): A list of function parameter names.

        Returns:
            str: The modified expression string with parameters replaced with temporary strings.
        """
        for idx, param in enumerate(params):
            pat = r"\b{}\b".format(param)
            temp_sub_str = "p_p_{}".format(idx)
            expr_str = re.sub(pat, temp_sub_str, expr_str)
        return expr_str

    @staticmethod
    def replace_temp_with_params(expr_str: str, params: List[str]) -> str:
        """
        Replaces temporary strings in a given expression string with function parameters.
        Args:
            expr_str (str): The expression string.
            params (List[str]): A list of function parameter names.

        Returns:
            str: The modified expression string with temporary strings replaced with function parameters.
        """
        for idx, param in enumerate(params):
            expr_str = expr_str.replace(f"p_{{p_{{{idx}}}}}", param)
        return expr_str

    @staticmethod
    def replace_latex_parens(expr_str: str) -> str:
        """
        Removes LaTeX-style parentheses from a mathematical expression string.

        Args:
            expr_str (str): The mathematical expression string to remove parentheses from.

        Returns:
            str: The modified expression string with parentheses removed.

        Note:
            It leaves behind the '(' and ')' from left and right respectively.
        """
        # Remove \left and \right commands from expression string
        expr_str = re.sub(r"\\(left|right)", "", expr_str)
        return expr_str

    @staticmethod
    def simplify_latex_expression(expr_str: str) -> str:
        """
        Simplifies a LaTeX expression string and returns the simplified expression as a LaTeX string.

        Args:
            expr_str (str): The LaTeX expression to simplify.

        Returns:
            str: The simplified LaTeX expression as a string.
        """
        expr = parse_latex(expr_str)
        simplified_expr = simplify(expr)
        simplified_latex_expr = str(latex(simplified_expr))
        return simplified_latex_expr

    @staticmethod
    def simplify_body(expr: ExpressionBuffer) -> None:
        expr.body = Expression.simplify_latex_expression(expr.body)

    @staticmethod
    def simplify_function_expression(expr: ExpressionBuffer) -> None:
        expr.body = Expression.replace_params_with_temp(expr.body, expr.signature)
        _ = Expression.simplify_body(expr=expr)
        expr.body = Expression.replace_temp_with_params(expr.body, expr.signature)

    @staticmethod
    def simplify_expression(expr: ExpressionBuffer) -> None:

        match expr.expr_type:
            case ExpressionType.FUNCTION:
                Expression.simplify_function_expression(expr)
            case _:
                Expression.simplify_body(expr)


@dataclass
class FunctionExpressionBuffer(ExpressionBuffer):
    name: str
    signature_str: str
    signature: List[str]
    body: str


@dataclass
class StatementExpressionBuffer(ExpressionBuffer):
    body: str


@dataclass
class AssignmentExpressionBuffer(ExpressionBuffer):
    name: str
    body: str
