rimport

#!/glade/u/apps/derecho/24.12/opt/view/bin/python
# TODO: Move all the Python into new file rimport.py for simpler testing. Keep rimport as a
# convenience wrapper.
"""
Copy files from CESM inputdata directory to a publishing directory, then replace the original with a
symlink to the copy.

Do `rimport --help` for more information.
"""
from __future__ import annotations

import argparse
import os
import pwd
import shutil
import sys
from pathlib import Path
from typing import Iterable, List
from urllib.request import Request, urlopen
from urllib.error import HTTPError

from relink import replace_one_file_with_symlink

import shared
INDENT = shared.INDENT

DEFAULT_INPUTDATA_ROOT = Path(shared.DEFAULT_INPUTDATA_ROOT)
DEFAULT_STAGING_ROOT = Path(shared.DEFAULT_STAGING_ROOT)
STAGE_OWNER = "cesmdata"
INPUTDATA_URL = "https://osdf-data.gdex.ucar.edu/ncar/gdex/d651077/cesmdata/inputdata"

# Configure logging
logger = shared.logger


def build_parser() -> argparse.ArgumentParser:
    """Build and configure the argument parser for rimport.

    Returns:
        argparse.ArgumentParser: Configured parser ready to parse command-line arguments.
    """
    parser = argparse.ArgumentParser(
        description=(
            f"Copy files from CESM inputdata directory ({DEFAULT_INPUTDATA_ROOT}) to a publishing"
            " directory, then replace the original with a symlink to the copy."
        ),
        add_help=False,  # Disable automatic help to add custom -help flag
    )

    parser.add_argument(
        "--file",
        "-file",
        dest="file",
        metavar="filename",
        help="Provide a file to import. Must be in the CESM inputdata directory.",
    )

    parser.add_argument(
        "--list",
        "-list",
        dest="filelist",
        metavar="filelist",
        help=(
            "Provide a file that contains a list of filenames to import. All filenames in the list"
            " must be in the CESM inputdata directory."
        ),
    )

    parser.add_argument(
        "items_to_process",
        nargs="*",
        help="One or more files to process. (Optional; can use --file instead to process just one.)"
    )

    # Add inputdata_root option flags
    shared.add_inputdata_root(parser)

    parser.add_argument(
        "--check",
        "-check",
        "-c",
        action="store_true",
        help="Check whether file(s) is/are already published.",
    )

    # Add verbosity options
    shared.add_parser_verbosity_group(parser)

    # Add help text
    shared.add_help(parser)

    return parser


def read_filelist(list_path: Path) -> List[str]:
    """Read a file list and return non-empty, non-comment lines.

    Reads a text file containing a list of filenames, filtering out:
    - Blank lines (empty or whitespace-only)
    - Comment lines (starting with '#')

    Args:
        list_path: Path to the file containing the list of filenames.

    Returns:
        List of stripped, non-empty lines that are not comments.
    """
    lines: List[str] = []
    with list_path.open("r", encoding="utf-8") as f:
        for raw in f:
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            lines.append(line)
    return lines


def normalize_paths(root: Path, relnames: Iterable[str]) -> List[Path]:
    """Convert relative or absolute path names to normalized absolute Paths.

    For each name in relnames:
    - If the name is relative, it is assumed to be relative to `root` and made absolute
    All paths are then normalized to their absolute form, replacing . and .. as needed.

    Note that symlinks are NOT resolved.

    Args:
        root: Base directory under which relative paths are assumed to be.
        relnames: Iterable of path names (relative or absolute) to normalize.

    Returns:
        List of normalized absolute Path objects.
    """
    paths: List[Path] = []
    for name in relnames:
        p = root / name if not Path(name).is_absolute() else Path(name)
        p = Path(os.path.normpath(p.absolute()))
        paths.append(p)
    return paths


def check_relink_worked(src: Path, dst: Path) -> None:
    """Check whether relink worked

    Args:
        src (Path): Source file (should have been converted to symlink)
        dst (Path): Destination file (symlink target)

    Raises:
        RuntimeError: If src is not a symlink pointing to dst.
    """
    if not (src.is_symlink() and src.resolve() == dst):
        raise RuntimeError("Error relinking during rimport")


def stage_data(
    src: Path, inputdata_root: Path, staging_root: Path, check: bool = False
) -> None:
    """Stage a file by mirroring its path under `staging_root`, then replace with symlink to staged.

    Destination path is computed by replacing the `inputdata_root` prefix of `src`
    with `staging_root`, i.e.:
        dst = staging_root / src.relative_to(inputdata_root)

    Args:
        src: Source file path to stage.
        inputdata_root: Root directory of the inputdata tree.
        staging_root: Root directory where files will be staged.
        check: If True, just check whether the file is already published.

    Raises:
        RuntimeError: If `src` is a live symlink pointing outside staging, or if `src` is outside
                      the inputdata root, or if `src` is already under staging directory.
        RuntimeError: If `src` is a broken symlink.
        RuntimeError: If it failed to replace `src` with a symlink to the staged file.
        FileNotFoundError: If `src` does not exist.

    Guardrails:
        * Raise if `src` is a *live* symlink to a file outside staging root ("outside staging").
        * Raise if `src` is a broken symlink or is outside the inputdata root.
    """
    if src.is_symlink():
        if not os.path.exists(src.resolve()):
            raise RuntimeError(f"Source is a broken symlink: {src}")
        if not src.resolve().is_relative_to(staging_root.resolve()):
            raise RuntimeError(
                f"Source is a symlink, but target ({src.resolve()}) is outside staging directory "
                f"({staging_root})"
            )
        logger.info("%sFile is already published and linked.", INDENT)
        print_can_file_be_downloaded(
            can_file_be_downloaded(src.resolve(), staging_root)
        )
        return

    if not src.exists():
        raise FileNotFoundError(f"source not found: {src}")

    try:
        rel = src.resolve().relative_to(inputdata_root.resolve())
    except ValueError as exc:
        if src.resolve().is_relative_to(staging_root.resolve()):
            raise RuntimeError(
                f"Source file '{src.name}' is already under staging directory '{staging_root}'."
            ) from exc
        raise RuntimeError(
            f"source not under inputdata root: {src} not in {inputdata_root}"
        ) from exc

    dst = staging_root / rel

    if dst.exists():
        msg = "File is already published but NOT linked"
        if check:
            logger.info("%s; would link.", msg)
            print_can_file_be_downloaded(can_file_be_downloaded(rel, staging_root))
        else:
            logger.info("%s; linking now.", msg)
            replace_one_file_with_symlink(inputdata_root, staging_root, str(src))
            check_relink_worked(src, dst)
            print_can_file_be_downloaded(can_file_be_downloaded(rel, staging_root))
        return

    if check:
        logger.info("%sFile is not already published", INDENT)
        return

    # Copy file to destination
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dst)
    logger.info("%s[rimport] staged %s -> %s", INDENT, src, dst)

    # Replace original with symlink to destination
    replace_one_file_with_symlink(inputdata_root, staging_root, str(src))
    check_relink_worked(src, dst)


def ensure_running_as(target_user: str, argv: list[str]) -> None:
    """Ensure the script is running as the target user, re-executing via sudo if needed.

    If not running as `target_user`, re-exec via sudo -u target_user (handles 2FA via PAM).
    This function will not return if re-execution is needed; it replaces the current process.

    Args:
        target_user: Username to run as (e.g., 'cesmdata').
        argv: Command-line arguments to pass to the re-executed process.

    Raises:
        SystemExit: If the target user is not found on the system (exit code 2).
        SystemExit: If not running interactively and authentication is required (exit code 2).

    Note:
        If re-execution is needed, this function calls os.execvp() and does not return.
    """
    try:
        target_uid = pwd.getpwnam(target_user).pw_uid
    except KeyError as exc:
        logger.error("rimport: target user '%s' not found on this system", target_user)
        raise SystemExit(2) from exc

    if os.geteuid() != target_uid:
        try:
            assert sys.stdin.isatty()
        except AssertionError as exc:
            logger.error(
                "rimport: need interactive TTY to authenticate as '%s' (2FA).\n"
                "         Try: sudo -u %s rimport …",
                target_user,
                target_user,
            )
            raise SystemExit(2) from exc
        # Re-exec under target user; this invokes sudo’s normal password/2FA flow.
        os.execvp("sudo", ["sudo", "-u", target_user, "--"] + argv)


def get_staging_root() -> Path:
    """Return the staging root directory path.

    Uses $RIMPORT_STAGING if set, otherwise returns the default staging root.

    Returns:
        Path: Resolved absolute path to the staging root directory.
    """
    env = os.getenv("RIMPORT_STAGING")
    if env:
        return Path(env).expanduser().resolve()
    return DEFAULT_STAGING_ROOT


def can_file_be_downloaded(file_relpath: Path, staging_root: Path, timeout: float = 10):
    """Check whether a file is available for download from the CESM inputdata server.

    Sends a HEAD request to the CESM inputdata URL to verify if the file exists and is
    accessible without downloading the entire file.

    Args:
        file_relpath: Relative path to the file (relative to staging_root), or an absolute
            path that will be made relative to staging_root.
        staging_root: Root directory of the staging area, used to compute relative path
            if file_relpath is absolute.
        timeout: Maximum time in seconds to wait for the server response. Default is 10.

    Returns:
        bool: True if the file is accessible (HTTP status 2xx or 3xx), False otherwise
            (including 404, network errors, timeouts, etc.).
    """
    # Get URL
    if file_relpath.is_absolute():
        file_relpath = file_relpath.relative_to(staging_root)
    url = os.path.join(INPUTDATA_URL, file_relpath)

    # Check whether URL can be accessed
    req = Request(url, method="HEAD")
    try:
        with urlopen(req, timeout=timeout) as resp:
            return 200 <= resp.status < 400
    except HTTPError:
        # Server reached, but resource doesn't exist (404, 410, etc.)
        return False


def print_can_file_be_downloaded(file_can_be_downloaded: bool):
    """Print a message indicating whether a file is available for download.

    Args:
        file_can_be_downloaded: Boolean indicating if the file can be downloaded.
    """
    if file_can_be_downloaded:
        logger.info("%sFile is available for download.", INDENT)
    else:
        logger.info("%sFile is not (yet) available for download.", INDENT)


def get_files_to_process(file: str, filelist: str, items_to_process: list):
    """Get list of files to process.

    Uses --file and/or --filelist arguments, as well as positional items_to_process if given.

    Args:
        file (str): Single file to process.
        filelist (str): File containing list of files to process.
        items_to_process (list): List of files to process.

    Returns:
        list: List of files to process
        int:  Result code
    """
    if file is not None:
        files_to_process = [file]
    else:
        files_to_process = []

    if filelist is not None:
        list_path = Path(filelist).expanduser().resolve()
        if not list_path.exists():
            logger.error("rimport: list file not found: %s", list_path)
            return None, 2
        files_in_list = read_filelist(list_path)
        if not files_in_list:
            logger.error("rimport: no filenames found in list: %s", list_path)
            return None, 2
        files_to_process.extend(files_in_list)

    if items_to_process:
        files_to_process.extend(items_to_process)

    if not files_to_process:
        logger.error("rimport: At least one of --file or --filelist is required")
        return None, 2

    return files_to_process, 0


def main(argv: List[str] | None = None) -> int:
    """Main entry point for the rimport tool.

    Copies and relinks files from the CESM inputdata directory to a staging/publishing directory,
    preserving the directory structure. Ensures the script runs as the correct user
    (STAGE_OWNER) and handles both single files and file lists.

    Args:
        argv: Command-line arguments to parse. If None, uses sys.argv.

    Returns:
        int: Exit code (0 for success, 1 if any files had errors, 2 for fatal errors).

    Environment Variables:
        RIMPORT_SKIP_USER_CHECK: Set to "1" to skip automatic user switching.
        RIMPORT_STAGING: Override the default staging root directory.

    Exit Codes:
        0: All files staged successfully.
        1: One or more files failed to stage or relink (errors printed to stderr).
        2: Fatal error (missing inputdata directory, missing file list, etc.).
    """
    parser = build_parser()
    args = parser.parse_args(argv)

    # Configure logging based on verbosity flags
    log_level = shared.get_log_level(quiet=args.quiet, verbose=args.verbose)
    shared.configure_logging(log_level)

    # Ensure we are running as the STAGE_OWNER account before touching the tree
    # Set env var RIMPORT_SKIP_USER_CHECK=1 if you prefer to run `sudox -u STAGE_OWNER rimport …`
    # explicitly (or for testing).
    if not args.check and os.getenv("RIMPORT_SKIP_USER_CHECK") != "1":
        ensure_running_as(STAGE_OWNER, sys.argv)

    root = Path(args.inputdata_root).expanduser().resolve()
    if not root.exists():
        logger.error("rimport: inputdata directory does not exist: %s", root)
        return 2

    # Determine the list of relative filenames to handle
    files_to_process, status = get_files_to_process(args.file, args.filelist, args.items_to_process)
    if status:
        return status

    # Resolve to full paths (keep accepting absolute names too)
    paths = normalize_paths(root, files_to_process)
    staging_root = get_staging_root()
    # Execute the new action per file
    errors = 0
    for p in paths:
        logger.info("'%s':", p)
        try:
            stage_data(p, root, staging_root, args.check)
        except Exception as e:  # pylint: disable=broad-exception-caught
            # General Exception keeps CLI robust for batch runs
            errors += 1
            logger.error("%srimport: error processing %s: %s", INDENT, p, e)

    if errors:
        return 1
    if not args.check:
        logger.info("\nNo need to run relink.py")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())