from functools import partial
from typing import List, Optional

from pydantic import BaseModel, Field, ConfigDict
from pydantic.alias_generators import to_camel

from persona_models.common import create_uuid


class Connection(BaseModel):
    id: str = Field(default_factory=partial(create_uuid, "connection"))
    source_id: str
    target_id: str
    source_handle: str = "source"
    target_handle: str = "target"

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)


class WorkflowResult(BaseModel):
    workflow_id: str
    steps: list = Field(default_factory=list)

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)


class ExecutionStatus(BaseModel):
    id: str = Field(default_factory=create_uuid)
    workflow_id: str
    status: str = "pending"
    result: Optional[dict] = None
    error: Optional[str] = None

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)


class Node(BaseModel):
    id: str = Field(default_factory=create_uuid)
    type: str
    label: str | None = None
    data: dict = Field(default_factory=dict)

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)


class Workflow(BaseModel):
    id: str = Field(default_factory=create_uuid)
    project_id: str | None = None
    name: str | None = None
    description: str | None = None
    root_node: Optional[Node] = None
    nodes: list[Node] = Field(default_factory=list)
    connections: list[Connection] = Field(default_factory=list)
    features: Optional[list] = None

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)

    def get_node_by_id(self, node_id: str) -> Node | None:
        all_nodes = ([self.root_node] if self.root_node else []) + list(self.nodes)
        for node in all_nodes:
            if node.id == node_id:
                return node
        return None


class Execution(BaseModel):
    id: str = Field(default_factory=create_uuid)
    workflow_id: str
    status: str = "pending"
    input_data: Optional[dict] = None
    output_data: Optional[dict] = None
    error: Optional[str] = None

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)
