# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

"""
Workflow
--------
```python
from datatailr import workflow, task
@task()
def get_number():
    return 42

@task()
def add(a, b):
    return a + b

@workflow("My Workflow")
def my_workflow():
    num = get_number()
    result = add(num, 8)

if __name__ == "__main__":
    my_workflow(local_run=True)
```


Running this script will execute the workflow locally. To deploy it to the Datatailr platform,
simply remove the `local_run=True` argument.

Note that the workflow is executed as a directed acyclic graph (DAG), so tasks can run in parallel
when there are no dependencies.

Arguments to tasks can be a mix of literals and other task outputs, like in the example above.

The resulting DAG will look like this:

```mermaid
graph LR
    A[get_number] --> B[add]
    8((8)) --> B
```

More complex workflows can be built using simple function composition:

```python
@task()
def reduce_list(*numbers):
    print(f"Reducing list: {numbers}")
    return sum(numbers)

@workflow("Complex Workflow")
def complex_workflow():
    total_1 = add(get_number().alias('my number'), 13).alias('first addition')
    total_2 = reduce_list(*[get_number().alias(f"num_{i}") for i in range(5)])
    add(total_2, 1).alias('add one')
    add(total_2, 2).alias('add two')
    add(total_2, 3).alias('add three')
    add(total_2, 4).alias('add four')
    add(total_2, 5).alias('add five')
```
This will generate the following graph and deploy it for execution on the Datatailr platform:

```mermaid
graph LR
    A[my_number] --> B[first addition];
    13((13)) --> B;
    D[num_0] --> E[reduce_list];
    D1[num_1] --> E;
    D2[num_2] --> E;
    D3[num_3] --> E;
    D4[num_4] --> E;
    E --> G1[add one];
    E --> G2[add two];
    E --> G3[add three];
    E --> G4[add four];
    E --> G5[add five];
```
Note the use of the `alias()` method to give tasks specific names in the workflow graph.
This is necessary when the same task is invoked multiple times to avoid name collisions.

Controlling task properties can be done on one of two levels:

- Globally for the entire workflow via arguments to the `@workflow` decorator.
- Individually for each task via arguments to the `@task` decorator.

These properties include:

- Image
- Resources
- Run as user
- ACL
- Python version
- Python requirements
- Build scripts


In addition, some properties can be updated on a task call basis.

For example, within the same workflow you can override resources or env vars for a single call:

```python
@workflow("Complex Workflow")
def complex_workflow():
    total_1 = add(get_number().alias('my number'), 13).alias('first addition')
    total_2 = reduce_list(*[get_number().alias(f"num_{i}") for i in range(5)])
    add(total_2, 1).set_resources(memory="512m", cpu=0.5).alias('add one')
    add(total_2, 2).update_env_vars({"MODE": "fast"}).alias('add two')
    add(total_2, 3).alias('add three')
```

To generate the JSON representation of the workflow without deploying or running it, simply pass the `to_json=True` argument when invoking the workflow function:
```python
workflow_json = my_workflow(to_json=True)
print(workflow_json)
```
"""

import os
import datetime
from typing import Callable, Optional, Dict, Union, List
import functools
from datatailr.logging import DatatailrLogger
from datatailr.scheduler import Workflow, Resources
from datatailr.scheduler.schedule import Schedule
from datatailr import Environment, Image, ACL, User

logger = DatatailrLogger(__name__).get_logger()


def workflow(
    name: Optional[str] = None,
    schedule: Optional[Schedule] = None,
    image: Optional[Image] = None,
    run_as: Optional[Union[str, User]] = None,
    resources: Resources = Resources(memory="100m", cpu=1),
    acl: Optional[ACL] = None,
    python_version: str = "auto",
    python_requirements: str | List[str] = "",
    build_script_pre: str = "",
    build_script_post: str = "",
    env_vars: Dict[str, str | int | float | bool] = {},
    fail_after: datetime.timedelta | str | None = None,
    expire_after: datetime.timedelta | str | None = None,
):
    """Decorator that turns a function into a Datatailr DAG workflow.

    The decorated function should compose ``@task``-decorated calls.  Task
    dependencies are inferred automatically from the data flow between
    function calls.

    All parameters supplied to this decorator set the **default** values for
    every task in the workflow.  Individual tasks can override these via
    their own ``@task`` arguments.

    Args:
        name: Display name of the workflow.  Defaults to the decorated
            function name, title-cased.
        schedule: An optional ``Schedule`` for recurring execution.
        image: Pre-configured container ``Image`` applied to all tasks.
        run_as: User or username under which the workflow runs.
        resources: Default ``Resources`` (CPU / memory) for each task.
        acl: Access control list for the workflow.
        python_version: Python version for the container images.
        python_requirements: Python dependencies (see ``Image``).
        build_script_pre: Dockerfile commands before pip install.
        build_script_post: Dockerfile commands after pip install.
        env_vars: Environment variables passed to every task container.
        fail_after: Maximum wall-clock duration before the workflow is
            marked as failed (e.g. ``timedelta(hours=2)`` or ``"2h"``).
        expire_after: Duration after which a completed workflow's
            resources are cleaned up.

    Returns:
        A wrapper function. Calling it deploys the workflow to the
        platform.  Pass ``local_run=True`` to execute locally, or
        ``to_json=True`` to obtain the JSON representation without
        deploying.  Pass ``save_only=True`` to save without deploying.

    Raises:
        ValueError: If ``local_run=True`` and a schedule is set.
    """
    _name = name
    _schedule = schedule
    _image = image
    _run_as = run_as
    _resources = resources
    _acl = acl
    _python_version = python_version
    _python_requirements = python_requirements
    _build_script_pre = build_script_pre
    _build_script_post = build_script_post
    _env_vars = env_vars
    _fail_after = fail_after
    _expire_after = expire_after

    def decorator(func) -> Callable:
        @functools.wraps(func)
        def wrapper(
            *args,
            local_run: bool = False,
            to_json: bool = False,
            save_only: bool = False,
            schedule: Optional[Schedule] = None,
            image: Optional[Image] = None,
            run_as: Optional[Union[str, User]] = None,
            resources: Optional[Resources] = None,
            acl: Optional[ACL] = None,
            python_version: Optional[str] = None,
            python_requirements: Optional[str] = None,
            build_script_pre: Optional[str] = None,
            build_script_post: Optional[str] = None,
            env_vars: Optional[Dict[str, str | int | float | bool]] = None,
            fail_after: Optional[datetime.timedelta | str] = None,
            expire_after: Optional[datetime.timedelta | str] = None,
            **kwargs,
        ):
            workflow_file_path = func.__code__.co_filename
            # If the workflow is being invoked from a package installed in the container then raise a warning and return
            if (
                workflow_file_path.startswith(
                    "/opt/datatailr/usr/lib/python/site-packages/"
                )
                and os.getenv("DATATAILR_JOB_TYPE") == "batch"
            ) or os.getenv("DATATAILR_BATCH_DONT_RUN_WORKFLOW") == "true":
                return
            __schedule = schedule or _schedule
            __image = image or _image
            __run_as = run_as or _run_as
            __resources = resources or _resources
            __acl = acl or _acl
            __python_version = python_version or _python_version
            __python_requirements = python_requirements or _python_requirements
            __build_script_pre = build_script_pre or _build_script_pre
            __build_script_post = build_script_post or _build_script_post
            __env_vars = env_vars or _env_vars
            __fail_after = fail_after or _fail_after
            __expire_after = expire_after or _expire_after

            if local_run and (__schedule is not None):
                raise ValueError("Cannot set schedule for local run.")

            dag = Workflow(
                name=_name or func.__name__.replace("_", " ").title(),
                environment=Environment.DEV,
                schedule=__schedule,
                image=__image,
                run_as=__run_as,
                resources=__resources,
                acl=__acl,
                local_run=local_run,
                python_version=__python_version,
                python_requirements=__python_requirements,
                build_script_pre=__build_script_pre,
                build_script_post=__build_script_post,
                env_vars=__env_vars,
                fail_after=__fail_after,
                expire_after=__expire_after,
            )
            dag.set_autorun(False)
            with dag:
                func(*args, **kwargs)
                if to_json:
                    dag.set_autorun(False)
                    return dag.to_json()
                elif save_only:
                    dag.set_autorun(False)
                    dag.save()
                    print(
                        f"Workflow '{dag.name}' saved successfully without deployment."
                    )

        return wrapper

    return decorator
