"""Python source-code generator for the Metaflow-Prefect integration.

``generate_prefect_file`` is the single public function.  It takes a
``FlowSpec`` plus a ``PrefectFlowConfig`` and returns the complete text of a
self-contained Python file that, when executed, runs the Metaflow flow as a
Prefect flow.

Design constraints
------------------
- The generated file must be self-contained — no imports from this package.
- It uses only stdlib + metaflow + prefect (already required by this package).
- Linear, split/join, and foreach graph shapes are all handled.
- Foreach body tasks run sequentially by default; users can switch to a
  concurrent task-runner if they need parallel splits.
"""

from __future__ import annotations

from datetime import datetime
from typing import Sequence

from metaflow_extensions.prefect.plugins.prefect._types import (
    FlowSpec,
    NodeType,
    ParameterSpec,
    PrefectFlowConfig,
    StepSpec,
)

# ---------------------------------------------------------------------------
# Tiny code-builder
# ---------------------------------------------------------------------------

_INDENT = "    "


class _CB:
    """Incrementally builds Python source, tracking indentation level."""

    def __init__(self) -> None:
        self._lines: list[str] = []
        self._level: int = 0

    # -- public API ----------------------------------------------------------

    def emit(self, line: str = "") -> "_CB":
        pad = _INDENT * self._level
        self._lines.append(pad + line if line else "")
        return self

    def indent(self) -> "_CB":
        self._level += 1
        return self

    def dedent(self) -> "_CB":
        self._level = max(0, self._level - 1)
        return self

    def build(self) -> str:
        return "\n".join(self._lines)


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


def generate_prefect_file(spec: FlowSpec, cfg: PrefectFlowConfig) -> str:
    """Return the full Python source of a runnable Prefect flow file.

    Args:
        spec: Analysed flow description (from ``_graph.analyze_graph``).
        cfg:  User-supplied deployment configuration.

    Returns:
        A multi-line Python source string ready to be written to disk.
    """
    cb = _CB()
    _emit_header(cb, spec, cfg)
    _emit_helpers(cb, cfg)
    for step in spec.steps:
        cb.emit()
        cb.emit()
        _emit_task(cb, step, spec, cfg)
    cb.emit()
    cb.emit()
    _emit_flow(cb, spec, cfg)
    cb.emit()
    cb.emit("if __name__ == '__main__':")
    cb.indent()
    cb.emit("%s()" % _python_name(spec.name))
    cb.dedent()
    return cb.build()


# ---------------------------------------------------------------------------
# Header + helpers block
# ---------------------------------------------------------------------------


def _emit_header(cb: _CB, spec: FlowSpec, cfg: PrefectFlowConfig) -> None:
    cb.emit("# Generated by metaflow-prefect on %s" % datetime.now().isoformat(timespec="seconds"))
    cb.emit("# Metaflow flow: %s" % spec.name)
    cb.emit("# Do not edit this file — regenerate with: python %s prefect create <file>" % cfg.flow_file)
    cb.emit()
    cb.emit("from __future__ import annotations")
    cb.emit()
    cb.emit("import json")
    cb.emit("import os")
    cb.emit("import subprocess")
    cb.emit("import sys")
    cb.emit("import tempfile")
    cb.emit("import uuid")
    cb.emit("from typing import Any")
    cb.emit()
    cb.emit("from prefect import flow, task, get_run_logger")
    cb.emit("from prefect.artifacts import create_markdown_artifact")
    cb.emit("from prefect.context import get_run_context")
    cb.emit()
    cb.emit("# ---------------------------------------------------------------------------")
    cb.emit("# Compile-time configuration")
    cb.emit("# ---------------------------------------------------------------------------")
    cb.emit("FLOW_FILE: str = %r" % cfg.flow_file)
    cb.emit("FLOW_NAME: str = %r" % spec.name)
    cb.emit("DATASTORE_TYPE: str = %r" % cfg.datastore_type)
    cb.emit("METADATA_TYPE: str = %r" % cfg.metadata_type)
    cb.emit("USERNAME: str = %r" % cfg.username)
    cb.emit("TAGS: list[str] = %r" % list(spec.tags))
    cb.emit("NAMESPACE: str | None = %r" % spec.namespace)
    cb.emit("SCHEDULE_CRON: str | None = %r" % spec.schedule_cron)
    cb.emit("WITH_DECORATORS: list[str] = %r" % list(cfg.with_decorators))


def _emit_helpers(cb: _CB, cfg: PrefectFlowConfig) -> None:
    cb.emit()
    cb.emit()
    cb.emit("# ---------------------------------------------------------------------------")
    cb.emit("# Runtime helpers (embedded — no external imports needed)")
    cb.emit("# ---------------------------------------------------------------------------")
    cb.emit()
    cb.emit("def _foreach_info_path(run_id: str, step_name: str) -> str:")
    cb.indent()
    cb.emit('"""Return the temp file path where foreach cardinality is written."""')
    cb.emit("safe_run = run_id.replace('/', '_').replace(':', '_')")
    cb.emit("return os.path.join(tempfile.gettempdir(), f'mf_prefect_foreach_{safe_run}_{step_name}.json')")
    cb.dedent()
    cb.emit()
    cb.emit()
    cb.emit("def _read_foreach_info(path: str) -> int:")
    cb.indent()
    cb.emit('"""Read foreach split-count written by the prefect_internal decorator."""')
    cb.emit("try:")
    cb.indent()
    cb.emit("with open(path) as _f:")
    cb.indent()
    cb.emit("return int(json.load(_f)['num_splits'])")
    cb.dedent()
    cb.dedent()
    cb.emit("except Exception:")
    cb.indent()
    cb.emit("return 0")
    cb.dedent()
    cb.dedent()
    cb.emit()
    cb.emit()
    cb.emit("def _run_cmd(cmd: list[str], extra_env: dict[str, str] | None = None) -> None:")
    cb.indent()
    cb.emit('"""Execute *cmd* as a subprocess, inheriting stdout/stderr."""')
    cb.emit("env = os.environ.copy()")
    cb.emit("if extra_env:")
    cb.indent()
    cb.emit("env.update(extra_env)")
    cb.dedent()
    cb.emit("subprocess.run(cmd, env=env, check=True)")
    cb.dedent()
    cb.emit()
    cb.emit()
    cb.emit("def _mf_artifact_names(run_id: str, step_name: str, task_id: str) -> list[str]:")
    cb.indent()
    cb.emit('"""Return user-defined artifact names from the Metaflow datastore (no values loaded)."""')
    cb.emit("try:")
    cb.indent()
    cb.emit("from metaflow.datastore import FlowDataStore")
    cb.emit("from metaflow.plugins import DATASTORES")
    cb.emit("_impl = next(d for d in DATASTORES if d.TYPE == DATASTORE_TYPE)")
    cb.emit("_root = _impl.get_datastore_root_from_config(lambda *a: None)")
    cb.emit("_fds = FlowDataStore(FLOW_NAME, None, storage_impl=_impl, ds_root=_root)")
    cb.emit("_tds = _fds.get_task_datastore(run_id, step_name, task_id, attempt=0, mode='r')")
    cb.emit("_SKIP = {'name', 'input'}  # Metaflow internal artifact names")
    cb.emit("return [n for n in _tds if not n.startswith('_') and n not in _SKIP]")
    cb.dedent()
    cb.emit("except Exception:")
    cb.indent()
    cb.emit("return []")
    cb.dedent()
    cb.dedent()
    cb.emit()
    cb.emit()
    cb.emit("def _step_cmd(")
    cb.indent()
    cb.emit("step_name: str,")
    cb.emit("run_id: str,")
    cb.emit("task_id: str,")
    cb.emit("input_paths: str,")
    cb.emit("retry_count: int = 0,")
    cb.emit("max_user_code_retries: int = 0,")
    cb.emit("split_index: int | None = None,")
    cb.dedent()
    cb.emit(") -> list[str]:")
    cb.indent()
    cb.emit('"""Build the metaflow `step` command list."""')
    cb.emit("cmd: list[str] = [")
    cb.indent()
    cb.emit("sys.executable, FLOW_FILE,")
    cb.emit('"--datastore", DATASTORE_TYPE,')
    cb.emit('"--metadata", METADATA_TYPE,')
    cb.emit('"--no-pylint",')
    cb.emit('"--with=prefect_internal",')
    cb.emit('"step", step_name,')
    cb.emit('"--run-id", run_id,')
    cb.emit('"--task-id", task_id,')
    cb.emit('"--retry-count", str(retry_count),')
    cb.emit('"--max-user-code-retries", str(max_user_code_retries),')
    cb.emit('"--input-paths", input_paths,')
    cb.dedent()
    cb.emit("]")
    cb.emit("for _tag in TAGS:")
    cb.indent()
    cb.emit('cmd += ["--tag", _tag]')
    cb.dedent()
    cb.emit("for _deco in WITH_DECORATORS:")
    cb.indent()
    cb.emit('cmd += [f"--with={_deco}"]')
    cb.dedent()
    cb.emit("if NAMESPACE:")
    cb.indent()
    cb.emit('cmd += ["--namespace", NAMESPACE]')
    cb.dedent()
    cb.emit("if split_index is not None:")
    cb.indent()
    cb.emit('cmd += ["--split-index", str(split_index)]')
    cb.dedent()
    cb.emit("return cmd")
    cb.dedent()


# ---------------------------------------------------------------------------
# Per-step @task functions
# ---------------------------------------------------------------------------


def _task_decorator(step: StepSpec) -> str:
    """Build the @task(...) decorator string for *step*."""
    parts = ['name="%s"' % step.name, "retries=%d" % step.max_user_code_retries]
    if step.timeout_seconds is not None:
        parts.append("timeout_seconds=%d" % step.timeout_seconds)
    if step.retry_delay_seconds is not None:
        parts.append("retry_delay_seconds=%d" % step.retry_delay_seconds)
    return "@task(%s)" % ", ".join(parts)


def _emit_task(cb: _CB, step: StepSpec, spec: FlowSpec, cfg: PrefectFlowConfig) -> None:
    """Emit the ``@task`` function for *step*."""
    cb.emit(_task_decorator(step))

    # The Metaflow "start" step always has step.name == "start".  Its
    # node_type may be START (linear), FOREACH, or SPLIT depending on how
    # self.next() is called — so we use the name, not the type, to detect it.
    is_start = step.name == "start"

    # Immediate child of a FOREACH node: needs split_index to select its item.
    _foreach_body_set = {
        s.out_funcs[0]
        for s in spec.steps
        if s.node_type == NodeType.FOREACH and s.out_funcs
    }
    is_foreach_body = step.name in _foreach_body_set

    # --- function signature ---
    if is_start and step.node_type == NodeType.FOREACH:
        # start step that also fans out via foreach= → returns (task_id, num_splits)
        cb.emit("def %s(run_id: str, parameters: dict[str, Any]) -> tuple[str, int]:" % _task_fn(step.name))
    elif is_start:
        # start step with linear or split next → returns task_id
        cb.emit("def %s(run_id: str, parameters: dict[str, Any]) -> str:" % _task_fn(step.name))
    elif step.is_foreach_join:
        cb.emit("def %s(run_id: str, parent_step: str, task_ids: list[str]) -> str:" % _task_fn(step.name))
    elif step.is_split_join:
        cb.emit("def %s(run_id: str, parent_task_ids: dict[str, str]) -> str:" % _task_fn(step.name))
    elif step.node_type == NodeType.FOREACH:
        cb.emit("def %s(run_id: str, prev_task_id: str) -> tuple[str, int]:" % _task_fn(step.name))
    elif is_foreach_body:
        # body step of a foreach: receives split_index from the list comprehension
        cb.emit("def %s(run_id: str, prev_task_id: str, split_index: int = 0) -> str:" % _task_fn(step.name))
    else:
        # linear, split, end
        cb.emit("def %s(run_id: str, prev_task_id: str) -> str:" % _task_fn(step.name))

    cb.indent()
    cb.emit('logger = get_run_logger()')
    cb.emit("task_id: str = uuid.uuid4().hex[:16]")

    # --- build input_paths ---
    # Each path is run_id/step_name/task_id — comma-separated for joins.
    if is_start:
        _emit_start_init(cb, spec)
        cb.emit("input_paths: str = f\"{run_id}/_parameters/{param_task_id}\"")
    elif step.is_foreach_join:
        # All body task outputs share the same parent step name.
        cb.emit(
            "input_paths: str = \",\".join("
            "f\"{run_id}/{parent_step}/{tid}\" for tid in task_ids)"
        )
    elif step.is_split_join:
        # One path per branch: run_id/branch_step/task_id
        parts = "\",\".join([" + ", ".join(
            'f"{run_id}/%s/{parent_task_ids[%r]}"' % (p, p) for p in step.in_funcs
        ) + "])"
        cb.emit("input_paths: str = %s" % parts)
    else:
        # single upstream: run_id/parent_step/prev_task_id
        parent = step.in_funcs[0] if step.in_funcs else "start"
        cb.emit("input_paths: str = f\"{run_id}/%s/{prev_task_id}\"" % parent)

    # --- foreach side-car setup ---
    if step.node_type == NodeType.FOREACH:
        cb.emit("foreach_path: str = _foreach_info_path(run_id, %r)" % step.name)
        cb.emit("_extra_env: dict[str, str] = {'METAFLOW_PREFECT_FOREACH_INFO_PATH': foreach_path}")
    else:
        cb.emit("_extra_env: dict[str, str] = {}")

    # --- inject Prefect context into subprocess env ---
    cb.emit("try:")
    cb.indent()
    cb.emit("_ctx = get_run_context()")
    cb.emit("_extra_env['METAFLOW_PREFECT_FLOW_RUN_ID'] = str(_ctx.flow_run.id)")
    cb.emit("_extra_env['METAFLOW_PREFECT_TASK_RUN_ID'] = str(_ctx.task_run.id)")
    cb.dedent()
    cb.emit("except Exception:")
    cb.indent()
    cb.emit("pass")
    cb.dedent()

    # --- @environment vars: merge into subprocess env ---
    if step.env_vars:
        cb.emit("_extra_env.update(%r)" % dict(step.env_vars))

    # --- run the step ---
    cb.emit("logger.info(f\"Metaflow step '%s' task_id={task_id}\")" % step.name)
    cb.emit("cmd = _step_cmd(")
    cb.indent()
    cb.emit("%r, run_id, task_id, input_paths," % step.name)
    cb.emit("max_user_code_retries=%d," % step.max_user_code_retries)
    if is_foreach_body:
        cb.emit("split_index=split_index,")
    cb.dedent()
    cb.emit(")")
    cb.emit("_run_cmd(cmd, extra_env=_extra_env)")

    # --- Prefect artifact: list Metaflow self.* keys + retrieval snippet ---
    artifact_key = step.name.replace("_", "-")
    cb.emit("_art_names = _mf_artifact_names(run_id, %r, task_id)" % step.name)
    cb.emit("_md = f\"## `%s` — {run_id}\\n\\n\"" % step.name)
    cb.emit("if _art_names:")
    cb.indent()
    cb.emit("for _n in _art_names:")
    cb.indent()
    cb.emit("_md += f\"`{_n}`\\n\"")
    cb.emit("_md += \"```python\\n\"")
    cb.emit("_md += f\"Task('{FLOW_NAME}/{run_id}/%s/{task_id}')['{_n}'].data\\n\"" % step.name)
    cb.emit("_md += \"```\\n\\n\"")
    cb.dedent()
    cb.dedent()
    cb.emit("else:")
    cb.indent()
    cb.emit("_md += \"*(no user artifacts)*\\n\"")
    cb.dedent()
    cb.emit("create_markdown_artifact(key=%r, markdown=_md)" % artifact_key)

    # --- return value ---
    if step.node_type == NodeType.FOREACH:
        cb.emit("num_splits: int = _read_foreach_info(foreach_path)")
        cb.emit("return task_id, num_splits")
    else:
        cb.emit("return task_id")

    cb.dedent()


def _emit_start_init(cb: _CB, spec: FlowSpec) -> None:
    """Emit the _parameters init call inside the start task."""
    cb.emit("# --- _parameters init task ---")
    cb.emit("param_task_id: str = uuid.uuid4().hex[:16]")
    cb.emit("init_cmd: list[str] = [")
    cb.indent()
    cb.emit("sys.executable, FLOW_FILE,")
    cb.emit('"--datastore", DATASTORE_TYPE,')
    cb.emit('"--metadata", METADATA_TYPE,')
    cb.emit('"--no-pylint",')
    cb.emit('"init",')
    cb.emit('"--run-id", run_id,')
    cb.emit('"--task-id", param_task_id,')
    cb.dedent()
    cb.emit("]")
    cb.emit("for _tag in TAGS:")
    cb.indent()
    cb.emit('init_cmd += ["--tag", _tag]')
    cb.dedent()
    cb.emit("init_env: dict[str, str] = os.environ.copy()")
    cb.emit("if parameters:")
    cb.indent()
    cb.emit('init_env["METAFLOW_PARAMETERS"] = json.dumps(parameters)')
    cb.dedent()
    cb.emit("subprocess.run(init_cmd, env=init_env, check=True)")


# ---------------------------------------------------------------------------
# Top-level @flow function
# ---------------------------------------------------------------------------


def _flow_decorator(spec: FlowSpec, cfg: PrefectFlowConfig) -> str:
    """Build the @flow(...) decorator string."""
    parts = ["name=%r" % spec.name, "description=%r" % (spec.description or spec.name)]
    if cfg.workflow_timeout is not None:
        parts.append("timeout_seconds=%d" % cfg.workflow_timeout)
    return "@flow(%s)" % ", ".join(parts)


def _emit_flow(cb: _CB, spec: FlowSpec, cfg: PrefectFlowConfig) -> None:
    """Emit the top-level ``@flow`` function."""
    # Schedule is registered at deploy time via the CLI; the @flow decorator
    # itself carries only the name and description here.
    cb.emit(_flow_decorator(spec, cfg))
    sig = _flow_signature(spec.parameters)
    cb.emit("def %s(%s) -> None:" % (_python_name(spec.name), sig))
    cb.indent()

    # derive run_id from Prefect's flow-run context
    cb.emit("ctx = get_run_context()")
    cb.emit("run_id: str = f\"prefect-{ctx.flow_run.id}\"")

    # collect parameters into a dict
    if spec.parameters:
        param_items = ", ".join('%r: %s' % (p.name, p.name) for p in spec.parameters)
        cb.emit("parameters: dict[str, Any] = {%s}" % param_items)
    else:
        cb.emit("parameters: dict[str, Any] = {}")

    cb.emit()

    # wire up task calls (topological order)
    task_id_vars: dict[str, str] = {}   # step_name -> Python variable name

    # Track foreach nodes so we can emit dynamic list comprehensions
    foreach_nodes = {s.name for s in spec.steps if s.node_type == NodeType.FOREACH}
    # Map foreach_step_name -> body_step_name
    foreach_body: dict[str, str] = {}
    for s in spec.steps:
        if s.node_type == NodeType.FOREACH and s.out_funcs:
            foreach_body[s.name] = s.out_funcs[0]

    for step in spec.steps:
        var = "_tid_%s" % step.name
        is_start = step.name == "start"

        if is_start and step.node_type == NodeType.FOREACH:
            # start step that is also a foreach fan-out
            cb.emit(
                "%s_pair: tuple[str, int] = %s(run_id, parameters)"
                % (var, _task_fn(step.name))
            )
            cb.emit("%s: str = %s_pair[0]" % (var, var))
            cb.emit("%s_nsplits: int = %s_pair[1]" % (var, var))
            body_name = foreach_body[step.name]
            body_var = "_tid_%s_list" % body_name
            futures_var = "_futures_%s" % body_name
            cb.emit(
                "%s = [%s.submit(run_id, %s, split_index=_i) for _i in range(%s_nsplits)]"
                % (futures_var, _task_fn(body_name), var, var)
            )
            cb.emit("%s: list[str] = [_f.result() for _f in %s]" % (body_var, futures_var))
            task_id_vars[body_name] = body_var

        elif is_start:
            # start step with linear or static-split next
            cb.emit("%s: str = %s(run_id, parameters)" % (var, _task_fn(step.name)))

        elif step.is_foreach_join:
            # The body step immediately precedes this join
            foreach_step = spec.steps[
                next(i for i, s in enumerate(spec.steps) if s.name == step.split_parents[-1])
            ]
            body_name = foreach_body[foreach_step.name]
            body_var = "_tid_%s_list" % body_name
            cb.emit(
                "%s: str = %s(run_id, %r, %s)"
                % (var, _task_fn(step.name), body_name, body_var)
            )

        elif step.is_split_join:
            parent_ids_literal = "{%s}" % ", ".join(
                "%r: %s" % (p, task_id_vars[p]) for p in step.in_funcs
            )
            cb.emit("%s: str = %s(run_id, %s)" % (var, _task_fn(step.name), parent_ids_literal))

        elif step.node_type == NodeType.FOREACH:
            # non-start foreach step always has a predecessor
            parent = step.in_funcs[0]
            cb.emit(
                "%s_pair: tuple[str, int] = %s(run_id, %s)"
                % (var, _task_fn(step.name), task_id_vars[parent])
            )
            cb.emit("%s: str = %s_pair[0]" % (var, var))
            cb.emit("%s_nsplits: int = %s_pair[1]" % (var, var))

            # Immediately emit the concurrent body-step submissions
            body_name = foreach_body[step.name]
            body_var = "_tid_%s_list" % body_name
            futures_var = "_futures_%s" % body_name
            cb.emit(
                "%s = [%s.submit(run_id, %s, split_index=_i) for _i in range(%s_nsplits)]"
                % (futures_var, _task_fn(body_name), var, var)
            )
            cb.emit("%s: list[str] = [_f.result() for _f in %s]" % (body_var, futures_var))
            # register body var so the foreach join can find it
            task_id_vars[body_name] = body_var

        elif step.name in foreach_body.values():
            # This is a foreach body step — already emitted in the foreach block above.
            task_id_vars[step.name] = "_tid_%s_list" % step.name
            continue

        else:
            # linear, split, end
            parent = step.in_funcs[0]
            cb.emit("%s: str = %s(run_id, %s)" % (var, _task_fn(step.name), task_id_vars[parent]))

        task_id_vars[step.name] = var

    cb.dedent()


# ---------------------------------------------------------------------------
# Small helpers
# ---------------------------------------------------------------------------


def _task_fn(step_name: str) -> str:
    return "_step_%s" % step_name


def _python_name(flow_name: str) -> str:
    """Convert CamelCase flow name to snake_case for Python identifiers."""
    name = flow_name
    result = []
    for i, ch in enumerate(name):
        if ch.isupper() and i > 0:
            result.append("_")
        result.append(ch.lower())
    return "".join(result)


def _flow_signature(params: Sequence[ParameterSpec]) -> str:
    """Build the Python function parameter string from flow parameters."""
    parts: list[str] = []
    for p in params:
        if isinstance(p.default, str):
            parts.append('%s: str = %r' % (p.name, p.default))
        elif isinstance(p.default, bool):
            # bool before int — bool IS a subtype of int in Python
            parts.append('%s: bool = %r' % (p.name, p.default))
        elif isinstance(p.default, int):
            parts.append('%s: int = %r' % (p.name, p.default))
        elif isinstance(p.default, float):
            parts.append('%s: float = %r' % (p.name, p.default))
        else:
            parts.append('%s: Any = %r' % (p.name, p.default))
    return ", ".join(parts)
