from __future__ import annotations

from typing import Any

from flownet.v1.compiler.errors import FlowDefinitionError
from flownet.v1.compiler.exception_scope import flow_program_scope
from flownet.v1.compiler.streams import ConnectedStreams, DataStream, SinkStream
from flownet.v1.core import FlowProgram


def compile_flow_program(flow_dsl: Any) -> FlowProgram:
    if isinstance(flow_dsl, FlowProgram):
        return flow_dsl
    flowfunction = _resolve_flow_callable(flow_dsl)
    return _compile_program(flowfunction)


def _resolve_flow_callable(flow_dsl: Any) -> Any:
    if callable(flow_dsl):
        return flow_dsl

    raise TypeError(
        "compile_flow_program() expects a callable flow DSL or v1 FlowProgram.",
    )


def _compile_program(flowfunction: Any) -> FlowProgram:
    program = FlowProgram()
    init_stream = DataStream(program, transformation=None, init=True)

    with flow_program_scope(program):
        output = flowfunction(init_stream)

    if output is None:
        program.no_return = True
    elif isinstance(output, SinkStream):
        program.no_return = True
    elif isinstance(output, ConnectedStreams):
        output.set_return()
    elif isinstance(output, DataStream):
        output.set_return()
    else:
        raise FlowDefinitionError(
            "Flow must return DataStream/ConnectedStreams/SinkStream/None, "
            f"got {type(output)}"
        )

    return program


__all__ = [
    "compile_flow_program",
]
