import re
import random
from comfy_api.latest import ComfyExtension, io

WEB_DIRECTORY = "./web"

class SimpleStringRepository(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="SimpleStringRepository",
            display_name="Simple String Repository",
            category="utils",
            description="Store and select strings from a multiline list. Supports index prefixes, strength, and {a|b|c} random syntax.",
            inputs=[
                io.String.Input("string", multiline=True, default="1. red\n2. blue\n3. green"),
                io.Combo.Input("delimiter", options=[",", " "], default=","),
                io.Int.Input("target1", default=1, min=0, max=100000),
                io.Float.Input("strength1", default=1.0, min=0.0, max=10.0, step=0.1),
                io.Int.Input("random1", default=-1, min=-1, max=100),
                io.Int.Input("target2", default=0, min=0, max=100000),
                io.Float.Input("strength2", default=1.0, min=0.0, max=10.0, step=0.1),
                io.Int.Input("random2", default=-1, min=-1, max=100),
                io.Int.Input("target3", default=0, min=0, max=100000),
                io.Float.Input("strength3", default=1.0, min=0.0, max=10.0, step=0.1),
                io.Int.Input("random3", default=-1, min=-1, max=100),
                io.Int.Input("target4", default=0, min=0, max=100000),
                io.Float.Input("strength4", default=1.0, min=0.0, max=10.0, step=0.1),
                io.Int.Input("random4", default=-1, min=-1, max=100),
                io.Int.Input("target5", default=0, min=0, max=100000),
                io.Float.Input("strength5", default=1.0, min=0.0, max=10.0, step=0.1),
                io.Int.Input("random5", default=-1, min=-1, max=100),
            ],
            outputs=[
                io.String.Output(display_name="parsed_string"),
            ],
        )

    @classmethod
    def fingerprint_inputs(cls, **kwargs):
        # Force re-execution when any random selector is set to -1 (auto-random)
        for i in range(1, 6):
            r = kwargs.get(f"random{i}", -1)
            t = kwargs.get(f"target{i}", 0)
            if t and r == -1:
                return float("nan")  # always re-run
        return None  # normal caching

    @classmethod
    def execute(cls, string, delimiter, target1, strength1, random1,
                target2, strength2, random2, target3, strength3, random3,
                target4, strength4, random4, target5, strength5, random5):
        inputs = [
            {"t": target1, "s": strength1, "r": random1},
            {"t": target2, "s": strength2, "r": random2},
            {"t": target3, "s": strength3, "r": random3},
            {"t": target4, "s": strength4, "r": random4},
            {"t": target5, "s": strength5, "r": random5},
        ]
        result = cls._process(string, delimiter, inputs)
        return io.NodeOutput(result)

    @staticmethod
    def _process(string, delimiter, inputs):
        lines = string.splitlines()
        repo = {}
        idx = 0
        for line in lines:
            line = line.strip()
            if not line:
                continue
            idx += 1
            parts = line.split(".", 1)
            if len(parts) == 2 and parts[0].isdigit():
                key = int(parts[0])
                text = parts[1].strip()
            else:
                key = idx
                text = line
            if key not in repo:
                repo[key] = text

        parts_out = []
        for item in inputs:
            t = item["t"]
            if t not in repo or t == 0:
                continue
            base = repo[t]

            # Random syntax {a|b|c}
            if re.fullmatch(r"^\{.+\}$", base):
                options = base[1:-1].split("|")
                r = item["r"]
                if r == -1:
                    r = random.randint(0, len(options) - 1)
                if 0 <= r < len(options):
                    base = options[r]
                else:
                    continue

            s = item["s"]
            if s == 0:
                continue
            if s != 1.0:
                base = f"({base}:{round(s, 2)})"

            parts_out.append(base)

        return delimiter.join(parts_out)


class SimpleStringRepositoryCompact(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="SimpleStringRepositoryCompact",
            display_name="Simple String Repository Compact",
            category="utils",
            description="Compact version. Each target is a string like '2,1.2,-1' (index,strength,random).",
            inputs=[
                io.String.Input("string", multiline=True, default="1. red\n2. blue\n3. green"),
                io.Combo.Input("delimiter", options=[",", " "], default=","),
                io.String.Input("target1", default="1,1,-1"),
                io.String.Input("target2", default="0,1,-1"),
                io.String.Input("target3", default="0,1,-1"),
                io.String.Input("target4", default="0,1,-1"),
                io.String.Input("target5", default="0,1,-1"),
            ],
            outputs=[
                io.String.Output(display_name="parsed_string"),
            ],
        )

    @classmethod
    def fingerprint_inputs(cls, **kwargs):
        # Force re-execution when any random selector is set to -1 (auto-random)
        for i in range(1, 6):
            r = kwargs.get(f"random{i}", -1)
            t = kwargs.get(f"target{i}", 0)
            if t and r == -1:
                return float("nan")  # always re-run
        return None  # normal caching

    @classmethod
    def execute(cls, string, delimiter, target1, target2, target3, target4, target5):
        raw = [target1, target2, target3, target4, target5]
        inputs = []
        for tstr in raw:
            parts = [p.strip() for p in tstr.split(",")]
            try:
                t = int(parts[0]) if parts else 0
                s = float(parts[1]) if len(parts) > 1 else 1.0
                r = int(parts[2]) if len(parts) > 2 else -1
            except ValueError:
                continue
            inputs.append({"t": t, "s": s, "r": r})

        result = SimpleStringRepository._process(string, delimiter, inputs)
        return io.NodeOutput(result)


class SimpleStringRepositoryExtension(ComfyExtension):
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            SimpleStringRepository,
            SimpleStringRepositoryCompact,
        ]


async def comfy_entrypoint() -> SimpleStringRepositoryExtension:
    return SimpleStringRepositoryExtension()