"""Author a workflow with the SDK and execute it via the Orcheo server.

Prerequisites:
    1. Install dependencies: `uv sync --all-groups`
    2. Run the backend locally: `make dev-server`

Once the server is running, execute this script with `uv run python
examples/quickstart/sdk_server_trigger.py` to see the workflow updates
streaming back over the websocket connection.
"""

from __future__ import annotations
import asyncio
import json
import websockets
from pydantic import BaseModel
from websockets import exceptions as ws_exceptions
from orcheo_sdk import (
    OrcheoClient,
    Workflow,
    WorkflowNode,
)


class FieldTransformConfig(BaseModel):
    """Single field mapping for DataTransformNode."""

    source: str
    target: str
    transform: str = "identity"


class DataTransformConfig(BaseModel):
    """Configuration schema for the DataTransformNode node."""

    input_data: dict[str, str]
    transforms: list[FieldTransformConfig]


class DataTransformWorkflowNode(WorkflowNode[DataTransformConfig]):
    """Convenience wrapper that exports DataTransformNode nodes from the SDK."""

    type_name = "DataTransformNode"


def build_workflow() -> Workflow:
    """Create a multi-step workflow that maps and formats a user name."""
    workflow = Workflow(name="sdk-websocket-demo")

    workflow.add_node(
        DataTransformWorkflowNode(
            "greet_user",
            DataTransformConfig(
                input_data={"name": "{{inputs.name}}"},
                transforms=[
                    FieldTransformConfig(
                        source="name",
                        target="message",
                        transform="string",
                    )
                ],
            ),
        )
    )

    workflow.add_node(
        DataTransformWorkflowNode(
            "format_message",
            DataTransformConfig(
                input_data={"greeting": "{{results.greet_user.result.message}}"},
                transforms=[
                    FieldTransformConfig(
                        source="greeting",
                        target="shout",
                        transform="upper",
                    )
                ],
            ),
        ),
        depends_on=["greet_user"],
    )

    return workflow


async def run() -> None:
    """Connect to the server, trigger the workflow, and stream updates."""
    workflow = build_workflow()
    graph_config = workflow.to_graph_config()

    client = OrcheoClient(base_url="http://localhost:8000")
    workflow_identifier = "sdk-websocket-demo"
    websocket_url = client.websocket_url(workflow_identifier)
    payload = client.build_payload(
        graph_config,
        inputs={"name": "Ada Lovelace"},
    )

    try:
        async with websockets.connect(
            websocket_url,
            open_timeout=5,
            close_timeout=5,
        ) as websocket:
            await websocket.send(json.dumps(payload))

            async for message in websocket:
                update = json.loads(message)
                print(f"Update: {update}")
                status = update.get("status")
                if status == "error":
                    error_detail = update.get("error") or "Unknown error"
                    print(f"Workflow execution failed: {error_detail}")
                    break
                if status == "completed":
                    break
    except (ConnectionRefusedError, OSError) as exc:
        print(
            "Failed to connect to the Orcheo server. "
            "Ensure `make dev-server` is running before executing this script."
        )
        print(f"Connection error: {exc}")
    except TimeoutError:
        print(
            "Timed out while establishing or closing the WebSocket connection. "
            "Retry once the server is reachable."
        )
    except ws_exceptions.InvalidStatusCode as exc:
        print(
            "The server rejected the WebSocket handshake. Verify the workflow "
            "identifier and backend availability."
        )
        print(f"HTTP status: {exc.status_code}")
    except ws_exceptions.WebSocketException as exc:
        print(f"WebSocket communication error: {exc}")


if __name__ == "__main__":
    asyncio.run(run())
