"""
Base node classes and interfaces for YBAgent workflows.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from enum import Enum
from pydantic import BaseModel, Field


class NodeType(str, Enum):
    """Node type enumeration."""
    INPUT = "input"
    OUTPUT = "output"
    LLM = "llm"
    LLM_OPENAI = "llm_openai"
    LLM_OLLAMA = "llm_ollama"
    LLM_HUGGINGFACE = "llm_huggingface"
    LLM_PROMPT = "llm_prompt"
    CONDITIONAL = "conditional"
    SWITCH = "switch"
    LOOP = "loop"
    MERGE = "merge"
    API_CALL = "api_call"
    TRANSFORM = "transform"
    SCRIPT = "script"
    VARIABLE = "variable"
    CUSTOM = "custom"
    
    # Data & Tools
    WEB_SCRAPER = "web_scraper"
    FILE_READ = "file_read"
    FILE_WRITE = "file_write"
    EMAIL = "email"
    DATABASE = "database"
    
    # Advanced AI
    VECTOR_SEARCH = "vector_search"
    SUMMARIZATION = "summarization"
    TRANSLATION = "translation"
    SENTIMENT = "sentiment"
    CLASSIFICATION = "classification"
    NER = "ner"
    VISION = "vision"
    
    # Code Agent
    CODE_AGENT = "code_agent"
    DATA_ANALYSIS = "data_analysis"
    REACT_AGENT = "react_agent"
    
    # Control & UX
    RETRY = "retry"
    ERROR_HANDLING = "error_handling"
    LOGGER = "logger"
    DEBUG = "debug"
    NOTIFICATION = "notification"
    FAIL = "fail"


class NodeInput(BaseModel):
    """Node input data structure."""
    name: str
    type: str = "any"
    required: bool = True
    default: Optional[Any] = None
    description: Optional[str] = None


class NodeOutput(BaseModel):
    """Node output data structure."""
    name: str
    type: str = "any"
    description: Optional[str] = None


class NodeMetadata(BaseModel):
    """Node metadata and configuration."""
    id: str
    node_type: NodeType
    name: Optional[str] = None
    description: Optional[str] = None
    inputs: List[NodeInput] = Field(default_factory=list)
    outputs: List[NodeOutput] = Field(default_factory=list)
    config: Dict[str, Any] = Field(default_factory=dict)
    retry_config: Dict[str, Any] = Field(default_factory=dict) # New field for retry settings
    position: Optional[Dict[str, float]] = None  # For visual builder: {x, y}


class BaseNode(ABC):
    """
    Abstract base class for all workflow nodes.
    
    All custom nodes should inherit from this class and implement the execute method.
    """
    
    def __init__(
        self,
        id: str,
        node_type: NodeType = NodeType.CUSTOM,
        name: Optional[str] = None,
        description: Optional[str] = None,
        **config
    ):
        """
        Initialize a node.
        
        Args:
            id: Unique identifier for the node
            node_type: Type of the node
            name: Human-readable name
            description: Node description
            **config: Additional configuration parameters
        """
        # Extract retry configuration from general config
        retry_config = {
            "max_retries": config.pop("max_retries", 0),
            "retry_delay": config.pop("retry_delay", 1.0),  # seconds
            "retry_on_error": config.pop("retry_on_error", True)
        }
        
        # Create clean config without node_type (and now retry params) to avoid conflicts
        clean_config = {k: v for k, v in config.items() if k != 'node_type'}
        
        self.metadata = NodeMetadata(
            id=id,
            node_type=node_type,
            name=name or id,
            description=description,
            config=clean_config,
            retry_config=retry_config
        )
        self._define_inputs_outputs()
    
    def _define_inputs_outputs(self):
        """Define node inputs and outputs. Override in subclasses."""
        pass
    
    @abstractmethod
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute the node logic.
        
        Args:
            inputs: Dictionary of input values from connected nodes
            context: Workflow execution context (shared state, variables, etc.)
        
        Returns:
            Dictionary of output values to pass to connected nodes
        """
        pass
    
    def validate_inputs(self, inputs: Dict[str, Any]) -> bool:
        """
        Validate that required inputs are present.
        
        Args:
            inputs: Input values to validate
        
        Returns:
            True if inputs are valid
        
        Raises:
            ValueError: If required inputs are missing
        """
        for input_def in self.metadata.inputs:
            if input_def.required and input_def.name not in inputs:
                if input_def.default is not None:
                    inputs[input_def.name] = input_def.default
                else:
                    raise ValueError(
                        f"Required input '{input_def.name}' missing for node '{self.metadata.id}'"
                    )
        return True
    
    def to_dict(self) -> Dict[str, Any]:
        """
        Serialize node to dictionary for saving workflows.
        
        Returns:
            Dictionary representation of the node
        """
        # Merge retry_config into the main config for serialization if it's not empty
        serialized_config = self.metadata.config.copy()
        if self.metadata.retry_config:
            serialized_config.update(self.metadata.retry_config)

        return {
            "id": self.metadata.id,
            "type": self.metadata.node_type.value,
            "name": self.metadata.name,
            "description": self.metadata.description,
            "config": serialized_config,
            "position": self.metadata.position,
            "inputs": [inp.model_dump() for inp in self.metadata.inputs],
            "outputs": [out.model_dump() for out in self.metadata.outputs],
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "BaseNode":
        """
        Deserialize node from dictionary.
        
        Args:
            data: Dictionary representation of the node
        
        Returns:
            Node instance
        """
        # Get config and remove node_type to avoid conflicts
        config = data.get("config", {}).copy()
        config.pop('node_type', None)
        
        return cls(
            id=data["id"],
            node_type=NodeType(data["type"]),
            name=data.get("name"),
            description=data.get("description"),
            **config
        )
    
    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}(id='{self.metadata.id}', type='{self.metadata.node_type.value}')>"
