from __future__ import annotations

import asyncio
import inspect
import traceback
import types
from collections.abc import AsyncIterator, Callable, Iterator, Mapping
from enum import Enum
from typing import TYPE_CHECKING, Any

from ag_ui.core import StepFinishedEvent, StepStartedEvent

from lfx.exceptions.component import ComponentBuildError
from lfx.graph.schema import INPUT_COMPONENTS, OUTPUT_COMPONENTS, InterfaceComponentTypes, ResultData
from lfx.graph.utils import UnbuiltObject, UnbuiltResult, emit_build_start_event, log_transaction
from lfx.graph.vertex.param_handler import ParameterHandler
from lfx.interface import initialize
from lfx.interface.listing import lazy_load_dict
from lfx.log.logger import logger
from lfx.schema.artifact import ArtifactType
from lfx.schema.data import Data
from lfx.schema.message import Message
from lfx.schema.properties import Usage
from lfx.schema.schema import INPUT_FIELD_NAME, OutputValue, build_output_logs
from lfx.utils.schemas import ChatOutputResponse
from lfx.utils.util import sync_to_async

if TYPE_CHECKING:
    from uuid import UUID

    from lfx.custom.custom_component.component import Component
    from lfx.events.event_manager import EventManager
    from lfx.graph.edge.base import CycleEdge, Edge
    from lfx.graph.graph.base import Graph
    from lfx.graph.vertex.schema import NodeData

    # Simple log type for tracing
    Log = dict


class VertexStates(str, Enum):
    """Vertex are related to it being active, inactive, or in an error state."""

    ACTIVE = "ACTIVE"
    INACTIVE = "INACTIVE"
    ERROR = "ERROR"


class Vertex:
    def __init__(
        self,
        data: NodeData,
        graph: Graph,
        *,
        base_type: str | None = None,
        is_task: bool = False,
        params: dict | None = None,
    ) -> None:
        # is_external means that the Vertex send or receives data from
        # an external source (e.g the chat)
        self._lock: asyncio.Lock | None = None
        self.will_stream = False
        self.updated_raw_params = False
        self.id: str = data["id"]
        self.base_name = self.id.split("-")[0]
        self.is_state = False
        # TODO: This won't be enough in the long term
        # we need to have a better way to determine if a vertex is an input or an output
        type_strings = [self.id.split("-")[0], data["data"]["type"]]
        self.is_input = any(input_component_name in type_strings for input_component_name in INPUT_COMPONENTS)
        self.is_output = any(output_component_name in type_strings for output_component_name in OUTPUT_COMPONENTS)
        self._is_loop = None
        self.has_session_id = None
        self.custom_component = None
        self.has_external_input = False
        self.has_external_output = False
        self.graph = graph
        self.full_data = data.copy()
        self.base_type: str | None = base_type
        self.outputs: list[dict] = []
        self.parse_data()
        self.built_object: Any = UnbuiltObject()
        self.built_result: Any = None
        self.built = False
        self._successors_ids: list[str] | None = None
        self.artifacts: dict[str, Any] = {}
        self.artifacts_raw: dict[str, Any] | None = {}
        self.artifacts_type: dict[str, str] = {}
        self.steps: list[Callable] = [self._build]
        self.steps_ran: list[Callable] = []
        self.task_id: str | None = None
        self.is_task = is_task
        self.params = params or {}
        self.parent_node_id: str | None = self.full_data.get("parent_node_id")
        self.load_from_db_fields: list[str] = []
        self.parent_is_top_level = False
        self.layer = None
        self.result: ResultData | None = None
        self.results: dict[str, Any] = {}
        self.outputs_logs: dict[str, OutputValue] = {}
        self.logs: dict[str, list[Log]] = {}
        self.has_cycle_edges = False
        try:
            self.is_interface_component = self.vertex_type in InterfaceComponentTypes
        except ValueError:
            self.is_interface_component = False

        self.use_result = False
        self.build_times: list[float] = []
        self.state = VertexStates.ACTIVE
        self.output_names: list[str] = [
            output["name"] for output in self.outputs if isinstance(output, dict) and "name" in output
        ]
        self._incoming_edges: list[CycleEdge] | None = None
        self._outgoing_edges: list[CycleEdge] | None = None

    @property
    def lock(self):
        """Lazy initialization of asyncio.Lock."""
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    @property
    def is_loop(self) -> bool:
        """Check if any output allows looping."""
        if self._is_loop is None:
            self._is_loop = any(output.get("allows_loop", False) for output in self.outputs)
        return self._is_loop

    def set_input_value(self, name: str, value: Any) -> None:
        if self.custom_component is None:
            msg = f"Vertex {self.id} does not have a component instance."
            raise ValueError(msg)
        self.custom_component.set_input_value(name, value)

    def to_data(self):
        return self.full_data

    def add_component_instance(self, component_instance: Component) -> None:
        component_instance.set_vertex(self)
        self.custom_component = component_instance

    def add_result(self, name: str, result: Any) -> None:
        self.results[name] = result

    def set_state(self, state: str) -> None:
        self.state = VertexStates[state]
        if self.state == VertexStates.INACTIVE and self.graph.in_degree_map[self.id] <= 1:
            # If the vertex is inactive and has only one in degree
            # it means that it is not a merge point in the graph
            self.graph.inactivated_vertices.add(self.id)
        elif self.state == VertexStates.ACTIVE and self.id in self.graph.inactivated_vertices:
            self.graph.inactivated_vertices.remove(self.id)

    def is_active(self):
        return self.state == VertexStates.ACTIVE

    @property
    def avg_build_time(self):
        return sum(self.build_times) / len(self.build_times) if self.build_times else 0

    def add_build_time(self, time) -> None:
        self.build_times.append(time)

    def set_result(self, result: ResultData) -> None:
        self.result = result

    def get_built_result(self):
        # If the Vertex.type is a power component
        # then we need to return the built object
        # instead of the result dict
        if self.is_interface_component and not isinstance(self.built_object, UnbuiltObject):
            result = self.built_object
            # if it is not a dict or a string and hasattr model_dump then
            # return the model_dump
            if not isinstance(result, dict | str) and hasattr(result, "content"):
                return result.content
            return result
        if isinstance(self.built_object, str):
            self.built_result = self.built_object

        if isinstance(self.built_result, UnbuiltResult):
            return {}

        return self.built_result if isinstance(self.built_result, dict) else {"result": self.built_result}

    def set_artifacts(self) -> None:
        pass

    @property
    def edges(self) -> list[CycleEdge]:
        return self.graph.get_vertex_edges(self.id)

    @property
    def outgoing_edges(self) -> list[CycleEdge]:
        if self._outgoing_edges is None:
            self._outgoing_edges = [edge for edge in self.edges if edge.source_id == self.id]
        return self._outgoing_edges

    @property
    def incoming_edges(self) -> list[CycleEdge]:
        if self._incoming_edges is None:
            self._incoming_edges = [edge for edge in self.edges if edge.target_id == self.id]
        return self._incoming_edges

    # Get edge connected to an output of a certain name
    def get_incoming_edge_by_target_param(self, target_param: str) -> str | None:
        return next((edge.source_id for edge in self.incoming_edges if edge.target_param == target_param), None)

    @property
    def edges_source_names(self) -> set[str | None]:
        return {edge.source_handle.name for edge in self.edges}

    @property
    def predecessors(self) -> list[Vertex]:
        return self.graph.get_predecessors(self)

    @property
    def successors(self) -> list[Vertex]:
        return self.graph.get_successors(self)

    @property
    def successors_ids(self) -> list[str]:
        return self.graph.successor_map.get(self.id, [])

    def __getstate__(self):
        state = self.__dict__.copy()
        state["_lock"] = None  # Locks are not serializable
        state["built_object"] = None if isinstance(self.built_object, UnbuiltObject) else self.built_object
        state["built_result"] = None if isinstance(self.built_result, UnbuiltResult) else self.built_result
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._lock = asyncio.Lock()  # Reinitialize the lock
        self.built_object = state.get("built_object") or UnbuiltObject()
        self.built_result = state.get("built_result") or UnbuiltResult()

    def set_top_level(self, top_level_vertices: list[str]) -> None:
        self.parent_is_top_level = self.parent_node_id in top_level_vertices

    def parse_data(self) -> None:
        self.data = self.full_data["data"]
        if self.data["node"]["template"]["_type"] == "Component":
            if "outputs" not in self.data["node"]:
                msg = f"Outputs not found for {self.display_name}"
                raise ValueError(msg)
            self.outputs = self.data["node"]["outputs"]
        else:
            self.outputs = self.data["node"].get("outputs", [])
            self.output = self.data["node"]["base_classes"]

        self.display_name: str = self.data["node"].get("display_name", self.id.split("-")[0])
        self.icon: str = self.data["node"].get("icon", self.id.split("-")[0])

        self.description: str = self.data["node"].get("description", "")
        self.frozen: bool = self.data["node"].get("frozen", False)

        self.is_input = self.data["node"].get("is_input") or self.is_input
        self.is_output = self.data["node"].get("is_output") or self.is_output
        template_dicts = {key: value for key, value in self.data["node"]["template"].items() if isinstance(value, dict)}

        self.has_session_id = "session_id" in template_dicts

        self.required_inputs: list[str] = []
        self.optional_inputs: list[str] = []
        for value_dict in template_dicts.values():
            list_to_append = self.required_inputs if value_dict.get("required") else self.optional_inputs

            if "type" in value_dict:
                list_to_append.append(value_dict["type"])
            if "input_types" in value_dict:
                list_to_append.extend(value_dict["input_types"])

        template_dict = self.data["node"]["template"]
        self.vertex_type = (
            self.data["type"]
            if "Tool" not in [type_ for out in self.outputs for type_ in out["types"]]
            or template_dict["_type"].islower()
            else template_dict["_type"]
        )

        if self.base_type is None:
            for base_type, value in lazy_load_dict.all_types_dict.items():
                if self.vertex_type in value:
                    self.base_type = base_type
                    break

    def get_value_from_output_names(self, key: str):
        if key in self.output_names:
            return self.graph.get_vertex(key)
        return None

    def get_value_from_template_dict(self, key: str):
        template_dict = self.data.get("node", {}).get("template", {})

        if key not in template_dict:
            msg = f"Key {key} not found in template dict"
            raise ValueError(msg)
        return template_dict.get(key, {}).get("value")

    def _set_params_from_normal_edge(self, params: dict, edge: Edge, template_dict: dict):
        param_key = edge.target_param

        # If the param_key is in the template_dict and the edge.target_id is the current node
        # We check this to make sure params with the same name but different target_id
        # don't get overwritten
        if param_key in template_dict and edge.target_id == self.id:
            if template_dict[param_key].get("list"):
                if param_key not in params:
                    params[param_key] = []
                params[param_key].append(self.graph.get_vertex(edge.source_id))
            elif edge.target_id == self.id:
                if isinstance(template_dict[param_key].get("value"), dict):
                    # we don't know the key of the dict but we need to set the value
                    # to the vertex that is the source of the edge
                    param_dict = template_dict[param_key]["value"]
                    if not param_dict or len(param_dict) != 1:
                        params[param_key] = self.graph.get_vertex(edge.source_id)
                    else:
                        params[param_key] = {key: self.graph.get_vertex(edge.source_id) for key in param_dict}

                else:
                    params[param_key] = self.graph.get_vertex(edge.source_id)
        elif param_key in self.output_names:
            #  if the loop is run the param_key item will be set over here
            # validate the edge
            params[param_key] = self.graph.get_vertex(edge.source_id)
        return params

    def build_params(self) -> None:
        """Build parameters for the vertex using the ParameterHandler."""
        if self.graph is None:
            msg = "Graph not found"
            raise ValueError(msg)

        if self.updated_raw_params:
            # Don't reset the flag - keep it True to protect against multiple build_params() calls
            # The flag will be reset when _build_each_vertex_in_params_dict() processes the params
            return

        # Create parameter handler with lazy storage service initialization
        param_handler = ParameterHandler(self, storage_service=None)

        # Process edge parameters
        edge_params = param_handler.process_edge_parameters(self.edges)

        # Process field parameters
        field_params, load_from_db_fields = param_handler.process_field_parameters()

        # Combine parameters, edge_params take precedence
        self.params = {**field_params, **edge_params}
        self.load_from_db_fields = load_from_db_fields
        self.raw_params = self.params.copy()

    def update_raw_params(self, new_params: Mapping[str, str | list[str]], *, overwrite: bool = False) -> None:
        """Update the raw parameters of the vertex with the given new parameters.

        Args:
            new_params (Dict[str, Any]): The new parameters to update.
            overwrite (bool, optional): Whether to overwrite the existing parameters.
                Defaults to False.

        Raises:
            ValueError: If any key in new_params is not found in self.raw_params.
        """
        # First check if the input_value in raw_params is not a vertex
        if not new_params:
            return
        if any(isinstance(self.raw_params.get(key), Vertex) for key in new_params):
            return
        if not overwrite:
            for key in new_params.copy():  # type: ignore[attr-defined]
                if key not in self.raw_params:
                    new_params.pop(key)  # type: ignore[attr-defined]
        self.raw_params.update(new_params)
        self.params = self.raw_params.copy()
        self.updated_raw_params = True

    def instantiate_component(self, user_id=None) -> None:
        if not self.custom_component:
            self.custom_component, _ = initialize.loading.instantiate_class(
                user_id=user_id,
                vertex=self,
            )

    async def _build(
        self,
        fallback_to_env_vars,
        user_id=None,
        event_manager: EventManager | None = None,
    ) -> None:
        """Initiate the build process."""
        await logger.adebug(f"Building {self.display_name}")
        await self._build_each_vertex_in_params_dict()
        if self.base_type is None:
            msg = f"Base type for vertex {self.display_name} not found"
            raise ValueError(msg)

        if not self.custom_component:
            custom_component, custom_params = initialize.loading.instantiate_class(
                user_id=user_id, vertex=self, event_manager=event_manager
            )
        else:
            custom_component = self.custom_component
            if hasattr(self.custom_component, "set_event_manager"):
                self.custom_component.set_event_manager(event_manager)
            custom_params = initialize.loading.get_params(self.params)

        await self._build_results(
            custom_component=custom_component,
            custom_params=custom_params,
            fallback_to_env_vars=fallback_to_env_vars,
            base_type=self.base_type,
        )

        self._validate_built_object()

        self.built = True

    def extract_messages_from_artifacts(self, artifacts: dict[str, Any]) -> list[dict]:
        """Extracts messages from the artifacts.

        Args:
            artifacts (Dict[str, Any]): The artifacts to extract messages from.

        Returns:
            List[str]: The extracted messages.
        """
        try:
            text = artifacts["text"]
            sender = artifacts.get("sender")
            sender_name = artifacts.get("sender_name")
            session_id = artifacts.get("session_id")
            stream_url = artifacts.get("stream_url")
            files = [{"path": file} if isinstance(file, str) else file for file in artifacts.get("files", [])]
            component_id = self.id
            type_ = self.artifacts_type

            if isinstance(sender_name, Data | Message):
                sender_name = sender_name.get_text()

            messages = [
                ChatOutputResponse(
                    message=text,
                    sender=sender,
                    sender_name=sender_name,
                    session_id=session_id,
                    stream_url=stream_url,
                    files=files,
                    component_id=component_id,
                    type=type_,
                ).model_dump(exclude_none=True)
            ]
        except KeyError:
            messages = []

        return messages

    def _get_all_upstream_vertices(self) -> list[Vertex]:
        """Walk all upstream vertices using edges, deduplicating by ID."""
        visited: set[str] = set()
        result: list[Vertex] = []
        stack = [edge.source_id for edge in self.graph.edges if edge.target_id == self.id]

        while stack:
            vid = stack.pop()
            if vid in visited:
                continue
            visited.add(vid)
            vertex = self.graph.get_vertex(vid)
            result.append(vertex)
            stack.extend(edge.source_id for edge in self.graph.edges if edge.target_id == vid)

        return result

    def _accumulate_upstream_token_usage(self) -> Usage | None:
        """Accumulate token usage from all upstream vertices.

        Walks all recursive predecessors via edges, deduplicates by vertex ID,
        and sums their token usage into a single total.
        """
        predecessors = self._get_all_upstream_vertices()
        total_input = 0
        total_output = 0
        has_data = False

        for predecessor in predecessors:
            if predecessor.result and predecessor.result.token_usage:
                usage = predecessor.result.token_usage
                total_input += usage.input_tokens or 0
                total_output += usage.output_tokens or 0
                has_data = True

        # Include own token usage if present
        if self.custom_component:
            own_usage = self.custom_component._token_usage  # noqa: SLF001
            if own_usage:
                total_input += own_usage.input_tokens or 0
                total_output += own_usage.output_tokens or 0
                has_data = True

        if not has_data:
            return None

        return Usage(
            input_tokens=total_input,
            output_tokens=total_output,
            total_tokens=total_input + total_output,
        )

    def _extract_token_usage(self) -> Usage | None:
        """Extract token usage from the custom component if available.

        Output vertices don't show token usage on the node badge because
        the accumulated total is displayed on the chat message instead.
        """
        if self.is_output:
            return None
        if self.custom_component and self.custom_component._token_usage:  # noqa: SLF001
            return self.custom_component._token_usage  # noqa: SLF001
        return None

    def finalize_build(self) -> None:
        result_dict = self.get_built_result()
        # We need to set the artifacts to pass information
        # to the frontend
        self.set_artifacts()
        artifacts = self.artifacts_raw
        messages = self.extract_messages_from_artifacts(artifacts) if isinstance(artifacts, dict) else []
        token_usage = self._extract_token_usage()
        result_dict = ResultData(
            results=result_dict,
            artifacts=artifacts,
            outputs=self.outputs_logs,
            logs=self.logs,
            messages=messages,
            component_display_name=self.display_name,
            component_id=self.id,
            token_usage=token_usage,
        )
        self.set_result(result_dict)

    async def _build_each_vertex_in_params_dict(self) -> None:
        """Iterates over each vertex in the params dictionary and builds it."""
        for key, value in self.raw_params.items():
            if self._is_vertex(value):
                if value == self:
                    del self.params[key]
                    continue
                await self._build_vertex_and_update_params(
                    key,
                    value,
                )
            elif isinstance(value, list) and self._is_list_of_vertices(value):
                await self._build_list_of_vertices_and_update_params(key, value)
            elif isinstance(value, dict):
                await self._build_dict_and_update_params(
                    key,
                    value,
                )
            elif key not in self.params or self.updated_raw_params:
                self.params[key] = value

        # Reset the flag after processing raw_params
        if self.updated_raw_params:
            self.updated_raw_params = False

    async def _build_dict_and_update_params(
        self,
        key,
        vertices_dict: dict[str, Vertex],
    ) -> None:
        """Iterates over a dictionary of vertices, builds each and updates the params dictionary."""
        for sub_key, value in vertices_dict.items():
            if not self._is_vertex(value):
                self.params[key][sub_key] = value
            else:
                result = await value.get_result(self, target_handle_name=key)
                self.params[key][sub_key] = result

    @staticmethod
    def _is_vertex(value):
        """Checks if the provided value is an instance of Vertex."""
        return isinstance(value, Vertex)

    def _is_list_of_vertices(self, value):
        """Checks if the provided value is a list of Vertex instances."""
        return all(self._is_vertex(vertex) for vertex in value)

    async def get_result(self, requester: Vertex, target_handle_name: str | None = None) -> Any:
        """Retrieves the result of the vertex.

        This is a read-only method so it raises an error if the vertex has not been built yet.

        Returns:
            The result of the vertex.
        """
        async with self.lock:
            return await self._get_result(requester, target_handle_name)

    async def _log_transaction_async(
        self,
        flow_id: str | UUID,
        source: Vertex,
        status: str,
        target: Vertex | None = None,
        error: str | Exception | None = None,
        outputs: dict[str, Any] | None = None,
    ) -> None:
        """Log a transaction asynchronously.

        Args:
            flow_id: The ID of the flow
            source: Source vertex
            status: Transaction status
            target: Optional target vertex
            error: Optional error information
            outputs: Optional explicit outputs dict (component execution results)
        """
        try:
            await log_transaction(flow_id, source, status, target, error, outputs)
        except Exception as exc:  # noqa: BLE001
            logger.debug(f"Error logging transaction: {exc!s}")

    async def _get_result(
        self,
        requester: Vertex,  # noqa: ARG002
        target_handle_name: str | None = None,  # noqa: ARG002
    ) -> Any:
        """Retrieves the result of the built component.

        If the component has not been built yet, a ValueError is raised.

        Returns:
            The built result if use_result is True, else the built object.
        """
        if not self.built:
            msg = f"Component {self.display_name} has not been built yet"
            raise ValueError(msg)

        return self.built_result if self.use_result else self.built_object

    async def _build_vertex_and_update_params(self, key, vertex: Vertex) -> None:
        """Builds a given vertex and updates the params dictionary accordingly."""
        result = await vertex.get_result(self, target_handle_name=key)
        self._handle_func(key, result)
        if isinstance(result, list):
            self._extend_params_list_with_result(key, result)
        self.params[key] = result

    async def _build_list_of_vertices_and_update_params(
        self,
        key,
        vertices: list[Vertex],
    ) -> None:
        """Iterates over a list of vertices, builds each and updates the params dictionary."""
        self.params[key] = []
        for vertex in vertices:
            result = await vertex.get_result(self, target_handle_name=key)
            # Weird check to see if the params[key] is a list
            # because sometimes it is a Data and breaks the code
            if not isinstance(self.params[key], list):
                self.params[key] = [self.params[key]]

            if isinstance(result, list):
                self.params[key].extend(result)
            else:
                try:
                    if self.params[key] == result:
                        continue

                    self.params[key].append(result)
                except AttributeError as e:
                    await logger.aexception(e)
                    msg = (
                        f"Params {key} ({self.params[key]}) is not a list and cannot be extended with {result}"
                        f"Error building Component {self.display_name}: \n\n{e}"
                    )
                    raise ValueError(msg) from e

    def _handle_func(self, key, result) -> None:
        """Handles 'func' key by checking if the result is a function and setting it as coroutine."""
        if key == "func":
            if not isinstance(result, types.FunctionType):
                if hasattr(result, "run"):
                    result = result.run
                elif hasattr(result, "get_function"):
                    result = result.get_function()
            elif inspect.iscoroutinefunction(result):
                self.params["coroutine"] = result
            else:
                self.params["coroutine"] = sync_to_async(result)

    def _extend_params_list_with_result(self, key, result) -> None:
        """Extends a list in the params dictionary with the given result if it exists."""
        if isinstance(self.params[key], list):
            self.params[key].extend(result)

    async def _build_results(
        self,
        custom_component,
        custom_params,
        base_type: str,
        *,
        fallback_to_env_vars=False,
    ) -> None:
        try:
            result = await initialize.loading.get_instance_results(
                custom_component=custom_component,
                custom_params=custom_params,
                vertex=self,
                fallback_to_env_vars=fallback_to_env_vars,
                base_type=base_type,
            )

            self.outputs_logs = build_output_logs(self, result)

            self._update_built_object_and_artifacts(result)
        except Exception as exc:
            tb = traceback.format_exc()
            await logger.aexception(exc)
            # Log transaction error
            flow_id = self.graph.flow_id
            if flow_id:
                await self._log_transaction_async(
                    str(flow_id), source=self, target=None, status="error", error=str(exc)
                )
            msg = f"Error building Component {self.display_name}: \n\n{exc}"
            raise ComponentBuildError(msg, tb) from exc

    def _update_built_object_and_artifacts(self, result: Any | tuple[Any, dict] | tuple[Component, Any, dict]) -> None:
        """Updates the built object and its artifacts."""
        if isinstance(result, tuple):
            if len(result) == 2:  # noqa: PLR2004
                self.built_object, self.artifacts = result
            elif len(result) == 3:  # noqa: PLR2004
                self.custom_component, self.built_object, self.artifacts = result
                self.logs = self.custom_component.get_output_logs()
                self.artifacts_raw = self.artifacts.get("raw", None)
                self.artifacts_type = {
                    self.outputs[0]["name"]: self.artifacts.get("type", None) or ArtifactType.UNKNOWN.value
                }
                self.artifacts = {self.outputs[0]["name"]: self.artifacts}
        else:
            self.built_object = result

    def _validate_built_object(self) -> None:
        """Checks if the built object is None and raises a ValueError if so."""
        if isinstance(self.built_object, UnbuiltObject):
            msg = f"{self.display_name}: {self.built_object_repr()}"
            raise TypeError(msg)
        if self.built_object is None:
            message = f"{self.display_name} returned None."
            if self.base_type == "custom_components":
                message += " Make sure your build method returns a component."

            logger.warning(message)
        elif isinstance(self.built_object, Iterator | AsyncIterator):
            if self.display_name == "Text Output":
                msg = f"You are trying to stream to a {self.display_name}. Try using a Chat Output instead."
                raise ValueError(msg)

    def _reset(self) -> None:
        self.built = False
        self.built_object = UnbuiltObject()
        self.built_result = UnbuiltResult()
        self.artifacts = {}
        self.steps_ran = []
        self.build_params()

    def _is_chat_input(self) -> bool:
        return False

    def build_inactive(self) -> None:
        # Just set the results to None
        self.built = True
        self.built_object = None
        self.built_result = None

    async def build(
        self,
        user_id=None,
        inputs: dict[str, Any] | None = None,
        files: list[str] | None = None,
        requester: Vertex | None = None,
        event_manager: EventManager | None = None,
        **kwargs,
    ) -> Any:
        # Add lazy loading check at the beginning
        # Check if we need to fully load this component first
        from lfx.interface.components import ensure_component_loaded
        from lfx.services.deps import get_settings_service

        settings_service = get_settings_service()
        if settings_service and settings_service.settings.lazy_load_components:
            component_name = self.id.split("-")[0]
            await ensure_component_loaded(self.vertex_type, component_name, settings_service)

        # Continue with the original implementation
        async with self.lock:
            if self.state == VertexStates.INACTIVE:
                # If the vertex is inactive, return None
                self.build_inactive()
                return None

            # Loop components should always run, even when frozen,
            # because they need to iterate through their data
            is_loop_component = self.display_name == "Loop" or self.is_loop
            if self.frozen and self.built and not is_loop_component:
                return await self.get_requester_result(requester)
            if self.built and requester is not None:
                # This means that the vertex has already been built
                # and we are just getting the result for the requester
                return await self.get_requester_result(requester)
            self._reset()

            # Emit build_start event for webhook real-time feedback
            if self.graph and self.graph.flow_id:
                await emit_build_start_event(self.graph.flow_id, self.id)

            # inject session_id if it is not None
            if inputs is not None and "session" in inputs and inputs["session"] is not None and self.has_session_id:
                session_id_value = self.get_value_from_template_dict("session_id")
                if session_id_value == "":
                    self.update_raw_params({"session_id": inputs["session"]}, overwrite=True)
            if self._is_chat_input() and (inputs or files):
                chat_input = {}
                if (
                    inputs
                    and isinstance(inputs, dict)
                    and "input_value" in inputs
                    and inputs.get("input_value") is not None
                ):
                    chat_input.update({"input_value": inputs.get(INPUT_FIELD_NAME, "")})
                if files:
                    chat_input.update({"files": files})

                self.update_raw_params(chat_input, overwrite=True)

            # Run steps
            for step in self.steps:
                if step not in self.steps_ran:
                    await step(user_id=user_id, event_manager=event_manager, **kwargs)
                    self.steps_ran.append(step)

            self.finalize_build()

            # Log transaction after successful build
            flow_id = self.graph.flow_id
            if flow_id:
                # Extract outputs from outputs_logs for transaction logging
                outputs_dict = None
                if self.outputs_logs:
                    outputs_dict = {
                        k: v.model_dump() if hasattr(v, "model_dump") else v for k, v in self.outputs_logs.items()
                    }
                await self._log_transaction_async(
                    str(flow_id), source=self, target=None, status="success", outputs=outputs_dict
                )

        return await self.get_requester_result(requester)

    async def get_requester_result(self, requester: Vertex | None):
        # If the requester is None, this means that
        # the Vertex is the root of the graph
        if requester is None:
            return self.built_object

        # Get the requester edge
        requester_edge = next((edge for edge in self.edges if edge.target_id == requester.id), None)
        # Return the result of the requester edge
        return (
            None
            if requester_edge is None
            else await requester_edge.get_result_from_source(source=self, target=requester)
        )

    def add_edge(self, edge: CycleEdge) -> None:
        if edge not in self.edges:
            self.edges.append(edge)

    def __repr__(self) -> str:
        return f"Vertex(display_name={self.display_name}, id={self.id}, data={self.data})"

    def __eq__(self, /, other: object) -> bool:
        try:
            if not isinstance(other, Vertex):
                return False
            # We should create a more robust comparison
            # for the Vertex class
            ids_are_equal = self.id == other.id
            # self.data is a dict and we need to compare them
            # to check if they are equal
            data_are_equal = self.data == other.data
        except AttributeError:
            return False
        else:
            return ids_are_equal and data_are_equal

    def __hash__(self) -> int:
        return id(self)

    def built_object_repr(self) -> str:
        # Add a message with an emoji, stars for success,
        return "Built successfully ✨" if self.built_object is not None else "Failed to build 😵‍💫"

    def apply_on_outputs(self, func: Callable[[Any], Any]) -> None:
        """Applies a function to the outputs of the vertex."""
        if not self.custom_component or not self.custom_component.outputs:
            return
        # Apply the function to each output
        [func(output) for output in self.custom_component.get_outputs_map().values()]

    # AGUI/AG UI Event Streaming Callbacks/Methods - (Optional, see Observable decorator)
    def raw_event_metrics(self, optional_fields: dict | None) -> dict:
        """This method is used to get the metrics of the vertex by the Observable decorator.

        If the vertex has a get_metrics method, it will be called, and the metrics will be captured
        to stream back to the user in an AGUI compliant format.
        Additional fields/metrics to be captured can be modified in this method, or in the callback methods,
        which are before_callback_event and after_callback_event before returning the AGUI event.
        """
        if optional_fields is None:
            optional_fields = {}
        import time

        return {"timestamp": time.time(), **optional_fields}

    def before_callback_event(self, *args, **kwargs) -> StepStartedEvent:  # noqa: ARG002
        """Should be a AGUI compatible event.

        VERTEX class generates a StepStartedEvent event.
        """
        metrics = {}
        if hasattr(self, "raw_event_metrics"):
            metrics = self.raw_event_metrics({"component_id": self.id})

        return StepStartedEvent(step_name=self.display_name, raw_event={"langflow": metrics})

    def after_callback_event(self, result, *args, **kwargs) -> StepFinishedEvent:  # noqa: ARG002
        """Should be a AGUI compatible event.

        VERTEX class generates a StepFinishedEvent event.
        """
        metrics = {}
        if hasattr(self, "raw_event_metrics"):
            metrics = self.raw_event_metrics({"component_id": self.id})
        return StepFinishedEvent(step_name=self.display_name, raw_event={"langflow": metrics})
