from .helper_functions import (
    parse_expr,
    get_v_variable,
    get_f_variable,
    checkLazyNew
)
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
from comfy_api.latest import io
import torch
from .Stack import MrmthStack
import copy
from .ParseTree import MrmthParseTree

class StringMathNode(io.ComfyNode):
    """
    Enables math expressions on String values.

    Inputs:
        V: Autogrow string inputs (V0, V1, ...)
        F: Autogrow float inputs (F0, F1, ...)
        Expression: Math expression to evaluate
    """

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="mrmth_ag_StringMathNode",
            category="More math",
            display_name="String math",
            inputs=[
                io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.String.Input("values"), prefix="V", min=1, max=50)),
                io.Autogrow.Input(id="F", template=io.Autogrow.TemplatePrefix(io.Float.Input("float", default=0.0, optional=True, lazy=True, force_input=True), prefix="F", min=1, max=50)),
                io.MultiType.Input(
                    io.String.Input("Expression", default="", multiline=False),
                    types=[io.String,MrmthParseTree],
                    tooltip="Expression to apply on weights",
                ),
                io.Int.Input(id="batching", default=0,min=-1),
                io.Boolean.Input(
                    id="remember_stack",
                    default=False,
                    display_name="Remember stack across batch",
                    tooltip=(
                        "If enabled, stack is copied at output leading to changes being remembered during batch operations (node runs multiple times in sucession). If disabled each batch gets it's own copy of the stack."
                    ),
                ),
                MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True)
            ],
            outputs=[
                io.String.Output(is_output_list=True),
                MrmthStack.Output(),
            ],
        )

    @classmethod
    def check_lazy_status(cls, Expression, V, F,batching=0, remember_stack=False,stack={}):
        return checkLazyNew(Expression,V,F)


    @classmethod
    def execute(cls, V, F, Expression, batching=0, remember_stack=False, stack={}):

        variables = {
            "a": V.get("V0",""), "b": V.get("V1",""), "c": V.get("V2",""), "d": V.get("V3",""),
            "w": F.get("F0", 0.0) if F.get("F0") is not None else 0.0,
            "x": F.get("F1", 0.0) if F.get("F1") is not None else 0.0,
            "y": F.get("F2", 0.0) if F.get("F2") is not None else 0.0,
            "z": F.get("F3", 0.0) if F.get("F3") is not None else 0.0
        }

        v_stacked, v_cnt = get_v_variable(V, length_mismatch="do nothing")
        if v_stacked is not None:
             variables["V"] = v_stacked
             variables["Vcnt"] = float(v_cnt)
             variables["V_count"] = float(v_cnt)

        f_stacked, f_cnt = get_f_variable(F)
        if f_stacked is not None:
             variables["F"] = f_stacked
             variables["Fcnt"] = float(f_cnt)
             variables["F_count"] = float(f_cnt)

        for k, val in F.items():
            variables[k] = val if val is not None else 0.0
        for k, val in V.items():
            variables[k] = val if val is not None else ""

        tree = None
        if isinstance(Expression,str):
            tree = parse_expr(Expression)
        else:
            tree = Expression

        visitor = UnifiedMathVisitor(variables, len(V["V0"]),torch.device("cpu"),state_storage=stack)
        result = visitor.visit(tree)

        if batching == -1:
            return (result, stack)

        result = str(result)
        if batching and batching > 0:
            def chunks(s, n):
                """Produce `n`-character chunks from `s`."""
                for start in range(0, len(s), n):
                    yield s[start:start+n]

            stack = stack if remember_stack else copy.deepcopy(stack)
            return (chunks(result,batching), stack)
        else:
            stack = stack if remember_stack else copy.deepcopy(stack)

            return (result, stack)