from __future__ import annotations

from comfy_api.latest import ComfyExtension, io
from comfy_api.input import CurveInput, MonotoneCubicCurve, LinearCurve
from typing_extensions import override


class CurveEditor(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="CurveEditor",
            display_name="Curve Editor",
            category="utils",
            inputs=[
                io.Curve.Input("curve"),
                io.Histogram.Input("histogram", optional=True),
            ],
            outputs=[
                io.Curve.Output("curve"),
            ],
        )

    @classmethod
    def execute(cls, curve, histogram=None) -> io.NodeOutput:
        if isinstance(curve, CurveInput):
            result = curve
        else:
            raw_points = curve["points"] if isinstance(curve, dict) else curve
            points = [(float(x), float(y)) for x, y in raw_points]
            interpolation = curve.get("interpolation", "monotone_cubic") if isinstance(curve, dict) else "monotone_cubic"
            if interpolation == "linear":
                result = LinearCurve(points)
            else:
                result = MonotoneCubicCurve(points)

        ui = {}
        if histogram is not None:
            ui["histogram"] = histogram if isinstance(histogram, list) else list(histogram)

        return io.NodeOutput(result, ui=ui) if ui else io.NodeOutput(result)


class CurveExtension(ComfyExtension):
    @override
    async def get_node_list(self):
        return [CurveEditor]


async def comfy_entrypoint():
    return CurveExtension()