"""Tightly coupled reference implementations.

Tightly coupled means the process doing the work *is* a Dagster process: it
imports the code location, connects to the metadata database and writes run
events directly. Every execution pod therefore needs network reachability to the
orchestration runtime dependencies.

Three jobs are provided, differing only in their ``executor_def``. That single
construct is the code-level half of the execution-target binding; the other half
is the instance-level run launcher (see ``yaml/tightly-coupled/``).
"""

from __future__ import annotations

from dagster import graph, in_process_executor, multiprocess_executor
from dagster_k8s import k8s_job_executor

from distributed_execution.ops import (
    generate_work_units,
    process_work_unit,
    report_execution_target,
    summarise_results,
)

# Applied to run pods by the K8sRunLauncher; surfaces in the Dagster UI run tags.
COMMON_TAGS = {
    "execution_target": "tightly_coupled",
    "business_operation": "DISTRIBUTED_EXECUTION_REFERENCE",
}

# Per-step pod shape. Only honoured by k8s_job_executor.
STEP_K8S_CONFIG = {
    "container_config": {
        "resources": {
            "requests": {"cpu": "100m", "memory": "128Mi"},
            "limits": {"cpu": "500m", "memory": "512Mi"},
        },
    },
    "pod_spec_config": {
        "restart_policy": "Never",
    },
}


@graph
def distributed_execution_reference():
    """Shared topology, so the three jobs differ only by execution target."""
    target_report = report_execution_target()
    units = generate_work_units()
    results = units.map(process_work_unit).collect()
    return summarise_results(results, target_report)


# 1. Single process. Steps run inside the run worker itself - no fan-out at all.
tightly_coupled_in_process_job = distributed_execution_reference.to_job(
    name="tightly_coupled_in_process_job",
    description=(
        "Tightly coupled, single-process. Steps execute inside the run worker. "
        "Runs unchanged on a laptop and in Kubernetes."
    ),
    executor_def=in_process_executor,
    tags={**COMMON_TAGS, "executor": "in_process"},
)

# 2. Subprocesses on the run worker. Fan-out bounded by that one pod's resources.
tightly_coupled_local_job = distributed_execution_reference.to_job(
    name="tightly_coupled_local_job",
    description=(
        "Tightly coupled, multiprocess. Steps execute as subprocesses of the run worker; "
        "concurrency is bounded by the run pod's CPU and memory limits."
    ),
    executor_def=multiprocess_executor.configured({"max_concurrent": 2}),
    tags={**COMMON_TAGS, "executor": "multiprocess"},
)

# 3. One Kubernetes Job per step. Requires a cluster; each step pod connects to
#    the metadata database on its own, which is what makes this tightly coupled.
tightly_coupled_k8s_job = distributed_execution_reference.to_job(
    name="tightly_coupled_k8s_job",
    description=(
        "Tightly coupled, one Kubernetes Job per step. Each step pod must reach the metadata "
        "database, object storage and Vault. Requires a cluster - not runnable locally."
    ),
    executor_def=k8s_job_executor.configured(
        {
            "image_pull_policy": "IfNotPresent",
            "step_k8s_config": STEP_K8S_CONFIG,
        }
    ),
    tags={**COMMON_TAGS, "executor": "k8s_job"},
)

TIGHTLY_COUPLED_JOBS = [
    tightly_coupled_in_process_job,
    tightly_coupled_local_job,
    tightly_coupled_k8s_job,
]