"""
This module provides functions for rendering Slurm job scripts from task definitions.
"""

import base64
import logging
import pickle
from dataclasses import dataclass, field
from typing import Any, Dict, Tuple, Callable, List, Optional, TYPE_CHECKING

from ._serialization import dumps_pickled
import inspect
import sys
import traceback
from .packaging.base import PackagingStrategy
import pathlib
import shlex
import os
import importlib
from .callbacks.callbacks import BaseCallback
from .validation import validate_job_name

if TYPE_CHECKING:
    from .cluster import Cluster

logger = logging.getLogger(__name__)

# Define the expected output filename for results
RESULT_FILENAME = "result.pkl"
ARGS_FILENAME = "task_args.pkl"
KWARGS_FILENAME = "task_kwargs.pkl"
CALLBACKS_FILENAME = "callbacks.pkl"


def _get_importable_module_name(func: Callable) -> str:
    """Attempts to find the importable module path for a function.

    Special handling for __main__ module: derives module name from file path
    relative to sys.path entries to enable remote imports.
    """
    module_name = func.__module__
    if module_name == "__main__":
        try:
            filepath = inspect.getfile(func)
            abs_filepath = pathlib.Path(filepath).resolve()

            # Build candidate module names from all sys.path parents; prefer dotted package paths
            candidates = []  # (score, derived_module_name)
            for p in sys.path:
                try:
                    abs_sys_path = pathlib.Path(p).resolve()
                    if not abs_filepath.is_relative_to(abs_sys_path):
                        continue
                    relative_path = abs_filepath.relative_to(abs_sys_path)
                    module_parts = list(relative_path.parts)
                    if not module_parts:
                        continue
                    if module_parts[-1].endswith(".py"):
                        module_parts[-1] = module_parts[-1][:-3]
                    if module_parts[-1] == "__init__":
                        module_parts.pop()
                    derived_module_name = ".".join(module_parts)
                    dot_count = derived_module_name.count(".")
                    depth = len(module_parts)
                    score = (dot_count, -depth)
                    candidates.append((score, derived_module_name))
                except (ValueError, OSError, TypeError):
                    continue

            if candidates:
                candidates.sort(reverse=True)
                chosen = candidates[0][1]
                logger.debug(
                    "Derived module name '%s' for function in '__main__' from path %s",
                    chosen,
                    filepath,
                )
                return chosen
            else:
                logger.warning(
                    "Could not determine module path for function in '__main__' from path %s",
                    filepath,
                )

        except (TypeError, OSError) as e:
            logger.warning(
                "Could not inspect function file path for __main__ module: %s",
                e,
            )

        return module_name
    else:
        return module_name


@dataclass
class RenderContext:
    """Structured input for :func:`render_job_script`.

    Groups the parameters needed to render an sbatch script so that
    callers pass a single object instead of 12 positional arguments.
    """

    task_func: Callable[..., Any]
    task_args: Tuple[Any, ...]
    task_kwargs: Dict[str, Any]
    task_definition: Dict[str, Any]
    sbatch_overrides: Dict[str, Any]
    packaging_strategy: PackagingStrategy
    target_job_dir: str
    pre_submission_id: str
    callbacks: List[BaseCallback] = field(default_factory=list)
    cluster: Optional["Cluster"] = None
    is_array_job: bool = False
    array_items_file: Optional[str] = None


def render_job_script(ctx: RenderContext) -> str:
    """Render a SLURM sbatch script from a :class:`RenderContext`."""
    task_func = ctx.task_func
    task_args = ctx.task_args
    task_kwargs = ctx.task_kwargs
    packaging_strategy = ctx.packaging_strategy
    target_job_dir = ctx.target_job_dir
    pre_submission_id = ctx.pre_submission_id
    callbacks = ctx.callbacks
    cluster = ctx.cluster
    is_array_job = ctx.is_array_job
    array_items_file = ctx.array_items_file

    sbatch_params = dict(ctx.task_definition)
    sbatch_params.update(ctx.sbatch_overrides)

    # For array jobs, use %a (array task ID) in output/error paths
    if is_array_job:
        stdout_path = sbatch_params.get("output")
        if not stdout_path:
            stdout_path = f"{target_job_dir}/slurm_{pre_submission_id}_%a.out"
            sbatch_params["output"] = stdout_path

        stderr_path = sbatch_params.get("error")
        if not stderr_path:
            stderr_path = f"{target_job_dir}/slurm_{pre_submission_id}_%a.err"
            sbatch_params["error"] = stderr_path
    else:
        stdout_path = sbatch_params.get("output")
        if not stdout_path:
            stdout_path = f"{target_job_dir}/slurm_{pre_submission_id}.out"
            sbatch_params["output"] = stdout_path

        stderr_path = sbatch_params.get("error")
        if not stderr_path:
            stderr_path = f"{target_job_dir}/slurm_{pre_submission_id}.err"
            sbatch_params["error"] = stderr_path

    script_lines = ["#!/bin/bash"]

    job_name = sbatch_params.pop("job_name", None)
    if not job_name:
        job_name = task_func.__name__
    validate_job_name(job_name)
    script_lines.append(f"#SBATCH --job-name={shlex.quote(job_name)}")

    output_path = sbatch_params.get("output")
    error_path = sbatch_params.get("error")
    if not output_path:
        if is_array_job:
            output_path = (
                f"{shlex.quote(target_job_dir)}/slurm_{pre_submission_id}_%a.out"
            )
        else:
            output_path = f"{shlex.quote(target_job_dir)}/slurm_{pre_submission_id}.out"
        sbatch_params["output"] = output_path
    if not error_path:
        if is_array_job:
            error_path = (
                f"{shlex.quote(target_job_dir)}/slurm_{pre_submission_id}_%a.err"
            )
        else:
            error_path = f"{shlex.quote(target_job_dir)}/slurm_{pre_submission_id}.err"
        sbatch_params["error"] = error_path

    # Emit remaining directives in insertion order
    for key, value in sbatch_params.items():
        if key == "job_name":
            continue
        flag = key.replace("_", "-")
        value_to_emit = value
        if isinstance(value, str) and key in {"output", "error"}:
            value_to_emit = shlex.quote(value)
        if value_to_emit is None:
            script_lines.append(f"#SBATCH --{flag}")
        else:
            script_lines.append(f"#SBATCH --{flag}={value_to_emit}")

    script_lines.append("")
    script_lines.append(f'echo "Target Job Directory (from Python): {target_job_dir}"')
    script_lines.append(f"export JOB_DIR={shlex.quote(target_job_dir)}")
    # Export the job base directory for the runner to find result files from dependent jobs
    script_lines.append("export SLURM_JOBS_DIR=$(dirname $(dirname $JOB_DIR))")

    # Export cluster configuration for workflow support
    if cluster is not None:
        slurmfile_path = getattr(cluster, "slurmfile_path", None)
        env_name = getattr(cluster, "env_name", None)
        # The Slurmfile will be uploaded to the job directory as "Slurmfile.toml"
        is_workflow = getattr(task_func, "_is_workflow", False)
        if is_workflow or slurmfile_path:
            remote_slurmfile_path = f"{target_job_dir}/Slurmfile.toml"
            script_lines.append(
                f"export SLURM_SDK_SLURMFILE={shlex.quote(remote_slurmfile_path)}"
            )
        if env_name:
            script_lines.append(f"export SLURM_SDK_ENV={shlex.quote(env_name)}")

    # Export packaging configuration for child tasks to inherit
    # This is especially important for container packaging where the image reference
    # needs to be passed to nested tasks
    if packaging_strategy is not None:
        import json

        packaging_config = getattr(packaging_strategy, "config", {})
        if packaging_config:
            # Include the packaging type
            config_to_export = dict(packaging_config)
            config_to_export["type"] = (
                getattr(packaging_strategy, "__class__", type(packaging_strategy))
                .__name__.replace("PackagingStrategy", "")
                .lower()
            )

            # For container packaging, clean up the config for child tasks
            # Remove build-time fields and set the resolved image reference
            if hasattr(packaging_strategy, "_image_reference"):
                config_to_export["image"] = packaging_strategy._image_reference
                # Remove fields that would trigger a rebuild
                # Child tasks should use the pre-built image directly
                for key in ["dockerfile", "context", "registry", "name", "tag"]:
                    config_to_export.pop(key, None)
                # Don't push - image already exists in registry
                config_to_export["push"] = False

            # Serialize and encode for safe shell export
            config_json = json.dumps(config_to_export)
            config_b64 = base64.b64encode(config_json.encode()).decode()
            script_lines.append(
                f"export SLURM_SDK_PACKAGING_CONFIG={shlex.quote(config_b64)}"
            )

    # Export pre-built dependency images for workflows
    # This allows child tasks to use pre-built images instead of inheriting
    if cluster is not None:
        prebuilt_images = getattr(cluster, "_prebuilt_dependency_images", None)
        if prebuilt_images:
            import json

            images_json = json.dumps(prebuilt_images)
            images_b64 = base64.b64encode(images_json.encode()).decode()
            script_lines.append(
                f"export SLURM_SDK_PREBUILT_IMAGES={shlex.quote(images_b64)}"
            )

    script_lines.append("")
    script_lines.append('echo "SLURM Job ID: ${SLURM_JOB_ID:-}"')
    script_lines.append('echo "Running on host: $(hostname)"')
    script_lines.append('echo "Initial working directory: $(pwd)"')
    script_lines.append('echo "Job output directory JOB_DIR: $JOB_DIR"')
    script_lines.append("")
    script_lines.append("# Change to job directory")
    script_lines.append(
        "cd $JOB_DIR || { echo 'ERROR: Failed to cd to $JOB_DIR' >&2; exit 1; }"
    )
    script_lines.append('echo "Changed to working directory: $(pwd)"')
    script_lines.append("")

    try:
        setup_commands = packaging_strategy.generate_setup_commands(
            task=task_func,
            job_id=pre_submission_id,
            job_dir="$JOB_DIR",
        )
        if setup_commands:
            script_lines.extend(setup_commands)
            script_lines.append("echo 'Packaging setup commands executed.'")
        else:
            script_lines.append("echo 'No packaging setup commands generated.'")
    except Exception as e:
        logger.error("Error generating packaging setup commands: %s", e)
        traceback.print_exc(file=sys.stderr)
        script_lines.append(
            "echo 'ERROR: Failed to generate packaging setup commands!' >&2"
        )
        script_lines.append("exit 1")

    script_lines.append("")

    module_name = _get_importable_module_name(task_func)
    func_name = task_func.__name__

    # For array jobs, use %a in result filenames for per-task results
    if is_array_job:
        args_file_path_str = None  # Not used for array jobs
        kwargs_file_path_str = None  # Not used for array jobs
        result_file_path_str = f"slurm_job_{pre_submission_id}_%a_{RESULT_FILENAME}"
    else:
        args_file_path_str = f"slurm_job_{pre_submission_id}_{ARGS_FILENAME}"
        kwargs_file_path_str = f"slurm_job_{pre_submission_id}_{KWARGS_FILENAME}"
        result_file_path_str = f"slurm_job_{pre_submission_id}_{RESULT_FILENAME}"

    callbacks_file_path_str = f"slurm_job_{pre_submission_id}_{CALLBACKS_FILENAME}"

    try:
        # For array jobs, skip args/kwargs pickling (loaded from array_items_file)
        if not is_array_job:
            pickled_args = base64.b64encode(dumps_pickled(task_args)).decode()
            pickled_kwargs = base64.b64encode(dumps_pickled(task_kwargs)).decode()
        pickled_sys_path = base64.b64encode(dumps_pickled(sys.path)).decode()

        picklable_callbacks = []
        for cb in callbacks or []:
            needs_runner = True
            if hasattr(cb, "requires_runner_transport"):
                try:
                    needs_runner = cb.requires_runner_transport()
                except Exception as _cb_err:  # pragma: no cover - defensive
                    logger.debug(
                        "Callback %s failed requires_runner_transport check: %s",
                        type(cb).__name__,
                        _cb_err,
                    )
            if not needs_runner:
                continue
            if getattr(cb, "requires_pickling", True) is False:
                logger.debug(
                    "Skipping callback %s: requires_pickling=False",
                    type(cb).__name__,
                )
                continue
            try:
                pickle.dumps(cb)
                picklable_callbacks.append(cb)
            except Exception as _cb_err:
                logger.debug(
                    "Skipping non-picklable callback %s: %s", type(cb).__name__, _cb_err
                )

        pickled_callbacks = (
            base64.b64encode(dumps_pickled(picklable_callbacks)).decode()
            if picklable_callbacks
            else ""
        )

    except Exception as e:
        raise RuntimeError(
            f"Failed to serialize task arguments for cluster execution.\n\n"
            f"Error: {e}\n\n"
            "This usually means one or more of your task arguments cannot be pickled.\n"
            "Common non-picklable objects include:\n"
            "  - Open file handles or database connections\n"
            "  - Lambda functions or local functions\n"
            "  - Objects with __getstate__ that raises errors\n"
            "  - Thread locks or multiprocessing primitives\n\n"
            "To fix:\n"
            "  1. Pass file paths instead of open file objects\n"
            "  2. Use module-level functions instead of lambdas\n"
            "  3. Ensure all arguments are standard Python types or pickle-compatible objects"
        ) from e

    # For regular jobs, write args/kwargs to files
    # For array jobs, skip this (args/kwargs come from array_items_file)
    if not is_array_job:
        script_lines.append(f'base64 -d > "{args_file_path_str}" << "BASE64_ARGS"')
        script_lines.append(pickled_args)
        script_lines.append("BASE64_ARGS")
        script_lines.append(f'base64 -d > "{kwargs_file_path_str}" << "BASE64_KWARGS"')
        script_lines.append(pickled_kwargs)
        script_lines.append("BASE64_KWARGS")

    if pickled_callbacks:
        script_lines.append(f'base64 -d > "{callbacks_file_path_str}" << "BASE64_CBS"')
        script_lines.append(pickled_callbacks)
        script_lines.append("BASE64_CBS")
    else:
        script_lines.append(f'touch "{callbacks_file_path_str}"')
    script_lines.append("")

    script_lines.append("echo 'Executing Python task via packaged runner...' ")
    submission_sys_path = [p for p in sys.path if isinstance(p, str) and p]
    repo_root = os.getcwd()
    if repo_root not in submission_sys_path:
        submission_sys_path.insert(0, repo_root)
    try:
        _slurm_mod = importlib.import_module("slurm")
        slurm_parent = pathlib.Path(_slurm_mod.__file__).resolve().parent.parent
        slurm_parent_str = str(slurm_parent)
        if slurm_parent_str not in submission_sys_path:
            submission_sys_path.insert(0, slurm_parent_str)
    except Exception:
        pass
    pythonpath_contrib = ":".join(submission_sys_path)
    if pythonpath_contrib:
        script_lines.append(
            f"export PYTHONPATH={shlex.quote(pythonpath_contrib)}:${{PYTHONPATH:-}}"
        )

    script_lines.append("PY_EXEC_RESOLVED=${PY_EXEC:-python}")
    script_lines.append("export PY_EXEC_RESOLVED")

    # For array jobs, set up filenames with bash variable substitution (expand %a to actual array task ID)
    if is_array_job:
        script_lines.append("")
        script_lines.append("# Set up filenames for array job (expand array task ID)")
        script_lines.append(
            f'RESULT_FILE="slurm_job_{pre_submission_id}_${{SLURM_ARRAY_TASK_ID}}_{RESULT_FILENAME}"'
        )
        if stdout_path and "%a" in stdout_path:
            expanded_stdout = stdout_path.replace("%a", "${SLURM_ARRAY_TASK_ID}")
            script_lines.append(f'STDOUT_PATH="{expanded_stdout}"')
        if stderr_path and "%a" in stderr_path:
            expanded_stderr = stderr_path.replace("%a", "${SLURM_ARRAY_TASK_ID}")
            script_lines.append(f'STDERR_PATH="{expanded_stderr}"')

    # Construct the runner command that executes the user's task function.
    def _escape_quotes(value: str) -> str:
        return value.replace('"', '\\"')

    runner_parts = [
        '"$PY_EXEC_RESOLVED"',
        "-m slurm.runner",
        f'--module "{_escape_quotes(module_name)}"',
        f'--function "{_escape_quotes(func_name)}"',
    ]

    # For array jobs, provide array-specific arguments
    if is_array_job:
        runner_parts.append('--array-index "$SLURM_ARRAY_TASK_ID"')
        runner_parts.append(f'--array-items-file "{_escape_quotes(array_items_file)}"')
    else:
        # For regular jobs, provide args/kwargs files
        runner_parts.append(f'--args-file "{_escape_quotes(args_file_path_str)}"')
        runner_parts.append(f'--kwargs-file "{_escape_quotes(kwargs_file_path_str)}"')

    # Common arguments for both job types
    if is_array_job:
        # Use bash variable for array jobs (already expanded)
        runner_parts.append('--output-file "$RESULT_FILE"')
    else:
        runner_parts.append(f'--output-file "{_escape_quotes(result_file_path_str)}"')

    runner_parts.extend(
        [
            f'--callbacks-file "{_escape_quotes(callbacks_file_path_str)}"',
            f'--sys-path "{_escape_quotes(pickled_sys_path)}"',
        ]
    )

    if target_job_dir:
        runner_parts.append(f'--job-dir "{_escape_quotes(target_job_dir)}"')
    if stdout_path:
        if is_array_job and "%a" in stdout_path:
            runner_parts.append('--stdout-path "$STDOUT_PATH"')
        else:
            runner_parts.append(f'--stdout-path "{_escape_quotes(stdout_path)}"')
    if stderr_path:
        if is_array_job and "%a" in stderr_path:
            runner_parts.append('--stderr-path "$STDERR_PATH"')
        else:
            runner_parts.append(f'--stderr-path "{_escape_quotes(stderr_path)}"')
    runner_parts.append(f'--pre-submission-id "{_escape_quotes(pre_submission_id)}"')

    runner_command = " ".join(runner_parts)

    try:
        wrapped_command = packaging_strategy.wrap_execution_command(
            command=runner_command,
            task=task_func,
            job_id=pre_submission_id,
            job_dir='"$JOB_DIR"',
        )
        # Add debug output for container packaging to verify image reference
        if (
            hasattr(packaging_strategy, "_image_reference")
            and packaging_strategy._image_reference
        ):
            script_lines.append(
                f"echo 'Executing with container image: {packaging_strategy._image_reference}'"
            )
        runner_command = wrapped_command
    except Exception as e:
        logger.error("Error wrapping execution command: %s", e)
        traceback.print_exc(file=sys.stderr)
        script_lines.append("echo 'ERROR: Failed to wrap execution command!' >&2")
        script_lines.append("exit 1")
    else:
        script_lines.append(runner_command)
    script_lines.append("EXECUTION_STATUS=$?")
    script_lines.append("")

    # Skip cleanup for native array jobs since all tasks share the same environment
    if not is_array_job:
        try:
            cleanup_commands = packaging_strategy.generate_cleanup_commands(
                task=task_func,
                job_id=pre_submission_id,
                job_dir="$JOB_DIR",
            )
            if cleanup_commands:
                script_lines.extend(cleanup_commands)
                script_lines.append("echo 'Packaging cleanup commands executed.'")
            else:
                script_lines.append("echo 'No packaging cleanup commands generated.'")
        except Exception as e:
            logger.error("Error generating packaging cleanup commands: %s", e)
            traceback.print_exc(file=sys.stderr)
            script_lines.append(
                "echo 'WARNING: Failed to generate packaging cleanup commands!' >&2"
            )

    script_lines.append("")

    script_lines.append('echo "Job finished with status: $EXECUTION_STATUS"')
    script_lines.append("exit $EXECUTION_STATUS")

    final_script = "\n".join(
        line.rstrip("\r") for line in "\n".join(script_lines).splitlines()
    )
    return final_script
