import json
import os
import random
import re
from pathlib import Path
from typing import Any


USER_DIR = Path(os.getcwd()) / "user" / "default"
CONFIG_DIR = USER_DIR / "ComfyUI-Complex-Prompt"
CONFIG_PATH = CONFIG_DIR / "config.json"
DEFAULT_CONFIG: dict[str, Any] = {}
CONDITION_FUNCTIONS = {
    "float": float,
    "int": int,
    "str": str,
}


def log(message: str) -> None:
    print(f"[ComfyUI-Complex-Prompt]: {message}")


def read_json(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as file:
        return json.load(file)


def write_json(path: Path, data: dict[str, Any]) -> None:
    with path.open("w", encoding="utf-8") as file:
        json.dump(data, file, ensure_ascii=False, indent=2)
        file.write("\n")


def load_config() -> dict[str, Any]:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)

    if not CONFIG_PATH.exists():
        write_json(CONFIG_PATH, DEFAULT_CONFIG)
        log(f"Created config file at {CONFIG_PATH}")
        return DEFAULT_CONFIG.copy()

    log(f"Reading config from {CONFIG_PATH}")
    return read_json(CONFIG_PATH)


CONFIG = load_config()

ARTEMKO7V_COMPLEX_PROMPT_VARS = "ArtemKo7vComplexPromptVars"
VAR_PATTERN = re.compile(r"\$([A-Za-z_][A-Za-z0-9_]*)")
DYNAMIC_VARIANT_PATTERN = re.compile(r"\{[^{}]*\}")

_PROMPT_GENERATOR = None


class ConditionValue:
    def __init__(self, value: Any):
        self.value = value

    def __str__(self) -> str:
        return str(self.value)

    def __repr__(self) -> str:
        return repr(str(self))

    def __bool__(self) -> bool:
        return bool(str(self))

    def __int__(self) -> int:
        return int(str(self))

    def __float__(self) -> float:
        return float(str(self))

    def __eq__(self, other: Any) -> bool:
        return str(self) == str(self.unwrap(other))

    def __ne__(self, other: Any) -> bool:
        return not self == other

    def __lt__(self, other: Any) -> bool:
        return self.compare(other, lambda left, right: left < right)

    def __le__(self, other: Any) -> bool:
        return self.compare(other, lambda left, right: left <= right)

    def __gt__(self, other: Any) -> bool:
        return self.compare(other, lambda left, right: left > right)

    def __ge__(self, other: Any) -> bool:
        return self.compare(other, lambda left, right: left >= right)

    def unwrap(self, other: Any) -> Any:
        if isinstance(other, ConditionValue):
            return other.value

        return other

    def compare(self, other: Any, operation):
        other = self.unwrap(other)

        try:
            return operation(float(self), float(other))
        except (TypeError, ValueError):
            return operation(str(self), str(other))


def get_prompt_generator():
    global _PROMPT_GENERATOR

    if _PROMPT_GENERATOR is None:
        try:
            from dynamicprompts.generators import RandomPromptGenerator
        except ImportError as error:
            raise ImportError(
                "ComfyUI-Complex-Prompt requires the 'dynamicprompts' package. "
                "Install this extension's requirements.txt dependencies."
            ) from error

        _PROMPT_GENERATOR = RandomPromptGenerator()

    return _PROMPT_GENERATOR


def mask_vars(prompt: str) -> tuple[str, dict[str, str]]:
    placeholders: dict[str, str] = {}

    def replace(match: re.Match[str]) -> str:
        placeholder = f"AK7VCOMPLEXPROMPTVAR{len(placeholders)}TOKEN"
        placeholders[placeholder] = match.group(0)
        return placeholder

    return VAR_PATTERN.sub(replace, prompt), placeholders


def restore_vars(prompt: str, placeholders: dict[str, str]) -> str:
    for placeholder, variable in placeholders.items():
        prompt = prompt.replace(placeholder, variable)

    return prompt


def has_dynamic_prompt_syntax(prompt: str) -> bool:
    return bool(DYNAMIC_VARIANT_PATTERN.search(prompt)) or "__" in prompt


def generate_dynamic_prompt(prompt: str | None, seed: int) -> str:
    if prompt in (None, ""):
        return ""

    masked_prompt, placeholders = mask_vars(prompt)
    if not has_dynamic_prompt_syntax(masked_prompt):
        return prompt

    generator = get_prompt_generator()
    # dynamicprompts treats the integer 0 as an omitted seed; [0] is explicit.
    prompts = generator.generate(masked_prompt, 1, seeds=[0] if seed == 0 else seed)
    generated_prompt = prompts[0] if prompts else ""
    return restore_vars(generated_prompt, placeholders)



def apply_vars(prompt: str, vars: dict[str, str] | None = None) -> str:
    if not vars:
        return prompt

    def replace(match: re.Match[str]) -> str:
        variable_name = match.group(1)
        if variable_name not in vars:
            return match.group(0)

        return str(vars[variable_name])

    return VAR_PATTERN.sub(replace, prompt)


def resolve_prompt(
    prompt: str | None,
    seed: int,
    vars: dict[str, Any] | None = None,
) -> str:
    """Expand choices, then recursively resolve the selected variable values."""
    def resolve(value: str | None, resolving: frozenset[str]) -> str:
        generated = generate_dynamic_prompt(value, seed)
        if not vars:
            return generated

        def replace(match: re.Match[str]) -> str:
            name = match.group(1)
            # Keep unknown and cyclic references visible without expanding forever.
            if name not in vars or name in resolving:
                return match.group(0)

            return resolve(str(vars[name]), resolving | {name})

        return VAR_PATTERN.sub(replace, generated)

    return resolve(prompt, frozenset())


def is_json_variable_value(value: Any) -> bool:
    return isinstance(value, (int, float, str)) and not isinstance(value, bool)


def normalize_json_object_text(json_text: str) -> str:
    json_text = json_text.strip()
    if not json_text:
        return "{}"

    if json_text.startswith(("{", "[")):
        return json_text

    return "{" + json_text + "}"


def parse_json_vars(
    json_text: str,
    seed: int,
    vars: dict[str, Any] | None = None,
) -> dict[str, Any]:
    json_text = normalize_json_object_text(json_text)

    try:
        parsed = json.loads(json_text)
    except json.JSONDecodeError as error:
        raise ValueError(f"Invalid JSON: {error.msg}") from error

    if not isinstance(parsed, dict):
        raise ValueError("JSON must be an object with variable names as keys.")

    result = dict(vars or {})
    rng = random.Random(seed)

    for key, value in parsed.items():
        if is_json_variable_value(value):
            result[key] = value
            continue

        if (
            isinstance(value, list)
            and value
            and all(is_json_variable_value(item) for item in value)
        ):
            result[key] = rng.choice(value)

    return result


def parse_choice_map(choices: str) -> dict[str, str]:
    choices = choices.strip()
    if not choices:
        return {}

    if not choices.startswith("{"):
        choices = "{" + choices + "}"

    try:
        parsed = json.loads(choices)
    except json.JSONDecodeError as error:
        raise ValueError(f"Invalid choices JSON: {error.msg}") from error

    if not isinstance(parsed, dict):
        raise ValueError("Choices must be a JSON object.")

    result: dict[str, str] = {}
    for key, value in parsed.items():
        if is_json_variable_value(value):
            result[str(key)] = str(value)

    return result


def normalize_condition(condition: str) -> str:
    normalized = VAR_PATTERN.sub(r"\1", condition)
    normalized = normalized.replace("&&", " and ")
    normalized = normalized.replace("||", " or ")
    normalized = re.sub(r"!(?!=)", " not ", normalized)
    return normalized


def should_set_variable(
    condition: str,
    vars: dict[str, str] | None = None,
) -> bool:
    condition = condition.strip()
    if not condition:
        return True

    try:
        from simpleeval import NameNotDefined, SimpleEval
    except ImportError as error:
        raise ImportError(
            "ComfyUI-Complex-Prompt conditional variables require the "
            "'simpleeval' package. Install this extension's requirements.txt "
            "dependencies."
        ) from error

    evaluator = SimpleEval(
        functions=CONDITION_FUNCTIONS,
        names={
            key: ConditionValue(value)
            for key, value in (vars or {}).items()
            if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key)
        },
    )

    try:
        return bool(evaluator.eval(normalize_condition(condition)))
    except NameNotDefined:
        return False


class ArtemKo7vComplexPrompt:
    CATEGORY = "ArtemKo7v"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("string",)
    FUNCTION = "generate_prompt"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": (
                    "STRING",
                    {
                        "multiline": True,
                        "dynamicPrompts": False,
                    },
                ),
                "seed": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 18446744073709551615,
                        "control_after_generate": "randomize",
                    },
                ),
            },
            "optional": {
                "vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
            },
        }

    def generate_prompt(
        self,
        prompt: str,
        seed: int,
        vars: dict[str, str] | None = None,
    ):
        return (resolve_prompt(prompt, seed, vars),)


class ArtemKo7vComplexPromptSetVariable:
    CATEGORY = "ArtemKo7v"
    RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS, "BOOLEAN")
    RETURN_NAMES = ("vars", "wasSet")
    FUNCTION = "set_variable"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "variable_name": (
                    "STRING",
                    {
                        "default": "",
                    },
                ),
                "value": (
                    "STRING",
                    {
                        "multiline": True,
                        "dynamicPrompts": False,
                    },
                ),
                "condition": (
                    "STRING",
                    {
                        "default": "",
                    },
                ),
                "seed": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 18446744073709551615,
                        "control_after_generate": "randomize",
                    },
                ),
                "trim": (
                    "BOOLEAN",
                    {
                        "default": False,
                    },
                ),
            },
            "optional": {
                "vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
            },
        }

    def set_variable(
        self,
        variable_name: str,
        value: str | None,
        condition: str,
        seed: int,
        trim: bool = False,
        vars: dict[str, str] | None = None,
    ):
        result = dict(vars or {})
        if not should_set_variable(condition, vars):
            return (result, False)

        value = "" if value is None else value
        stored_value = resolve_prompt(value, seed, vars)
        if trim:
            stored_value = stored_value.strip()

        result[variable_name] = stored_value
        return (result, True)


class ArtemKo7vComplexPromptSetVariableByChoice:
    CATEGORY = "ArtemKo7v"
    RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS, "BOOLEAN")
    RETURN_NAMES = ("vars", "wasSet")
    FUNCTION = "set_variable_by_choice"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "variable_name": (
                    "STRING",
                    {
                        "default": "",
                    },
                ),
                "source_variable": (
                    "STRING",
                    {
                        "default": "",
                    },
                ),
                "choices": (
                    "STRING",
                    {
                        "multiline": True,
                        "dynamicPrompts": False,
                        "default": "",
                    },
                ),
                "seed": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 18446744073709551615,
                        "control_after_generate": "randomize",
                    },
                ),
                "trim": (
                    "BOOLEAN",
                    {
                        "default": False,
                    },
                ),
            },
            "optional": {
                "vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
            },
        }

    def set_variable_by_choice(
        self,
        variable_name: str,
        source_variable: str,
        choices: str,
        seed: int,
        trim: bool = False,
        vars: dict[str, str] | None = None,
    ):
        result = dict(vars or {})
        source_variable = source_variable.strip()
        if source_variable.startswith("$"):
            source_variable = source_variable[1:]

        if not source_variable or source_variable not in result:
            return (result, False)

        choice_map = parse_choice_map(choices)
        source_value = str(result[source_variable])
        value = choice_map.get(source_value, choice_map.get("__default__"))
        if value is None:
            return (result, False)

        stored_value = resolve_prompt(value, seed, vars)
        if trim:
            stored_value = stored_value.strip()

        result[variable_name] = stored_value
        return (result, True)


class ArtemKo7vComplexPromptParseJSON:
    CATEGORY = "ArtemKo7v"
    RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS,)
    RETURN_NAMES = ("vars",)
    FUNCTION = "parse_json"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_text": (
                    "STRING",
                    {
                        "multiline": True,
                        "dynamicPrompts": False,
                        "default": "",
                    },
                ),
                "seed": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 18446744073709551615,
                        "control_after_generate": "randomize",
                    },
                ),
            },
            "optional": {
                "vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
            },
        }

    def parse_json(
        self,
        json_text: str,
        seed: int,
        vars: dict[str, Any] | None = None,
    ):
        return (parse_json_vars(json_text, seed, vars),)


class ArtemKo7vComplexPromptCombineVars:
    """Combines vars inputs in order, with later inputs taking precedence."""

    CATEGORY = "ArtemKo7v"
    MAX_VAR_INDEX = 64
    RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS,)
    RETURN_NAMES = ("vars",)
    FUNCTION = "combine_vars"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vars_1": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
            },
        }

    def combine_vars(self, vars_1: dict[str, Any], **kwargs):
        result = dict(vars_1 or {})
        for index in range(2, self.MAX_VAR_INDEX + 1):
            vars = kwargs.get(f"vars_{index}")
            if vars:
                result.update(vars)

        return (result,)


class ArtemKo7vComplexPromptEmptyString:
    CATEGORY = "ArtemKo7v"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("string",)
    FUNCTION = "get_empty_string"

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {}}

    def get_empty_string(self):
        return ("",)


NODE_CLASS_MAPPINGS = {
    "ArtemKo7vComplexPrompt": ArtemKo7vComplexPrompt,
    "ArtemKo7vComplexPromptSetVariable": ArtemKo7vComplexPromptSetVariable,
    "ArtemKo7vComplexPromptSetVariableByChoice": (
        ArtemKo7vComplexPromptSetVariableByChoice
    ),
    "ArtemKo7vComplexPromptParseJSON": ArtemKo7vComplexPromptParseJSON,
    "ArtemKo7vComplexPromptCombineVars": ArtemKo7vComplexPromptCombineVars,
    "ArtemKo7vComplexPromptEmptyString": ArtemKo7vComplexPromptEmptyString,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ArtemKo7vComplexPrompt": "Complex Prompt",
    "ArtemKo7vComplexPromptSetVariable": "Complex Prompt Set Variable",
    "ArtemKo7vComplexPromptSetVariableByChoice": (
        "Complex Prompt Set Variable By Choice"
    ),
    "ArtemKo7vComplexPromptParseJSON": "Complex Prompt Parse JSON",
    "ArtemKo7vComplexPromptCombineVars": "Complex Prompt Combine Vars",
    "ArtemKo7vComplexPromptEmptyString": "Complex Prompt Empty String",
}