"""Configuration loading and processing."""

import json
import os
import tempfile
from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Any, Optional, Union
from urllib.parse import urlparse

import requests
import torch
import yaml
from transformers.utils import is_torch_bf16_gpu_available

from axolotl.integrations.base import PluginManager
from axolotl.telemetry.errors import send_errors
from axolotl.telemetry.manager import TelemetryManager
from axolotl.utils.comet_ import setup_comet_env_vars
from axolotl.utils.config import (
    normalize_cfg_datasets,
    normalize_config,
    validate_config,
)
from axolotl.utils.dict import DictDefault
from axolotl.utils.logging import get_logger
from axolotl.utils.mlflow_ import setup_mlflow_env_vars
from axolotl.utils.tee import prepare_debug_log
from axolotl.utils.trackio_ import setup_trackio_env_vars
from axolotl.utils.trainer import prepare_optim_env
from axolotl.utils.wandb_ import setup_wandb_env_vars

LOG = get_logger(__name__)


def _coerce_value(value: Any, existing: Optional[Any] = None) -> Any:
    """Coerce a string CLI value to its most likely Python type.

    If an existing value is present in the config, its type is used to guide
    casting.  Otherwise, YAML-style inference is applied: booleans, ints,
    floats, and None literals are recognised automatically.

    Args:
        value: The raw value (typically a string from the CLI).
        existing: An optional existing config value whose type guides coercion.

    Returns:
        The value cast to the inferred or expected type.
    """
    if not isinstance(value, str):
        return value

    # If the config already has a typed value, cast to match
    if existing is not None:
        if isinstance(existing, bool):
            return value.lower() in ("true", "1", "yes")
        if isinstance(existing, int):
            try:
                return int(value)
            except (ValueError, TypeError):
                return value
        if isinstance(existing, float):
            try:
                return float(value)
            except (ValueError, TypeError):
                return value
        # For other types (str, list, dict, etc.), return as-is
        return value

    # No existing value -- use YAML-style inference
    lower = value.lower()
    if lower in ("true", "yes"):
        return True
    if lower in ("false", "no"):
        return False
    if lower in ("null", "none", "~"):
        return None

    # Try int then float
    try:
        return int(value)
    except ValueError:
        pass
    try:
        return float(value)
    except ValueError:
        pass

    return value


API_KEY_FIELDS = {"comet_api_key"}

TELEMETRY_MANAGER = TelemetryManager.get_instance()


def check_remote_config(config: Union[str, Path]) -> Union[str, Path]:
    """
    First, determines if the passed config is a valid HTTPS URL. Then, attempts to query
    for it and parse its content, first as JSON, then as YAML (YAML is preferred).
    Finally, the parsed content is written to a local file and its path is returned.

    Args:
        config: HTTPS URL to a YAML or JSON file.

    Returns:
        Either the original `config` if it's not a valid HTTPS URL, or the path to the
        downloaded remote config.

    Raises:
        ValueError: If the remote configuration is neither valid JSON or YAML.
        RuntimeError: If some request-related exception occurs from the file download.
        Exception: Catch-all for any other exception.
    """
    # Check if the config is a valid HTTPS URL to a .yml or .yaml file
    if not (isinstance(config, str) and config.startswith("https://")):
        return config  # Return the original value if it's not a valid URL

    filename = os.path.basename(urlparse(config).path)
    temp_dir = tempfile.mkdtemp()

    try:
        response = requests.get(config, timeout=30)
        response.raise_for_status()  # Check for HTTP errors

        content = response.content
        try:
            # Try parsing as JSON first to catch cases where JSON content is mistakenly
            # considered YAML.
            json.loads(content)

            # Log a warning but do not raise an error; JSON is technically valid YAML.
            # This can happen when you forget to point to a raw GitHub link.
            LOG.warning(
                f"Warning: The content of the file at {config} is JSON, which is technically valid YAML but might not be intended."
            )
        except json.JSONDecodeError:
            # If it's not valid JSON, verify it's valid YAML
            try:
                yaml.safe_load(content)
            except yaml.YAMLError as err:
                raise ValueError(
                    f"Failed to parse the content at {config} as YAML: {err}"
                ) from err

        # Write the content to a file if it's valid YAML (or JSON treated as YAML)
        output_path = Path(temp_dir) / filename
        with open(output_path, "wb") as file:
            file.write(content)
        LOG.info(
            f"Using the following config obtained from {config}: \n\n{content.decode('utf-8')}\n"
        )
        return output_path

    except requests.RequestException as err:
        # This catches all requests-related exceptions including HTTPError
        raise RuntimeError(f"Failed to download {config}: {err}") from err
    except Exception as err:
        # Catch-all for any other exceptions
        raise err


def choose_config(path: Path) -> str:
    """
    Helper method for choosing a `axolotl` config YAML file (considering only files
    ending with `.yml` or `.yaml`). If more than one config file exists in the passed
    `path`, the user is prompted to choose one.

    Args:
        path: Directory in which config file(s) are stored.

    Returns:
        Path to either (1) the sole YAML file, or (2) if more than one YAML files exist,
        the user-selected YAML file.

    Raises:
        ValueError: If no YAML files are found in the given `path`.
    """
    yaml_files = list(path.glob("*.yml")) + list(path.glob("*.yaml"))

    if not yaml_files:
        raise ValueError(
            "No YAML config files found in the specified directory. Are you using a .yml extension?"
        )

    if len(yaml_files) == 1:
        LOG.info(f"Using default YAML file '{yaml_files[0]}'")
        return str(yaml_files[0])

    LOG.info("Choose a YAML file:")
    for idx, file in enumerate(yaml_files):
        LOG.info(f"{idx + 1}. {file}")

    chosen_file = None
    while chosen_file is None:
        try:
            choice = int(input("Enter the number of your choice: "))
            if 1 <= choice <= len(yaml_files):
                chosen_file = str(yaml_files[choice - 1])
            else:
                LOG.info("Invalid choice. Please choose a number from the list.")
        except ValueError:
            LOG.info("Invalid input. Please enter a number.")

    return chosen_file


def prepare_plugins(cfg: DictDefault):
    """
    Registers the plugins for the given configuration.

    Args:
        cfg: Dictionary mapping `axolotl` config keys to values.
    """
    if cfg.get("plugins"):
        plugin_manager = PluginManager.get_instance()
        for plugin_name in cfg["plugins"]:
            plugin_manager.register(plugin_name)
        for plugin in plugin_manager.plugins.values():
            plugin.register(cfg)


def plugin_set_cfg(cfg: DictDefault):
    if cfg.get("plugins"):
        plugin_manager = PluginManager.get_instance()
        plugin_manager.cfg = cfg


@send_errors
def load_cfg(
    config: str | Path | DictDefault = Path("examples/"), **kwargs
) -> DictDefault:
    """
    Loads the `axolotl` configuration stored at `config`, validates it, and performs
    various setup.

    Args:
        config: Path (local or remote) to `axolotl` config YAML file.
        kwargs: Additional keyword arguments to override config file values.

    Returns:
        `DictDefault` mapping configuration keys to values.
    """
    if isinstance(config, (str, Path)):
        config = check_remote_config(config)
        if Path(config).is_dir():
            config = choose_config(Path(config))

        # Load the config from the yaml file
        with open(config, encoding="utf-8") as file:
            cfg: DictDefault = DictDefault(yaml.safe_load(file))

        cfg.axolotl_config_path = config
    else:
        cfg = config
        with NamedTemporaryFile(
            mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
        ) as temp_file:
            temp_file.write(yaml.dump(config.to_dict()))
            temp_file.close()
        cfg.axolotl_config_path = temp_file.name

    TELEMETRY_MANAGER.send_event(event_type="config-loaded", properties=cfg)

    # If there are any options passed in the cli, if it is something that seems valid
    # from the yaml, then overwrite the value
    cfg_keys = cfg.keys()

    # Separate nested (dot-notation) kwargs from flat kwargs
    nested_kwargs: dict[str, dict[str, Any]] = {}
    flat_kwargs: dict[str, Any] = {}
    for key, value in kwargs.items():
        if "__" in key:
            parent, child = key.split("__", 1)
            nested_kwargs.setdefault(parent, {})[child] = value
        else:
            flat_kwargs[key] = value

    # Apply flat kwargs
    for key, value in flat_kwargs.items():
        # If not strict, allow writing to cfg even if it's not in the yml already
        if key in cfg_keys or not cfg.strict:
            cfg[key] = _coerce_value(value, cfg.get(key))

    # Apply nested kwargs (e.g., trl__beta -> cfg.trl.beta)
    for parent, children in nested_kwargs.items():
        if parent not in cfg_keys and cfg.strict:
            continue
        if cfg[parent] is None:
            cfg[parent] = {}
        if not isinstance(cfg[parent], dict):
            LOG.warning(
                "Overwriting non-dict value for '%s' with nested CLI overrides", parent
            )
            cfg[parent] = {}
        for child_key, child_value in children.items():
            existing_child = cfg[parent].get(child_key)
            cfg[parent][child_key] = _coerce_value(child_value, existing_child)

    try:
        device_props = torch.cuda.get_device_properties("cuda")
        gpu_version = "sm_" + str(device_props.major) + str(device_props.minor)
    except:
        gpu_version = None

    prepare_plugins(cfg)

    cfg = validate_config(
        cfg,
        capabilities={
            "bf16": is_torch_bf16_gpu_available(),
            "fp8": compute_supports_fp8(),
            "n_gpu": int(os.environ.get("WORLD_SIZE", 1)),
            "compute_capability": gpu_version,
        },
        env_capabilities={
            "torch_version": str(torch.__version__).split("+", maxsplit=1)[0]
        },
    )

    # NOTE(djsaunde): We start outputting to output_dir/debug.log at this point since we
    # have to wait for cfg.output to be resolved. We could call this earlier if we write
    # to a temporary file, and then move it later.
    prepare_debug_log(cfg)
    prepare_optim_env(cfg)
    normalize_config(cfg)
    normalize_cfg_datasets(cfg)
    setup_wandb_env_vars(cfg)
    setup_mlflow_env_vars(cfg)
    setup_comet_env_vars(cfg)
    setup_trackio_env_vars(cfg)
    plugin_set_cfg(cfg)

    TELEMETRY_MANAGER.send_event(event_type="config-processed", properties=cfg)
    cfg_to_log = {
        k: "[REDACTED]" if k in API_KEY_FIELDS else v
        for k, v in cfg.items()
        if v is not None
    }
    LOG.info(
        "config:\n%s",
        json.dumps(cfg_to_log, indent=2, default=str, sort_keys=True),
    )

    return cfg


def compute_supports_fp8() -> bool:
    try:
        compute_capability = torch.cuda.get_device_capability()
        return compute_capability >= (9, 0)
    except RuntimeError:
        return False
