from functools import lru_cache
from itertools import islice
import json


def _clone_container(value):
    if isinstance(value, dict):
        return dict(value)
    if isinstance(value, list):
        return list(value)
    return value


@lru_cache(maxsize=256)
def _parse_json_path(path):
    if not path.strip():
        return ()

    segments = []
    for raw_segment in path.split("."):
        segment = raw_segment.strip()
        if not segment:
            continue
        if segment.isdigit():
            segments.append(int(segment))
        else:
            segments.append(segment)
    return tuple(segments)


def _get_json_path_value(data, path):
    current = data
    for segment in _parse_json_path(path):
        if isinstance(segment, int):
            if not isinstance(current, list) or segment >= len(current):
                return None
            current = current[segment]
            continue

        if not isinstance(current, dict) or segment not in current:
            return None
        current = current[segment]
    return current


def _has_json_path(data, path):
    current = data
    for segment in _parse_json_path(path):
        if isinstance(segment, int):
            if not isinstance(current, list) or segment >= len(current):
                return False
            current = current[segment]
            continue

        if not isinstance(current, dict) or segment not in current:
            return False
        current = current[segment]
    return True


def _build_missing_container(next_segment):
    return [] if isinstance(next_segment, int) else {}


def _set_json_path_value_recursive(current, segments, value):
    if not segments:
        return value

    segment = segments[0]
    is_last = len(segments) == 1

    if isinstance(segment, int):
        if not isinstance(current, list):
            return None

        result = list(current)
        while segment >= len(result):
            result.append(None)

        if is_last:
            result[segment] = value
            return result

        next_current = result[segment]
        if not isinstance(next_current, (dict, list)):
            next_current = _build_missing_container(segments[1])
        updated = _set_json_path_value_recursive(next_current, segments[1:], value)
        if updated is None:
            return None
        result[segment] = updated
        return result

    if not isinstance(current, dict):
        return None

    result = dict(current)
    if is_last:
        result[segment] = value
        return result

    next_current = result.get(segment)
    if not isinstance(next_current, (dict, list)):
        next_current = _build_missing_container(segments[1])
    updated = _set_json_path_value_recursive(next_current, segments[1:], value)
    if updated is None:
        return None
    result[segment] = updated
    return result


def _set_json_path_value(data, path, value):
    if not isinstance(data, dict):
        return None

    segments = _parse_json_path(path)
    if not segments:
        return value

    return _set_json_path_value_recursive(data, segments, value)


def _remove_json_path_value_recursive(current, segments):
    if not segments:
        return _clone_container(current)

    segment = segments[0]
    is_last = len(segments) == 1

    if isinstance(segment, int):
        if not isinstance(current, list) or segment >= len(current):
            return None

        result = list(current)
        if is_last:
            del result[segment]
            return result

        updated = _remove_json_path_value_recursive(result[segment], segments[1:])
        if updated is None:
            return None
        result[segment] = updated
        return result

    if not isinstance(current, dict) or segment not in current:
        return None

    result = dict(current)
    if is_last:
        del result[segment]
        return result

    updated = _remove_json_path_value_recursive(result[segment], segments[1:])
    if updated is None:
        return None
    result[segment] = updated
    return result


def _remove_json_path_value(data, path):
    segments = _parse_json_path(path)
    if not segments:
        return _clone_container(data)
    return _remove_json_path_value_recursive(data, segments)


def _merge_json_objects(base_object, overlay_object):
    if not isinstance(base_object, dict) or not isinstance(overlay_object, dict):
        return None

    merged = dict(base_object)
    for key, value in overlay_object.items():
        base_value = merged.get(key)
        if isinstance(base_value, dict) and isinstance(value, dict):
            merged[key] = _merge_json_objects(base_value, value)
        else:
            merged[key] = _clone_container(value)
    return merged


def _flatten_json_value(value, prefix="", result=None):
    if result is None:
        result = {}

    if isinstance(value, dict):
        if not value and prefix:
            result[prefix] = {}
            return result
        for key, nested_value in value.items():
            next_prefix = f"{prefix}.{key}" if prefix else str(key)
            _flatten_json_value(nested_value, next_prefix, result)
        return result

    if isinstance(value, list):
        if not value and prefix:
            result[prefix] = []
            return result
        for index, nested_value in enumerate(value):
            next_prefix = f"{prefix}.{index}" if prefix else str(index)
            _flatten_json_value(nested_value, next_prefix, result)
        return result

    result[prefix] = value
    return result


def _unflatten_json_object(flat_object):
    if not isinstance(flat_object, dict):
        return None

    result = {}
    for path, value in flat_object.items():
        if not isinstance(path, str):
            return None
        result = _set_json_path_value(result, path, value)
        if result is None:
            return None
    return result


def _resolve_iteration_index(index, mode, total_items):
    if total_items <= 0:
        return -1

    if mode == "incr":
        index += 1
    elif mode == "decr":
        index -= 1

    return max(0, min(index, total_items - 1))


class DumasJSONStringToObjectNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "convert_string_to_json"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_string": ("STRING", {"multiline": True}),
            }
        }

    def convert_string_to_json(self, json_string):
        try:
            json_object = json.loads(json_string)
            return (json_object,)
        except json.JSONDecodeError as exc:
            print(f"[DumasNodes] Error decoding JSON: {exc}")
            return (None,)


class DumasJSONObjectToStringNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("STRING",)
    FUNCTION = "convert_object_to_string"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_object": ("JSON",),
                "pretty": ("BOOLEAN", {"default": True}),
                "sort_keys": ("BOOLEAN", {"default": False}),
            }
        }

    def convert_object_to_string(self, json_object, pretty, sort_keys):
        indent = 2 if pretty else None
        separators = None if pretty else (",", ":")
        return (json.dumps(json_object, indent=indent, sort_keys=sort_keys, separators=separators),)


class DumasJSONGetValueNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "get_value"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_object": ("JSON",),
                "path": ("STRING", {"default": "", "multiline": False}),
            }
        }

    def get_value(self, json_object, path):
        return (_get_json_path_value(json_object, path),)


class DumasJSONSetValueNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "set_value"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_object": ("JSON",),
                "path": ("STRING", {"default": "", "multiline": False}),
                "value_json": ("STRING", {"multiline": True, "default": "null"}),
            }
        }

    def set_value(self, json_object, path, value_json):
        try:
            value = json.loads(value_json)
        except json.JSONDecodeError as exc:
            print(f"[DumasNodes] Error decoding JSON value: {exc}")
            return (None,)

        return (_set_json_path_value(json_object, path, value),)


class DumasJSONHasKeyNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("BOOLEAN",)
    FUNCTION = "has_key"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_object": ("JSON",),
                "path": ("STRING", {"default": "", "multiline": False}),
            }
        }

    def has_key(self, json_object, path):
        return (_has_json_path(json_object, path),)


class DumasJSONRemoveKeyNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "remove_key"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_object": ("JSON",),
                "path": ("STRING", {"default": "", "multiline": False}),
            }
        }

    def remove_key(self, json_object, path):
        return (_remove_json_path_value(json_object, path),)


class DumasJSONPickFieldsNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "pick_fields"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_object": ("JSON",),
                "paths": ("STRING", {"multiline": True, "default": ""}),
            }
        }

    def pick_fields(self, json_object, paths):
        result = {}
        for raw_path in paths.splitlines():
            path = raw_path.strip()
            if not path or not _has_json_path(json_object, path):
                continue
            value = _get_json_path_value(json_object, path)
            result = _set_json_path_value(result, path, value)
            if result is None:
                return (None,)
        return (result,)


class DumasJSONMergeObjectsNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "merge_objects"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "base_object": ("JSON",),
                "overlay_object": ("JSON",),
            }
        }

    def merge_objects(self, base_object, overlay_object):
        return (_merge_json_objects(base_object, overlay_object),)


class DumasJSONKeysNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON", "INT")
    RETURN_NAMES = ("keys", "count")
    FUNCTION = "list_keys"

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

    def list_keys(self, json_object):
        if not isinstance(json_object, dict):
            return (None, 0)
        keys = list(json_object.keys())
        return (keys, len(keys))


class DumasJSONArrayLengthNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("INT",)
    FUNCTION = "get_length"

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

    def get_length(self, json_array):
        if not isinstance(json_array, list):
            return (0,)
        return (len(json_array),)


class DumasJSONArrayAppendNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "append_item"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_array": ("JSON",),
                "value_json": ("STRING", {"multiline": True, "default": "null"}),
            }
        }

    def append_item(self, json_array, value_json):
        if not isinstance(json_array, list):
            return (None,)
        try:
            value = json.loads(value_json)
        except json.JSONDecodeError as exc:
            print(f"[DumasNodes] Error decoding JSON value: {exc}")
            return (None,)
        result = list(json_array)
        result.append(value)
        return (result,)


class DumasJSONArraySliceNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "slice_array"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_array": ("JSON",),
                "start": ("INT", {"default": 0, "step": 1}),
                "end": ("INT", {"default": 0, "step": 1}),
                "step": ("INT", {"default": 1, "step": 1, "min": 1}),
            }
        }

    def slice_array(self, json_array, start, end, step):
        if not isinstance(json_array, list):
            return (None,)
        actual_end = None if end == 0 else end
        return (json_array[start:actual_end:step],)


class DumasJSONArrayIteratorNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON", "INT", "INT")
    RETURN_NAMES = ("item", "current_index", "total_items")
    FUNCTION = "iterate"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_input": ("JSON",),
                "index": ("INT", {"default": 0, "min": 0, "step": 1}),
                "mode": (["fixed", "incr", "decr"], {"default": "fixed"}),
            }
        }

    def iterate(self, json_input, index, mode):
        if not isinstance(json_input, list):
            return (None, -1, 0)

        total_items = len(json_input)
        if total_items == 0:
            return (None, -1, 0)

        current_index = _resolve_iteration_index(index, mode, total_items)
        return (json_input[current_index], current_index, total_items)


class DumasJSONObjectIteratorNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("STRING", "JSON", "INT", "INT")
    RETURN_NAMES = ("key", "value", "current_index", "total_items")
    FUNCTION = "iterate"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_input": ("JSON",),
                "index": ("INT", {"default": 0, "min": 0, "step": 1}),
                "mode": (["fixed", "incr", "decr"], {"default": "fixed"}),
            }
        }

    def iterate(self, json_input, index, mode):
        if not isinstance(json_input, dict):
            return (None, None, -1, 0)

        total_items = len(json_input)
        if total_items == 0:
            return (None, None, -1, 0)

        current_index = _resolve_iteration_index(index, mode, total_items)
        key, value = next(islice(json_input.items(), current_index, None))
        return (key, value, current_index, total_items)


class DumasJSONFlattenNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "flatten_json"

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

    def flatten_json(self, json_input):
        return (_flatten_json_value(json_input),)


class DumasJSONUnflattenNode:
    CATEGORY = "Dumas/JSON"
    RETURN_TYPES = ("JSON",)
    FUNCTION = "unflatten_json"

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

    def unflatten_json(self, flat_json_object):
        return (_unflatten_json_object(flat_json_object),)


NODE_CLASS_MAPPINGS = {
    "DumasJSONStringToObject": DumasJSONStringToObjectNode,
    "DumasJSONObjectToString": DumasJSONObjectToStringNode,
    "DumasJSONGetValue": DumasJSONGetValueNode,
    "DumasJSONSetValue": DumasJSONSetValueNode,
    "DumasJSONHasKey": DumasJSONHasKeyNode,
    "DumasJSONRemoveKey": DumasJSONRemoveKeyNode,
    "DumasJSONPickFields": DumasJSONPickFieldsNode,
    "DumasJSONMergeObjects": DumasJSONMergeObjectsNode,
    "DumasJSONKeys": DumasJSONKeysNode,
    "DumasJSONArrayLength": DumasJSONArrayLengthNode,
    "DumasJSONArrayAppend": DumasJSONArrayAppendNode,
    "DumasJSONArraySlice": DumasJSONArraySliceNode,
    "DumasJSONArrayIterator": DumasJSONArrayIteratorNode,
    "DumasJSONObjectIterator": DumasJSONObjectIteratorNode,
    "DumasJSONFlatten": DumasJSONFlattenNode,
    "DumasJSONUnflatten": DumasJSONUnflattenNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DumasJSONStringToObject": "Dumas JSON String to Object",
    "DumasJSONObjectToString": "Dumas JSON Object to String",
    "DumasJSONGetValue": "Dumas JSON Get Value",
    "DumasJSONSetValue": "Dumas JSON Set Value",
    "DumasJSONHasKey": "Dumas JSON Has Key",
    "DumasJSONRemoveKey": "Dumas JSON Remove Key",
    "DumasJSONPickFields": "Dumas JSON Pick Fields",
    "DumasJSONMergeObjects": "Dumas JSON Merge Objects",
    "DumasJSONKeys": "Dumas JSON Keys",
    "DumasJSONArrayLength": "Dumas JSON Array Length",
    "DumasJSONArrayAppend": "Dumas JSON Array Append",
    "DumasJSONArraySlice": "Dumas JSON Array Slice",
    "DumasJSONArrayIterator": "Dumas JSON Array Iterator",
    "DumasJSONObjectIterator": "Dumas JSON Object Iterator",
    "DumasJSONFlatten": "Dumas JSON Flatten",
    "DumasJSONUnflatten": "Dumas JSON Unflatten",
}