arguments.py

"""
arguments.py - Command-line argument parsing and YAML configuration loading.

Provides the CLI interface for FLocKit. Parses command-line arguments
(config file path, data path overrides, port), loads the YAML configuration
via OmegaConf, validates required fields, and applies any CLI overrides on
top of the YAML values. Environment variables (``FLOCKIT_*``) act as
fallback defaults so the same entry point works in interactive,
container, and CI contexts.

Public API:
    add_args        - Parse argparse namespace plus OmegaConf-style overrides.
    load_arguments  - Parse CLI, load YAML, validate, and return DictConfig.

Environment:
    FLOCKIT_CONF                - Default path to the YAML config file.
    FLOCKIT_PORT                - Default port for the SDK server.
    FLOCKIT_DATA_PATH           - Default override for ``data_args.data_path``.
    FLOCKIT_DATA_SOURCE         - Default override for ``data_args.data_source``.
    FLOCKIT_DATA_INDICES_PATH   - Default override for indices path.
    FLOCKIT_OVERRIDES           - Extra OmegaConf ``key=value`` overrides.

Usage:
    python main.py --conf configs/my_config.yaml --port 5000
"""

from __future__ import annotations

import argparse
import os
from typing import List, Tuple

from omegaconf import DictConfig, OmegaConf

from flockit_logger import get_logger
from constants import (
    HOMEPAGE,
    PRODUCT_DIST_NAME,
    PRODUCT_NAME,
    PRODUCT_TAGLINE,
    VENDOR,
    get_product_version,
)


# -----------------------------------------------------------------------------
# Environment helpers
# -----------------------------------------------------------------------------

def _parse_env_str(var_name: str) -> str | None:
    """Read an optional string from environment variables and normalize empty to None."""
    raw = os.getenv(var_name)
    if raw is None:
        return None
    value = raw.strip()
    return value if value else None


def _parse_env_int(var_name: str) -> int | None:
    """Read an optional integer from environment variables."""
    raw = os.getenv(var_name)
    if raw is None or raw == "":
        return None
    try:
        return int(raw)
    except ValueError as exc:
        raise ValueError(
            f"Environment variable {var_name} must be an integer, got: {raw}"
        ) from exc


def _parse_env_overrides() -> List[str]:
    """Read optional OmegaConf CLI overrides from FLOCKIT_OVERRIDES."""
    raw = os.getenv("FLOCKIT_OVERRIDES", "").strip()
    if not raw:
        return []

    tokens: List[str] = []
    for token in raw.replace("\n", ",").split(","):
        token = token.strip()
        if not token:
            continue
        if "=" not in token:
            get_logger(__name__).warning(
                "Ignoring malformed FLOCKIT_OVERRIDES token (expected key=value): {}", token
            )
            continue
        tokens.append(token)

    return tokens


# -----------------------------------------------------------------------------
# Public API
# -----------------------------------------------------------------------------

def add_args() -> Tuple[argparse.Namespace, List[str]]:
    """
    Parse command-line arguments for FLocKit.

    Returns:
        A tuple of:
        - Parsed argparse namespace.
        - OmegaConf-style CLI overrides (``key=value`` tokens).
    """
    default_conf = _parse_env_str("FLOCKIT_CONF")
    default_port = _parse_env_int("FLOCKIT_PORT")
    default_data_path = _parse_env_str("FLOCKIT_DATA_PATH")
    default_data_source = _parse_env_str("FLOCKIT_DATA_SOURCE")
    default_indices_path = _parse_env_str("FLOCKIT_DATA_INDICES_PATH")

    parser = argparse.ArgumentParser(
        prog=PRODUCT_DIST_NAME,
        description=(
            f"{PRODUCT_NAME} ({VENDOR}) — production-grade federated "
            "learning toolkit for the FLock protocol. Loads a template "
            "config, builds the task model, and starts the FLock SDK "
            "HTTP server."
        ),
        epilog=(
            f"{PRODUCT_TAGLINE}  ·  {HOMEPAGE}  ·  "
            "see templates/<name>/configs/*.yaml for ready-to-run examples."
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--version", "-V",
        action="version",
        version=f"%(prog)s {get_product_version()}  ({VENDOR})",
    )
    parser.add_argument(
        "--yaml_config_file",
        "--conf",
        "--config",
        help="Path to template configuration file (.yaml). Can also be set via FLOCKIT_CONF.",
        type=str,
        required=default_conf is None,
        default=default_conf,
        metavar="PATH",
        dest="yaml_config_file",
    )
    parser.add_argument(
        "--data_path",
        help=(
            "Path to the dataset file. Overrides data_args.data_path in YAML. "
            "Can also be set via FLOCKIT_DATA_PATH."
        ),
        type=str,
        required=False,
        default=default_data_path,
        metavar="PATH",
    )
    parser.add_argument(
        "--data_source",
        help=(
            "Dataset source override (e.g. 'local' or 'huggingface'). "
            "Can also be set via FLOCKIT_DATA_SOURCE."
        ),
        type=str,
        required=False,
        default=default_data_source,
        metavar="SOURCE",
    )
    parser.add_argument(
        "--data_indices_path",
        help=(
            "Path to a JSON/JSONL file with dataset row indices to select. "
            "Can also be set via FLOCKIT_DATA_INDICES_PATH."
        ),
        type=str,
        required=False,
        default=default_indices_path,
        metavar="PATH",
    )
    parser.add_argument(
        "--port",
        help="Port number for the FLock SDK Flask server. Can also be set via FLOCKIT_PORT.",
        type=int,
        default=default_port,
        required=False,
        metavar="PORT",
    )

    args, unknown = parser.parse_known_args()
    cli_overrides: List[str] = []
    ignored_unknown: List[str] = []
    for token in unknown:
        if "=" in token and not token.startswith("-"):
            cli_overrides.append(token)
        else:
            ignored_unknown.append(token)

    # Environment overrides are appended first and CLI overrides can still win.
    cli_overrides = _parse_env_overrides() + cli_overrides

    if ignored_unknown:
        raise ValueError(
            "Unknown CLI arguments: "
            + ", ".join(ignored_unknown)
            + ". For config overrides, use key=value syntax."
        )
    return args, cli_overrides


def load_arguments() -> DictConfig:
    """
    Load, validate, and return the merged configuration.

    Steps:
        1. Parse CLI arguments.
        2. Load the YAML configuration file via OmegaConf.
        3. Validate required fields (``common_args.template_name``).
        4. Apply CLI overrides (``data_path``, ``data_source``, ``port``, etc.).

    Returns:
        The final OmegaConf DictConfig ready for use by the application.

    Raises:
        FileNotFoundError: If the configuration file does not exist.
        ValueError: If the file is invalid or missing required fields.
    """
    cmd_args, cli_overrides = add_args()

    config_file = cmd_args.yaml_config_file
    if not config_file:
        raise ValueError(
            "Configuration file must be provided via --conf/--config or FLOCKIT_CONF"
        )
    if not os.path.exists(config_file):
        raise FileNotFoundError(f"Configuration file not found: {config_file}")
    if not os.path.isfile(config_file):
        raise FileNotFoundError(f"Path exists but is not a file: {config_file}")

    try:
        args = OmegaConf.load(config_file)
    except Exception as exc:
        raise ValueError(f"Failed to parse configuration file '{config_file}': {exc}") from exc

    _apply_command_line_overrides(args, cmd_args)
    if cli_overrides:
        args = OmegaConf.merge(args, OmegaConf.from_cli(cli_overrides))

    _validate_required_fields(args)
    return args


# -----------------------------------------------------------------------------
# Private helpers
# -----------------------------------------------------------------------------

def _validate_required_fields(args: DictConfig) -> None:
    """
    Ensure the config contains the minimum required sections and keys.

    Args:
        args: The loaded OmegaConf configuration.

    Raises:
        ValueError: If ``common_args`` or ``common_args.template_name`` is missing.
    """
    if not hasattr(args, "common_args"):
        raise ValueError("Configuration file must include 'common_args' section")
    if not hasattr(args.common_args, "template_name"):
        raise ValueError("Configuration file must include 'common_args.template_name'")
    if not args.common_args.template_name:
        raise ValueError("template_name cannot be empty")
    if hasattr(args.common_args, "run_mode"):
        run_mode = str(args.common_args.run_mode).strip().lower()
        if run_mode not in {"simulation", "production"}:
            raise ValueError(
                "common_args.run_mode must be 'simulation' or 'production', "
                f"got: {args.common_args.run_mode!r}"
            )
    if hasattr(args.common_args, "port") and args.common_args.port is not None:
        try:
            port = int(args.common_args.port)
        except (TypeError, ValueError) as exc:
            raise ValueError(
                f"common_args.port must be an integer, got: {args.common_args.port!r}"
            ) from exc
        if not (1 <= port <= 65535):
            raise ValueError(f"common_args.port must be in [1, 65535], got: {port}")
        args.common_args.port = port


def _apply_command_line_overrides(args: DictConfig, cmd_args: argparse.Namespace) -> None:
    """
    Overlay CLI-provided values on top of the YAML configuration.

    When FL-Alliance-Client passes a directory (e.g. staging dir), the CLI override
    would overwrite the YAML's data_path (e.g. "adults5.csv"). We resolve this by
    looking for the YAML's file inside the directory so the model gets the correct path.

    Args:
        args: The mutable OmegaConf configuration.
        cmd_args: Parsed CLI arguments.
    """
    if cmd_args.data_path:
        if not hasattr(args, "data_args"):
            args.data_args = OmegaConf.create({})
        override_path = cmd_args.data_path
        # When override is a directory, append YAML's data_path to it (e.g. dir + "adults5.csv").
        if os.path.isdir(override_path) and hasattr(args.data_args, "data_path"):
            yaml_path = args.data_args.data_path
            if yaml_path:
                override_path = os.path.normpath(os.path.join(override_path, str(yaml_path)))
        args.data_args.data_path = override_path

    if cmd_args.data_source:
        if not hasattr(args, "data_args"):
            args.data_args = OmegaConf.create({})
        args.data_args.data_source = cmd_args.data_source

    if cmd_args.data_indices_path:
        if not hasattr(args, "data_args"):
            args.data_args = OmegaConf.create({})
        args.data_args.data_indices_path = cmd_args.data_indices_path

    if cmd_args.port is not None:
        if not hasattr(args, "common_args"):
            args.common_args = OmegaConf.create({})
        args.common_args.port = cmd_args.port