###############################################################################
# (c) Copyright 2026 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
"""Rewrite CWL workflows to support custom step inputs with conditional execution.

This module rewrites generated CWL workflows before local execution so that
upstream steps can be skipped when a user provides custom input for a
downstream step. It uses CWL v1.2 features:

- ``when`` clauses to conditionally skip upstream steps
- ``pickValue: first_non_null`` to select between custom input and upstream output

This rewriting is done locally in LbAPLocal and does NOT affect the production
CWL generated by LbAPCommon. Production workflows remain clean and unchanged.

FIXME: This is brittle and should eventually be integrated into the CWL generation
logic in LbAPCommon.
"""

from __future__ import annotations

import logging
from pathlib import Path

from ruamel.yaml import YAML

logger = logging.getLogger(__name__)

CUSTOM_INPUT_ID = "custom-input"


def rewrite_cwl_for_custom_input(
    cwl_path: Path,
    custom_input_files: list[str],
    job_name: str,
) -> tuple[Path, Path]:
    """Rewrite a CWL workflow to support custom input for a specific step.

    When a user provides custom input files (e.g. pre-existing SPRUCE.DST),
    this function modifies the CWL workflow so that:

    1. Upstream steps that produce the input are skipped (via ``when``)
    2. The target step receives the custom input (via ``pickValue``)

    Args:
        cwl_path: Path to the generated CWL workflow file.
        custom_input_files: List of custom input file paths (local or LFN:).
        job_name: The job name from the CLI, used to identify the target step.

    Returns:
        Tuple of (modified_cwl_path, inputs_yaml_path).
    """
    yaml = YAML()
    yaml.default_flow_style = False
    yaml.width = 120

    with open(cwl_path) as f:
        cwl_data = yaml.load(f)

    # Find the target step and its containing transformation
    transform_idx, target_step_id = _find_target_step(cwl_data, job_name)
    transform_step = cwl_data["steps"][transform_idx]
    sub_workflow = transform_step["run"]

    logger.info(
        "Target step: %s (in transformation %d)",
        target_step_id,
        transform_idx + 1,
    )

    # Find upstream steps to skip
    upstream_ids = _find_upstream_steps(sub_workflow["steps"], target_step_id)
    logger.info("Upstream steps to skip: %s", upstream_ids)

    if not upstream_ids:
        logger.warning(
            "No upstream steps found for %s - custom input will replace "
            "the workflow input directly",
            target_step_id,
        )

    # --- 1. Add custom-input parameter at each nesting level ---

    # Main workflow inputs
    _add_optional_input(cwl_data["inputs"], CUSTOM_INPUT_ID)
    _make_input_optional(cwl_data["inputs"], "input-data")

    # Wire through to transformation step
    _add_step_input(transform_step["in"], CUSTOM_INPUT_ID, CUSTOM_INPUT_ID)

    # Transformation sub-workflow inputs
    _add_optional_input(sub_workflow["inputs"], CUSTOM_INPUT_ID)
    _make_input_optional(sub_workflow["inputs"], "input-data")

    # --- 2. Add `when` clauses to upstream steps ---

    for step in sub_workflow["steps"]:
        if step["id"] in upstream_ids:
            step["when"] = f"$(inputs['{CUSTOM_INPUT_ID}'] === null)"
            _add_step_input(step["in"], CUSTOM_INPUT_ID, CUSTOM_INPUT_ID)
            logger.info("Added when clause to step: %s", step["id"])

    # --- 3. Update target step's input-data to use pickValue ---

    for step in sub_workflow["steps"]:
        if step["id"] == target_step_id:
            for inp in step["in"]:
                if inp["id"] == "input-data":
                    original_source = inp["source"]
                    inp["source"] = [CUSTOM_INPUT_ID, original_source]
                    inp["pickValue"] = "first_non_null"
                    logger.info(
                        "Updated %s input-data: pickValue from [%s, %s]",
                        target_step_id,
                        CUSTOM_INPUT_ID,
                        original_source,
                    )
                    break

    # --- 4. Make outputs from skipped steps nullable ---

    for output in sub_workflow.get("outputs", []):
        output_source = output.get("outputSource")
        if output_source is None:
            continue
        # Check if outputSource references a skipped step
        source_str = output_source if isinstance(output_source, str) else None
        source_list = output_source if isinstance(output_source, list) else []

        sources_to_check = [source_str] if source_str else source_list
        for src in sources_to_check:
            if src and src.split("/")[0] in upstream_ids:
                _make_output_type_nullable(output)
                break

    # --- 5. Write modified CWL ---

    modified_cwl_path = cwl_path.parent / (cwl_path.stem + "_custom_input.cwl")
    with open(modified_cwl_path, "w") as f:
        yaml.dump(cwl_data, f)
    logger.info("Written modified CWL: %s", modified_cwl_path)

    # --- 6. Write inputs YAML ---

    inputs_yaml = _build_inputs_yaml(custom_input_files)
    inputs_yaml_path = cwl_path.parent / (cwl_path.stem + "_custom_inputs.yml")
    with open(inputs_yaml_path, "w") as f:
        yaml.dump(inputs_yaml, f)
    logger.info("Written inputs YAML: %s", inputs_yaml_path)

    return modified_cwl_path, inputs_yaml_path


def _find_target_step(cwl_data: dict, job_name: str) -> tuple[int, str]:
    """Find the target step within the workflow.

    Searches nested transformation sub-workflows for a step whose ID
    contains the job_name. Falls back to the last leaf step if no
    exact match is found.

    Returns:
        Tuple of (transformation_index, step_id).

    Raises:
        ValueError: If no suitable target step can be found.
    """
    # Normalise job_name for matching (replace commas/hyphens)
    job_name_normalised = job_name.replace(",", "_").replace("-", "_")

    for t_idx, transform_step in enumerate(cwl_data.get("steps", [])):
        sub_wf = transform_step.get("run", {})
        if sub_wf.get("class") != "Workflow":
            continue

        sub_steps = sub_wf.get("steps", [])
        if not sub_steps:
            continue

        # Try to match step ID against job_name
        for step in sub_steps:
            step_id = step["id"]
            step_id_normalised = step_id.replace("-", "_")
            if step_id_normalised.endswith(job_name_normalised):
                return t_idx, step_id

    # Fallback: use the last leaf step in the first transformation
    for t_idx, transform_step in enumerate(cwl_data.get("steps", [])):
        sub_wf = transform_step.get("run", {})
        if sub_wf.get("class") != "Workflow":
            continue

        sub_steps = sub_wf.get("steps", [])
        if not sub_steps:
            continue

        leaf_steps = _find_leaf_steps(sub_steps)
        if leaf_steps:
            target = leaf_steps[-1]
            logger.warning(
                "Could not match job_name '%s' to a step ID; "
                "falling back to last leaf step '%s'",
                job_name,
                target,
            )
            return t_idx, target

    raise ValueError(
        f"Could not find a target step for job_name '{job_name}' "
        f"in the CWL workflow"
    )


def _find_leaf_steps(steps: list[dict]) -> list[str]:
    """Find steps that no other step depends on (leaf nodes)."""
    all_ids = {step["id"] for step in steps}
    depended_on = set()

    for step in steps:
        for inp in step.get("in", []):
            source = inp.get("source", "")
            sources = [source] if isinstance(source, str) else source
            for src in sources:
                if "/" in src:
                    dep_id = src.split("/")[0]
                    if dep_id in all_ids:
                        depended_on.add(dep_id)

    return [s["id"] for s in steps if s["id"] not in depended_on]


def _find_upstream_steps(steps: list[dict], target_step_id: str) -> set[str]:
    """Find all steps upstream of the target (direct and transitive).

    Traces ``source`` references in step ``in`` bindings to build the
    dependency graph and walks it backwards from the target.
    """
    # Build dependency graph: step_id -> set of step_ids it depends on
    all_ids = {step["id"] for step in steps}
    deps: dict[str, set[str]] = {step["id"]: set() for step in steps}

    for step in steps:
        step_id = step["id"]
        for inp in step.get("in", []):
            source = inp.get("source", "")
            sources = [source] if isinstance(source, str) else source
            for src in sources:
                if "/" in src:
                    dep_id = src.split("/")[0]
                    if dep_id in all_ids:
                        deps[step_id].add(dep_id)

    # Walk backwards from target to find all ancestors
    ancestors: set[str] = set()
    to_visit = list(deps.get(target_step_id, set()))
    while to_visit:
        current = to_visit.pop()
        if current not in ancestors:
            ancestors.add(current)
            to_visit.extend(deps.get(current, set()))

    return ancestors


def _add_optional_input(inputs: list[dict], input_id: str) -> None:
    """Add an optional File[]? input to a workflow's inputs list."""
    # Check if already present
    for inp in inputs:
        if inp["id"] == input_id:
            return

    inputs.append(
        {
            "id": input_id,
            "doc": "Custom input files (skips upstream steps when provided)",
            "type": ["null", {"type": "array", "items": "File"}],
            "default": None,
        }
    )


def _make_input_optional(inputs: list[dict], input_id: str) -> None:
    """Change an input from required (File[]) to optional (File[]?)."""
    for inp in inputs:
        if inp["id"] == input_id:
            current_type = inp.get("type")
            if current_type == "File[]":
                inp["type"] = ["null", {"type": "array", "items": "File"}]
                inp["default"] = None
            break


def _add_step_input(step_ins: list[dict], input_id: str, source: str) -> None:
    """Add an input binding to a workflow step's ``in`` list."""
    for inp in step_ins:
        if inp["id"] == input_id:
            return
    step_ins.append({"id": input_id, "source": source})


def _make_output_type_nullable(output: dict) -> None:
    """Make an output type nullable if it isn't already."""
    current_type = output.get("type")
    if current_type is None:
        return

    # Already nullable (union type with null)
    if isinstance(current_type, list):
        return
    if isinstance(current_type, dict) and current_type.get("type") == "array":
        items = current_type.get("items", [])
        if isinstance(items, list) and "null" in items:
            return

    # Convert File[] to File[]?
    if current_type == "File[]":
        output["type"] = ["null", {"type": "array", "items": "File"}]


def _build_inputs_yaml(custom_input_files: list[str]) -> dict:
    """Build the CWL inputs YAML dict with custom input files."""
    file_entries = []
    for filepath in custom_input_files:
        if filepath.startswith("LFN:"):
            file_entries.append({"class": "File", "location": filepath})
        else:
            # Local file - use absolute path
            abs_path = str(Path(filepath).resolve())
            file_entries.append({"class": "File", "path": abs_path})

    return {CUSTOM_INPUT_ID: file_entries}
