from hashlib import sha512
import random
from typing import Any

from pydantic import BaseModel, Field, model_validator, validator, model_serializer

ComfyParameterKey = list[str]

FAL_OUTPUT_TYPE_NODES = [
    "SaveImage_fal",
]


class NodeConnection(BaseModel):
    from_node_id: str
    output_index: int

    @model_serializer
    def ser_model(self) -> list[str | int]:
        return [self.from_node_id, self.output_index]


PrimitiveInputType = str | int | float | bool | None
NodeInputType = PrimitiveInputType | NodeConnection


class NodeData(BaseModel):
    inputs: dict[str, NodeInputType]
    class_type: str

    def __init__(self, **kwargs):
        # if input type is list, convert it to a node input type
        for key, value in kwargs.get("inputs", {}).items():
            if isinstance(value, list):
                if len(value) != 2:
                    raise ValueError(
                        f"The length of the list should be 2. Got: {len(value)}"
                    )

                kwargs["inputs"][key] = NodeConnection(
                    from_node_id=value[0], output_index=value[1]
                )

        super().__init__(**kwargs)


class FalFile(BaseModel):
    key: ComfyParameterKey = Field(
        description="The key of the parameter that the file is associated with."
    )
    url: str = Field(description="The URL of the file.")


class FalInputDevInfo(BaseModel):
    key: ComfyParameterKey = Field(
        description="The key of the parameter that the input is associated with."
    )
    class_type: str = Field(
        description="The class type of the input node.",
    )


PromptType = dict[str, NodeData]
FalInputsType = dict[str, PrimitiveInputType]
FalInputDevInfoType = dict[str, FalInputDevInfo]


class ComfyInput(BaseModel):
    prompt: PromptType = Field(
        description="The prompt for the workflow.",
    )
    fal_inputs: FalInputsType | None = Field(
        description="""
            Inputs for the workflow. The keys are the names of the inputs and the values
            are the values of the inputs. The inputs will be passed to the workflow.
            The input substitution will be done based on the keys provided in the
            `fal_inputs_dev_info` field.
        """,
        default_factory=dict,
    )
    fal_inputs_dev_info: FalInputDevInfoType | None = Field(
        default_factory=dict,
        description="""
            Developer information for the FAL inputs.This will be used for
            various purposes like debugging, logging, parameter substitution, etc.
        """,
    )
    extra_data: dict | None = Field(
        default_factory=dict, description="Extra data for the workflow."
    )
    disable_metadata: bool = Field(
        default=False, description="Disable saving prompt metadata in files."
    )

    @validator("prompt")
    def validate_prompt(cls, v: PromptType):
        validate_fal_output_node_unique(v)
        return v

    @model_validator(mode="after")
    def validate_inputs(cls, values: "ComfyInput"):
        fal_inputs = values.fal_inputs
        fal_inputs_dev_info = values.fal_inputs_dev_info
        prompt = values.prompt

        # Validate fal_inputs
        if fal_inputs_dev_info:
            seed_fields = set(
                field for field in fal_inputs_dev_info if field.endswith("seed")
            )
            random_seed = random.randint(0, 2**31 - 1)

            for seed_field in seed_fields:
                if not fal_inputs.get(seed_field, None):
                    print(f"Setting random seed {random_seed} for {seed_field}")
                    fal_inputs[seed_field] = random_seed

        # Validate keys match
        if set(fal_inputs.keys()) != set(fal_inputs_dev_info.keys()):
            raise ValueError(
                "The keys in `fal_inputs_dev_info` should be equal to the keys in `fal_inputs`.\n"
                f"Keys in `fal_inputs`: {set(fal_inputs.keys())}\n"
                f"Keys in `fal_inputs_dev_info`: {set(fal_inputs_dev_info.keys())}\n"
            )

        # Validate class types
        if prompt:
            for fal_input in fal_inputs_dev_info.values():
                key = fal_input.key
                class_type = fal_input.class_type

                node_id = key[0]
                node_class_type = prompt[node_id].class_type

                if class_type != node_class_type:
                    raise ValueError(
                        f"The class type of the input node with key '{key}' should be equal to the class type provided in `fal_inputs_dev_info`."
                    )

        return values

    def __hash__(self):
        # Based on https://github.com/pydantic/pydantic/issues/1303
        # I did not benchmark this, but the computation time should be negligible
        return int.from_bytes(
            sha512(f"{self.__class__.__qualname__}::{self.json()}".encode()).digest(),
            byteorder="big",
        )


def validate_fal_output_node_unique(prompt: PromptType):
    output_nodes = [
        node for node in prompt.values() if node.class_type in FAL_OUTPUT_TYPE_NODES
    ]

    output_fields = set()

    for node in output_nodes:
        if "output_name" not in node.inputs:
            raise ValueError("The output nodes should have the 'output_name' field.")

        output_name = node.inputs["output_name"]
        output_fields.add(output_name)

    if len(output_fields) != len(output_nodes):
        raise ValueError("The output names of the fal output nodes should be unique.")


class ComfyOutput(BaseModel):
    outputs: dict[str, dict[str, Any]]
    fal_outputs: dict[str, dict[str, Any]]
    timings: dict[str, float]
