from typing import Optional

from core.app.entities.queue_entities import AppQueueEvent
from core.model_runtime.utils.encoders import jsonable_encoder
from core.workflow.callbacks.base_workflow_callback import BaseWorkflowCallback
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeType

_TEXT_COLOR_MAPPING = {
    "blue": "36;1",
    "yellow": "33;1",
    "pink": "38;5;200",
    "green": "32;1",
    "red": "31;1",
}


class WorkflowLoggingCallback(BaseWorkflowCallback):

    def __init__(self) -> None:
        self.current_node_id = None

    def on_workflow_run_started(self) -> None:
        """
        Workflow run started
        """
        self.print_text("\n[on_workflow_run_started]", color='pink')

    def on_workflow_run_succeeded(self) -> None:
        """
        Workflow run succeeded
        """
        self.print_text("\n[on_workflow_run_succeeded]", color='green')

    def on_workflow_run_failed(self, error: str) -> None:
        """
        Workflow run failed
        """
        self.print_text("\n[on_workflow_run_failed]", color='red')

    def on_workflow_node_execute_started(self, node_id: str,
                                         node_type: NodeType,
                                         node_data: BaseNodeData,
                                         node_run_index: int = 1,
                                         predecessor_node_id: Optional[str] = None) -> None:
        """
        Workflow node execute started
        """
        self.print_text("\n[on_workflow_node_execute_started]", color='yellow')
        self.print_text(f"Node ID: {node_id}", color='yellow')
        self.print_text(f"Type: {node_type.value}", color='yellow')
        self.print_text(f"Index: {node_run_index}", color='yellow')
        if predecessor_node_id:
            self.print_text(f"Predecessor Node ID: {predecessor_node_id}", color='yellow')

    def on_workflow_node_execute_succeeded(self, node_id: str,
                                           node_type: NodeType,
                                           node_data: BaseNodeData,
                                           inputs: Optional[dict] = None,
                                           process_data: Optional[dict] = None,
                                           outputs: Optional[dict] = None,
                                           execution_metadata: Optional[dict] = None) -> None:
        """
        Workflow node execute succeeded
        """
        self.print_text("\n[on_workflow_node_execute_succeeded]", color='green')
        self.print_text(f"Node ID: {node_id}", color='green')
        self.print_text(f"Type: {node_type.value}", color='green')
        self.print_text(f"Inputs: {jsonable_encoder(inputs) if inputs else ''}", color='green')
        self.print_text(f"Process Data: {jsonable_encoder(process_data) if process_data else ''}", color='green')
        self.print_text(f"Outputs: {jsonable_encoder(outputs) if outputs else ''}", color='green')
        self.print_text(f"Metadata: {jsonable_encoder(execution_metadata) if execution_metadata else ''}",
                        color='green')

    def on_workflow_node_execute_failed(self, node_id: str,
                                        node_type: NodeType,
                                        node_data: BaseNodeData,
                                        error: str,
                                        inputs: Optional[dict] = None,
                                        outputs: Optional[dict] = None,
                                        process_data: Optional[dict] = None) -> None:
        """
        Workflow node execute failed
        """
        self.print_text("\n[on_workflow_node_execute_failed]", color='red')
        self.print_text(f"Node ID: {node_id}", color='red')
        self.print_text(f"Type: {node_type.value}", color='red')
        self.print_text(f"Error: {error}", color='red')
        self.print_text(f"Inputs: {jsonable_encoder(inputs) if inputs else ''}", color='red')
        self.print_text(f"Process Data: {jsonable_encoder(process_data) if process_data else ''}", color='red')
        self.print_text(f"Outputs: {jsonable_encoder(outputs) if outputs else ''}", color='red')

    def on_node_text_chunk(self, node_id: str, text: str, metadata: Optional[dict] = None) -> None:
        """
        Publish text chunk
        """
        if not self.current_node_id or self.current_node_id != node_id:
            self.current_node_id = node_id
            self.print_text('\n[on_node_text_chunk]')
            self.print_text(f"Node ID: {node_id}")
            self.print_text(f"Metadata: {jsonable_encoder(metadata) if metadata else ''}")

        self.print_text(text, color="pink", end="")

    def on_event(self, event: AppQueueEvent) -> None:
        """
        Publish event
        """
        self.print_text("\n[on_workflow_event]", color='blue')
        self.print_text(f"Event: {jsonable_encoder(event)}", color='blue')

    def print_text(
            self, text: str, color: Optional[str] = None, end: str = "\n"
    ) -> None:
        """Print text with highlighting and no end characters."""
        text_to_print = self._get_colored_text(text, color) if color else text
        print(f'{text_to_print}', end=end)

    def _get_colored_text(self, text: str, color: str) -> str:
        """Get colored text."""
        color_str = _TEXT_COLOR_MAPPING[color]
        return f"\u001b[{color_str}m\033[1;3m{text}\u001b[0m"