# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""BYOB containerization: build Docker images from compiled BYOB benchmarks.

Generates launcher-compliant Docker images that embed the compiled
benchmark package, code, and data so they can be run as standalone
NeMo Evaluator containers.
"""

import copy
import os
import shutil
import subprocess
from pathlib import Path
from typing import List, Optional

import yaml

from nemo_evaluator.contrib.byob.defaults import DEFAULT_BASE_IMAGE
from nemo_evaluator.logging import get_logger

logger = get_logger(__name__)


def _find_nemo_evaluator_package_root() -> Optional[Path]:
    """Find the root directory of the nemo-evaluator package.

    Walks up from ``nemo_evaluator.__file__`` looking for ``pyproject.toml``.
    Returns None if the package is installed from a wheel (no source tree).
    """
    try:
        import nemo_evaluator

        pkg_init = Path(nemo_evaluator.__file__)
        # Walk up: __init__.py -> nemo_evaluator/ -> src/ -> package root
        for parent in pkg_init.parents:
            if (parent / "pyproject.toml").exists():
                return parent
    except (ImportError, AttributeError):
        pass
    return None


def generate_dockerfile(
    pkg_name: str,
    base_image: str = DEFAULT_BASE_IMAGE,
    user_requirements: Optional[List[str]] = None,
) -> str:
    """Generate a Dockerfile for a compiled BYOB benchmark package.

    The image layout follows the NeMo Evaluator launcher convention:

    - ``/opt/byob/code/`` — benchmark Python module(s)
    - ``/opt/byob/data/``  — dataset file(s)
    - ``/opt/byob_pkg/``   — compiled nemo_evaluator namespace package

    Args:
        pkg_name: Package name (e.g. ``byob_boolq``).
        base_image: Base Docker image.
        user_requirements: Extra pip requirements from the benchmark.

    Returns:
        Dockerfile content as a string.
    """
    user_reqs_line = ""
    if user_requirements:
        user_reqs_line = (
            "COPY requirements.txt /tmp/requirements.txt\n"
            "RUN pip install --no-cache-dir -r /tmp/requirements.txt"
        )

    return f"""\
# Auto-generated by nemo-evaluator-byob containerize
FROM {base_image}

RUN pip install --no-cache-dir --upgrade pip

# Install nemo-evaluator (with BYOB contrib modules)
COPY nemo_eval_pkg/ /opt/nemo_eval_pkg/
RUN pip install --no-cache-dir /opt/nemo_eval_pkg/

# Install compiled BYOB package
COPY pkg/ /opt/byob_pkg/
RUN pip install --no-cache-dir /opt/byob_pkg/

# Copy benchmark code and data
COPY code/ /opt/byob/code/
COPY data/ /opt/byob/data/

# Install user requirements
{user_reqs_line}

# Copy framework metadata for harness discovery
RUN python -c "from nemo_evaluator.core.input import _copy_fdfs; _copy_fdfs('/opt/metadata/')"

# Verify package import
RUN python -c "import nemo_evaluator.{pkg_name}; import nemo_evaluator.contrib.byob.runner"

# Labels for launcher compliance
LABEL com.nvidia.nemo-evaluator.pkg-name="{pkg_name}"
LABEL com.nvidia.nemo-evaluator.integration-type=".nemo_evaluator"
"""


def rewrite_fdf_paths(fdf: dict, pkg_name: str) -> dict:
    """Rewrite host-local paths in an FDF dict to container paths.

    Transforms ``extra.benchmark_module`` and ``extra.dataset`` from
    absolute host paths to container-relative paths under ``/opt/byob/``.

    Args:
        fdf: Framework Definition Format dict (will NOT be mutated).
        pkg_name: Package name for context.

    Returns:
        New FDF dict with rewritten paths.
    """
    fdf = copy.deepcopy(fdf)
    extra = fdf.get("defaults", {}).get("config", {}).get("params", {}).get("extra", {})

    benchmark_module = extra.get("benchmark_module", "")
    if benchmark_module:
        filename = os.path.basename(benchmark_module)
        extra["benchmark_module"] = f"/opt/byob/code/{filename}"

    dataset = extra.get("dataset", "")
    if dataset and not dataset.startswith(
        ("hf://", "s3://", "gs://", "http://", "https://")
    ):
        filename = os.path.basename(dataset)
        extra["dataset"] = f"/opt/byob/data/{filename}"

    return fdf


def prepare_build_context(
    pkg_dir: str,
    fdf: dict,
    context_dir: str,
) -> str:
    """Prepare the Docker build context directory.

    Creates the following layout inside ``context_dir``::

        context_dir/
            pkg/           <- compiled namespace package
            code/          <- benchmark .py module
            data/          <- dataset .jsonl file

    Also writes the rewritten ``framework.yml`` into the copied package.

    Args:
        pkg_dir: Path to the installed compiled package directory.
        fdf: Original FDF dict (host paths).
        context_dir: Target build context directory (created if needed).

    Returns:
        Path to the context directory.
    """
    context = Path(context_dir)
    context.mkdir(parents=True, exist_ok=True)

    # Copy compiled package
    pkg_dest = context / "pkg"
    if pkg_dest.exists():
        shutil.rmtree(pkg_dest)
    shutil.copytree(pkg_dir, str(pkg_dest))

    extra = fdf.get("defaults", {}).get("config", {}).get("params", {}).get("extra", {})

    # Copy benchmark module to code/
    code_dir = context / "code"
    code_dir.mkdir(parents=True, exist_ok=True)
    benchmark_module = extra.get("benchmark_module", "")
    if benchmark_module and os.path.isfile(benchmark_module):
        shutil.copy2(
            benchmark_module, str(code_dir / os.path.basename(benchmark_module))
        )

    # Copy or fetch dataset to data/
    data_dir = context / "data"
    data_dir.mkdir(parents=True, exist_ok=True)
    dataset = extra.get("dataset", "")
    if dataset:
        if os.path.isfile(dataset):
            # Local file — copy directly
            shutil.copy2(dataset, str(data_dir / os.path.basename(dataset)))
        elif dataset.startswith(("hf://", "s3://", "gs://", "http://", "https://")):
            # Remote URI — fetch via dataset fetcher, then copy the result
            from nemo_evaluator.contrib.byob.dataset import get_fetcher_for_uri

            try:
                fetcher = get_fetcher_for_uri(dataset)
                result = fetcher.fetch(dataset, cache_dir=data_dir)
                # Update the FDF to point to the local filename inside /opt/byob/data/
                local_name = result.local_path.name
                extra["dataset"] = f"/opt/byob/data/{local_name}"
                # Move/copy if fetched to a different location than data_dir
                if result.local_path.parent != data_dir:
                    shutil.copy2(str(result.local_path), str(data_dir / local_name))
                logger.info(
                    "Fetched remote dataset for containerization",
                    uri=dataset,
                    local=str(data_dir / local_name),
                )
            except Exception as e:
                raise RuntimeError(
                    f"Cannot fetch dataset '{dataset}' for containerization: {e}"
                ) from e

    # Copy nemo-evaluator package so the container has the BYOB runner
    nemo_eval_pkg_dir = _find_nemo_evaluator_package_root()
    if nemo_eval_pkg_dir:
        nemo_dest = context / "nemo_eval_pkg"
        if nemo_dest.exists():
            shutil.rmtree(nemo_dest)
        shutil.copytree(
            nemo_eval_pkg_dir,
            str(nemo_dest),
            ignore=shutil.ignore_patterns(
                "__pycache__",
                "*.pyc",
                ".git",
                "*.egg-info",
                "tests",
                "test_*",
                ".pytest_cache",
            ),
        )

    # Rewrite FDF and write into the copied package
    pkg_name = fdf.get("framework", {}).get("pkg_name", "")
    rewritten_fdf = rewrite_fdf_paths(fdf, pkg_name)
    framework_yml_candidates = list(
        (pkg_dest / "nemo_evaluator").rglob("framework.yml")
    )
    for fw_path in framework_yml_candidates:
        with open(fw_path, "w") as f:
            yaml.safe_dump(rewritten_fdf, f, default_flow_style=False, sort_keys=False)

    return str(context)


def build_image(
    context_dir: str,
    tag: str,
    base_image: str = DEFAULT_BASE_IMAGE,
    pkg_name: str = "",
    user_requirements: Optional[List[str]] = None,
) -> str:
    """Build a Docker image from a prepared build context.

    Generates the Dockerfile, writes it into the context, and shells
    out to ``docker build``.

    Args:
        context_dir: Path to the build context prepared by
            :func:`prepare_build_context`.
        tag: Docker image tag (e.g. ``"myrepo/byob_boolq:latest"``).
        base_image: Base Docker image.
        pkg_name: Package name for Dockerfile generation and labels.
        user_requirements: Extra pip requirements.

    Returns:
        The image tag on success.

    Raises:
        RuntimeError: If ``docker build`` fails.
    """
    dockerfile_content = generate_dockerfile(
        pkg_name=pkg_name,
        base_image=base_image,
        user_requirements=user_requirements,
    )
    dockerfile_path = os.path.join(context_dir, "Dockerfile")
    with open(dockerfile_path, "w") as f:
        f.write(dockerfile_content)

    # Write requirements.txt if user requirements are specified
    if user_requirements:
        reqs_path = os.path.join(context_dir, "requirements.txt")
        with open(reqs_path, "w") as f:
            f.write("\n".join(user_requirements) + "\n")

    logger.info("Building Docker image", tag=tag, context=context_dir)
    result = subprocess.run(
        ["docker", "build", "-t", tag, context_dir],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(
            f"docker build failed (exit {result.returncode}):\n{result.stderr}"
        )

    logger.info("Docker image built successfully", tag=tag)
    return tag


def push_image(tag: str) -> None:
    """Push a Docker image to a registry.

    Args:
        tag: Docker image tag to push.

    Raises:
        RuntimeError: If ``docker push`` fails.
    """
    logger.info("Pushing Docker image", tag=tag)
    result = subprocess.run(
        ["docker", "push", tag],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(
            f"docker push failed (exit {result.returncode}):\n{result.stderr}"
        )

    logger.info("Docker image pushed successfully", tag=tag)
