"""
Pydantic Schemas for NeuroFlow

Defines data models for Node, Edge, and Workflow objects
used throughout the application.
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Optional
from uuid import UUID, uuid4

from pydantic import BaseModel, Field


class NodeType(str, Enum):
    """Enumeration of all supported node types."""
    
    # Data Loading
    CSV_LOADER = "csv_loader"
    JSON_LOADER = "json_loader"
    
    # Preprocessing
    DATA_SCALER = "data_scaler"
    TRAIN_TEST_SPLIT = "train_test_split"
    FEATURE_SELECTOR = "feature_selector"
    MISSING_VALUE_HANDLER = "missing_value_handler"
    ENCODER = "encoder"
    
    # Training
    TRAINER = "trainer"
    
    # Evaluation
    EVALUATOR = "evaluator"
    
    # Output
    MODEL_EXPORTER = "model_exporter"
    PREDICTION = "prediction"


class ScalerType(str, Enum):
    """Types of scalers available."""
    STANDARD = "standard"
    MINMAX = "minmax"
    ROBUST = "robust"
    NORMALIZER = "normalizer"


class ModelType(str, Enum):
    """Types of ML models available."""
    # Classification
    LOGISTIC_REGRESSION = "logistic_regression"
    RANDOM_FOREST_CLASSIFIER = "random_forest_classifier"
    SVM_CLASSIFIER = "svm_classifier"
    KNN_CLASSIFIER = "knn_classifier"
    DECISION_TREE_CLASSIFIER = "decision_tree_classifier"
    GRADIENT_BOOSTING_CLASSIFIER = "gradient_boosting_classifier"
    
    # Regression
    LINEAR_REGRESSION = "linear_regression"
    RANDOM_FOREST_REGRESSOR = "random_forest_regressor"
    SVR = "svr"
    KNN_REGRESSOR = "knn_regressor"
    DECISION_TREE_REGRESSOR = "decision_tree_regressor"
    GRADIENT_BOOSTING_REGRESSOR = "gradient_boosting_regressor"


class ExecutionStatus(str, Enum):
    """Status of node/workflow execution."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


# ============== Position Schema ==============

class Position(BaseModel):
    """Position of a node on the canvas."""
    x: float = 0.0
    y: float = 0.0


# ============== Node Schemas ==============

class NodeParams(BaseModel):
    """Base parameters for a node - extensible dictionary."""
    
    class Config:
        extra = "allow"  # Allow additional fields


class CSVLoaderParams(NodeParams):
    """Parameters for CSV loader node."""
    file_path: str = ""
    delimiter: str = ","
    has_header: bool = True
    encoding: str = "utf-8"


class DataScalerParams(NodeParams):
    """Parameters for data scaler node."""
    scaler_type: ScalerType = ScalerType.STANDARD
    columns: list[str] = Field(default_factory=list)
    fit_on_train_only: bool = True


class TrainTestSplitParams(NodeParams):
    """Parameters for train/test split node."""
    test_size: float = 0.2
    random_state: Optional[int] = 42
    stratify_column: Optional[str] = None
    shuffle: bool = True


class TrainerParams(NodeParams):
    """Parameters for trainer node."""
    model_type: ModelType = ModelType.RANDOM_FOREST_CLASSIFIER
    target_column: str = ""
    hyperparameters: dict[str, Any] = Field(default_factory=dict)


class EvaluatorParams(NodeParams):
    """Parameters for evaluator node."""
    metrics: list[str] = Field(default_factory=lambda: ["accuracy", "f1"])
    cross_validation: bool = False
    cv_folds: int = 5


class NodeData(BaseModel):
    """Data associated with a node."""
    label: str = ""
    params: dict[str, Any] = Field(default_factory=dict)
    
    class Config:
        extra = "allow"


class Node(BaseModel):
    """
    Represents a single node in the workflow graph.
    
    Each node has a type that determines its behavior,
    and a params dictionary containing type-specific configuration.
    """
    id: str = Field(default_factory=lambda: str(uuid4()))
    type: str  # NodeType value as string to match React Flow
    position: Position = Field(default_factory=Position)
    data: NodeData = Field(default_factory=NodeData)
    
    # Execution state
    status: ExecutionStatus = ExecutionStatus.PENDING
    error_message: Optional[str] = None
    execution_time_ms: Optional[float] = None
    
    class Config:
        use_enum_values = True


class NodeCreate(BaseModel):
    """Schema for creating a new node."""
    type: str
    position: Position = Field(default_factory=Position)
    data: NodeData = Field(default_factory=NodeData)


class NodeUpdate(BaseModel):
    """Schema for updating a node."""
    position: Optional[Position] = None
    data: Optional[NodeData] = None


# ============== Edge Schemas ==============

class Edge(BaseModel):
    """
    Represents a connection between two nodes.
    
    Data flows from source node to target node.
    """
    id: str = Field(default_factory=lambda: str(uuid4()))
    source: str  # Source node ID
    target: str  # Target node ID
    source_handle: Optional[str] = None  # Output handle on source
    target_handle: Optional[str] = None  # Input handle on target
    
    class Config:
        use_enum_values = True


class EdgeCreate(BaseModel):
    """Schema for creating a new edge."""
    source: str
    target: str
    source_handle: Optional[str] = None
    target_handle: Optional[str] = None


# ============== Workflow Schemas ==============

class WorkflowMetadata(BaseModel):
    """Metadata for a workflow."""
    name: str = "Untitled Workflow"
    description: str = ""
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    tags: list[str] = Field(default_factory=list)
    version: str = "1.0.0"


class Workflow(BaseModel):
    """
    Represents a complete ML workflow as a directed acyclic graph (DAG).
    
    Contains nodes (operations) and edges (data flow connections).
    """
    id: str = Field(default_factory=lambda: str(uuid4()))
    metadata: WorkflowMetadata = Field(default_factory=WorkflowMetadata)
    nodes: list[Node] = Field(default_factory=list)
    edges: list[Edge] = Field(default_factory=list)
    
    # Execution state
    status: ExecutionStatus = ExecutionStatus.PENDING
    current_node_id: Optional[str] = None
    
    class Config:
        use_enum_values = True


class WorkflowCreate(BaseModel):
    """Schema for creating a new workflow."""
    metadata: Optional[WorkflowMetadata] = None
    nodes: list[NodeCreate] = Field(default_factory=list)
    edges: list[EdgeCreate] = Field(default_factory=list)


class WorkflowUpdate(BaseModel):
    """Schema for updating a workflow."""
    metadata: Optional[WorkflowMetadata] = None
    nodes: Optional[list[Node]] = None
    edges: Optional[list[Edge]] = None


# ============== Execution Schemas ==============

class ExecutionRequest(BaseModel):
    """Request to execute a workflow."""
    workflow: Workflow
    run_id: Optional[str] = None


class NodeExecutionResult(BaseModel):
    """Result of executing a single node."""
    node_id: str
    status: ExecutionStatus
    execution_time_ms: float
    output_summary: Optional[dict[str, Any]] = None
    error_message: Optional[str] = None


class ExecutionResult(BaseModel):
    """Result of executing a complete workflow."""
    run_id: str
    workflow_id: str
    status: ExecutionStatus
    started_at: datetime
    completed_at: Optional[datetime] = None
    total_execution_time_ms: float = 0.0
    node_results: list[NodeExecutionResult] = Field(default_factory=list)
    final_output: Optional[dict[str, Any]] = None
    error_message: Optional[str] = None


class ExecutionProgress(BaseModel):
    """Progress update during workflow execution."""
    run_id: str
    workflow_id: str
    current_node_id: str
    current_node_index: int
    total_nodes: int
    status: ExecutionStatus
    message: str = ""
    timestamp: datetime = Field(default_factory=datetime.utcnow)


# ============== API Response Schemas ==============

class APIResponse(BaseModel):
    """Standard API response wrapper."""
    success: bool
    message: str = ""
    data: Optional[Any] = None


class HealthResponse(BaseModel):
    """Health check response."""
    status: str = "healthy"
    version: str
    timestamp: datetime = Field(default_factory=datetime.utcnow)
