from typing import Any, Union
import uuid

from lmnr.cli.parser.nodes import Handle
from lmnr.cli.parser.nodes.code import CodeNode
from lmnr.cli.parser.nodes.condition import ConditionNode
from lmnr.cli.parser.nodes.input import InputNode
from lmnr.cli.parser.nodes.json_extractor import JsonExtractorNode
from lmnr.cli.parser.nodes.llm import LLMNode
from lmnr.cli.parser.nodes.output import OutputNode
from lmnr.cli.parser.nodes.router import Route, RouterNode
from lmnr.cli.parser.nodes.semantic_search import (
    Dataset,
    SemanticSearchNode,
)
from lmnr.types import NodeInput, ChatMessage


def node_input_from_json(json_val: Any) -> NodeInput:
    if isinstance(json_val, str):
        return json_val
    elif isinstance(json_val, list):
        return [ChatMessage.model_validate(msg) for msg in json_val]
    else:
        raise ValueError(f"Invalid NodeInput value: {json_val}")


Node = Union[
    InputNode,
    OutputNode,
    ConditionNode,
    LLMNode,
    RouterNode,
    SemanticSearchNode,
    CodeNode,
    JsonExtractorNode,
]


def node_from_dict(node_dict: dict) -> Node:
    if node_dict["type"] == "Input":
        return InputNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            input=node_input_from_json(node_dict["input"]),
        )
    elif node_dict["type"] == "Output":
        return OutputNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            inputs_mappings={
                uuid.UUID(k): uuid.UUID(v)
                for k, v in node_dict["inputsMappings"].items()
            },
        )
    elif node_dict["type"] == "Condition":
        return ConditionNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            inputs_mappings={
                uuid.UUID(k): uuid.UUID(v)
                for k, v in node_dict["inputsMappings"].items()
            },
            condition=node_dict["condition"],
        )
    elif node_dict["type"] == "LLM":
        return LLMNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
            dynamic_inputs=[
                Handle.from_dict(handle) for handle in node_dict["dynamicInputs"]
            ],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            inputs_mappings={
                uuid.UUID(k): uuid.UUID(v)
                for k, v in node_dict["inputsMappings"].items()
            },
            prompt=node_dict["prompt"],
            model=node_dict["model"],
            model_params=(
                node_dict["modelParams"] if "modelParams" in node_dict else None
            ),
            stream=False,
            structured_output_enabled=node_dict.get("structuredOutputEnabled", False),
            structured_output_max_retries=node_dict.get(
                "structuredOutputMaxRetries", 0
            ),
            structured_output_schema=node_dict.get("structuredOutputSchema", None),
            structured_output_schema_target=node_dict.get(
                "structuredOutputSchemaTarget", None
            ),
        )
    elif node_dict["type"] == "Router":
        return RouterNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            inputs_mappings={
                uuid.UUID(k): uuid.UUID(v)
                for k, v in node_dict["inputsMappings"].items()
            },
            routes=[Route(name=route["name"]) for route in node_dict["routes"]],
            has_default_route=node_dict["hasDefaultRoute"],
        )
    elif node_dict["type"] == "SemanticSearch":
        return SemanticSearchNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            inputs_mappings={
                uuid.UUID(k): uuid.UUID(v)
                for k, v in node_dict["inputsMappings"].items()
            },
            limit=node_dict["limit"],
            threshold=node_dict["threshold"],
            template=node_dict["template"],
            datasets=[Dataset.from_dict(ds) for ds in node_dict["datasets"]],
        )
    elif node_dict["type"] == "Code":
        return CodeNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            inputs_mappings={
                uuid.UUID(k): uuid.UUID(v)
                for k, v in node_dict["inputsMappings"].items()
            },
            code=node_dict["code"],
            fn_name=node_dict["fnName"],
        )
    elif node_dict["type"] == "JsonExtractor":
        return JsonExtractorNode(
            id=uuid.UUID(node_dict["id"]),
            name=node_dict["name"],
            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
            inputs_mappings={
                uuid.UUID(k): uuid.UUID(v)
                for k, v in node_dict["inputsMappings"].items()
            },
            template=node_dict["template"],
        )
    else:
        raise ValueError(f"Node type {node_dict['type']} not supported")
