"""Runner module for GEB functions."""

import cProfile
import importlib
import logging
import os
import platform
import pstats
import shutil
import subprocess
import sys
import tempfile
import zipfile
from collections.abc import Callable
from datetime import datetime
from operator import attrgetter
from pathlib import Path
from typing import Any, TextIO, TypeVar, cast

import geopandas as gpd
import memray
import yaml
from pydantic import BaseModel, ValidationError
from shapely.geometry import box

from geb import GEB_PACKAGE_DIR
from geb.build import GEBModel as GEBModelBuild
from geb.build.__init__ import (
    cluster_subbasins_following_coastline,
    create_cluster_visualization_map,
    create_multi_basin_configs,
    get_all_downstream_subbasins_in_geom,
    get_river_graph,
    save_clusters_as_merged_geometries,
    save_clusters_to_geoparquet,
)
from geb.build.data_catalog import NewDataCatalog
from geb.build.methods import build_method
from geb.config_schema import Config
from geb.model import GEBModel
from geb.workflows.io import WorkingDirectory, read_params
from geb.workflows.methods import multi_level_merge

OPTIMIZE_DEFAULT: bool = False
TIMING_DEFAULT: bool = False
WORKING_DIRECTORY_DEFAULT: Path = Path(".")
CONFIG_DEFAULT: Path = Path("model.yml")
UPDATE_DEFAULT: Path = Path("update.yml")
BUILD_DEFAULT: Path = Path("build.yml")
PROFILE_SPEED_DEFAULT: bool = False
PROFILE_RAM_DEFAULT: bool = False

DATA_CATALOG_DEFAULT: Path = GEB_PACKAGE_DIR / "data_catalog.yml"
DATA_PROVIDER_DEFAULT: str = os.environ.get("GEB_DATA_PROVIDER", "default")
DATA_ROOT_DEFAULT: Path = Path(
    os.environ.get(
        "GEB_DATA_ROOT",
        GEB_PACKAGE_DIR.parent.parent / "data_catalog",
    )
)
ALTER_FROM_MODEL_DEFAULT: Path = Path("../base")

ResultType = TypeVar("ResultType")


class DetectDuplicateKeysYamlLoader(yaml.SafeLoader):
    """Custom YAML loader that detects duplicate keys in mappings.

    Raises:
        ValueError: If a duplicate key is found in the YAML mapping.
    """

    def construct_mapping(
        self, node: yaml.nodes.MappingNode, deep: bool = False
    ) -> dict:
        """Construct a mapping from a YAML node, checking for duplicate keys.

        Args:
            node: The YAML node to construct the mapping from.
            deep: Whether to perform a deep construction of the mapping. Defaults to False.

        Raises:
            ValueError: If a duplicate key is found in the YAML mapping.

        Returns:
            dict: The constructed mapping.
        """
        mapping = {}
        for key_node, value_node in node.value:
            key = self.construct_object(key_node, deep=deep)
            if key in mapping:
                raise ValueError(f"Duplicate key found: {key}")
            mapping[key] = self.construct_object(value_node, deep=deep)
        return mapping


def parse_config(
    config_path: dict | Path | str,
    current_directory: Path | None = None,
    schema: type[BaseModel] | None = None,
) -> dict[str, Any]:
    """Parse config.

    This method recursively parses the config file and resolves any 'inherits' keys.

    Args:
        config_path: Path to the config file or a dict with the config.
        current_directory: Current directory to resolve relative paths.
            If None, the current working directory is used.
        schema: Pydantic schema to validate the config against.

    Returns:
        Full model configuation of the model without any remaining 'inherits' keys.
    """
    if current_directory is None:
        current_directory = Path.cwd()

    if isinstance(config_path, dict):
        config = config_path
    else:
        config: dict | None = yaml.load(
            open(current_directory / config_path, "r"),
            Loader=DetectDuplicateKeysYamlLoader,
        )
        if config is None:
            config = {}
        current_directory = current_directory / Path(config_path).parent

    if "inherits" in config:
        inherit_config_path = config["inherits"]
        inherit_config_path = inherit_config_path.format(**os.environ)
        # replace {VAR} with environment variable VAR if it exists
        inherit_config_path = os.environ.get(
            inherit_config_path, os.path.expandvars(inherit_config_path)
        )
        # if inherits is not an absolute path, we assume it is relative to the config file
        if not Path(inherit_config_path).is_absolute():
            inherit_config_path = current_directory / config["inherits"]
        inherited_config = yaml.load(
            open(inherit_config_path, "r"),
            Loader=yaml.FullLoader,
        )
        current_directory = current_directory / Path(inherit_config_path).parent
        del config[
            "inherits"
        ]  # remove inherits key from config to avoid infinite recursion
        config = multi_level_merge(inherited_config, config)
        config = parse_config(
            config, current_directory=current_directory, schema=schema
        )

    # Validate config
    if schema is not None:
        try:
            schema(**config)
        except ValidationError as e:
            # We warn instead of raising an error to allow for extra fields or partial configs
            # during development, but ideally this should be strict.
            logging.warning(f"Configuration validation failed: {e}")

    return config


def create_logger(name: str) -> logging.Logger:
    """Create logger with console and file handler.

    Args:
        name: Name of the logger.
    Returns:
        Logger instance.
    """
    logger: logging.Logger = logging.getLogger(name)

    if logger.handlers:
        return logger

    # set log level to debug
    logger.setLevel(logging.DEBUG)
    # create console handler and set level to debug
    ch: logging.StreamHandler[TextIO] = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)
    # create formatter
    formatter = logging.Formatter(
        "%(asctime)s - %(levelname)s - %(message)s", datefmt="%d-%m %H:%M:%S"
    )
    # add formatter to ch
    ch.setFormatter(formatter)
    # add ch to logger
    logger.addHandler(ch)

    # prevent double logging
    logger.propagate = False

    # add file handler
    folder = Path("logs")
    folder.mkdir(exist_ok=True, parents=True)
    fp = folder / f"{name}.log"
    if fp.exists():
        fp.unlink()  # remove existing log file if it exists
    fh = logging.FileHandler(fp)
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(formatter)
    logger.addHandler(fh)
    return logger


def run_model_with_method(
    method: str | None,
    config: dict | str | Path = CONFIG_DEFAULT,
    working_directory: Path = WORKING_DIRECTORY_DEFAULT,
    timing: bool = TIMING_DEFAULT,
    profile_speed: bool = PROFILE_SPEED_DEFAULT,
    profile_ram: bool = PROFILE_RAM_DEFAULT,
    optimize: bool = OPTIMIZE_DEFAULT,
    method_args: dict = {},
    close_after_run: bool = True,
) -> Any:
    """Run model with a specific method.

    Args:
        method: Method to run on the model. If None, the model is created but no method is run.
        config: Path to the model configuration file or a dict with the config.
        working_directory: Working directory for the model.
        timing: If True, run the model with timing, printing the time taken for specific methods.
        profile_speed: If True, run the model with speed profiling.
        profile_ram: If True, run the model with RAM profiling.
        optimize: If True, run the model in optimized mode, skipping asserts and water balance checks.
        method_args: Optional arguments to pass to the method.
        close_after_run: If True, close the model after running the method. Defaults to True.

    Returns:
        The result of the method run or the GEBModel instance if no method was run.

    Raises:
        SystemExit: If the model is restarted in optimized mode.
    """
    # check if we need to run the model in optimized mode
    # if the model is already running in optimized mode, we don't need to restart it
    # or else we start an infinite loop
    if optimize and sys.flags.optimize == 0:
        # If the script is not a .py file, we need to add the .exe extension
        if platform.system() == "Windows" and not sys.argv[0].endswith(".py"):
            sys.argv[0] = sys.argv[0] + ".exe"
        command: list[str] = [sys.executable, "-O"] + sys.argv
        raise SystemExit(subprocess.run(command).returncode)

    def run_operation(
        logger: logging.Logger, close_after_run: bool = close_after_run
    ) -> Any:
        from geb.config_schema import Config

        config_parsed: dict[str, Any] = parse_config(config, schema=Config)
        files: dict[str, Any] = parse_config(
            read_params(Path("input/files.yml"))
            if "files" not in config_parsed["general"]
            else config_parsed["general"]["files"]
        )

        geb = GEBModel(
            config=config_parsed,
            files=files,
            timing=timing,
            logger=logger,
        )

        result = geb
        if method is not None:
            result = getattr(geb, method)(**method_args)

        # If we are profiling, we don't close the model here, but return it
        # so it can be profiled while open. The profiler wrapper will handle closing.
        if (profile_speed or profile_ram) and close_after_run:
            return result, geb

        if close_after_run:
            geb.close()

        return result or geb

    with WorkingDirectory(working_directory):
        return _run_with_optional_profiling(
            profile_speed,
            profile_ram,
            run_operation,
            name=method if method is not None else "wo_method",
            logger=create_logger(name=method if method is not None else "wo_method"),
        )


def _dump_speed_profile(profile: cProfile.Profile, name: str, date: str) -> None:
    """
    Persist speed profiling statistics to disk in both binary and human-readable formats.

    Args:
        profile: Profile instance containing execution statistics.
        name: Identifier used for naming the output files.
        date: Date string for unique file naming.
    """
    profiling_folder = Path("profiling")
    profiling_folder.mkdir(exist_ok=True)

    # Save binary data for visualization in tools like SnakeViz or RunSnakeRun
    binary_path = profiling_folder / f"{name}_speed_{date}.prof"
    profile.dump_stats(binary_path)

    # Save human-readable summary
    text_path = profiling_folder / f"{name}_speed_summary_{date}.txt"
    with open(text_path, "w", encoding="utf-8") as stream:
        stats = pstats.Stats(profile, stream=stream)
        stats.strip_dirs().sort_stats("cumtime").print_stats()

    print(f"Speed profile saved to: {binary_path}")
    print(f"Speed profile summary saved to: {text_path}")


def _run_with_optional_profiling(
    profile_speed: bool,
    profile_ram: bool,
    operation: Callable[..., ResultType],
    name: str,
    logger: logging.Logger,
) -> ResultType:
    """Run an operation with optional speed and RAM profiling.

    Args:
        profile_speed: Whether speed profiling should be enabled.
        profile_ram: Whether RAM profiling should be enabled.
        operation: Callable containing the operation to execute.
        name: Name of the operation being executed, used for profiling output.
        logger: Logger instance to pass to the operation.

    Returns:
        Return value from the operation.
    """
    if not profile_speed and not profile_ram:
        return operation(logger=logger)

    date: str = datetime.now().strftime("%Y%m%d_%H%M%S")

    profiling_folder = Path("profiling")
    profiling_folder.mkdir(exist_ok=True)

    # Initialize profiles
    speed_profiler = None
    if profile_speed:
        speed_profiler = cProfile.Profile()
        speed_profiler.enable()

    ram_profiler_tracker = None
    if profile_ram:
        ram_path: Path = profiling_folder / f"{name}_ram_{date}.bin"
        if ram_path.exists():
            ram_path.unlink()  # Remove existing RAM profile if it exists
        ram_profiler_tracker = memray.Tracker(ram_path)
        ram_profiler_tracker.__enter__()

    try:
        if profile_ram:
            # We wrap the operation to allow capturing objects just before it finishes
            # but while the GEBModel instance is still alive.
            pass

        # If we are profiling, we tell the operation not to close the model
        # so we can profile it while it is still open.
        if profile_speed or profile_ram:
            run_output: Any = operation(close_after_run=False, logger=logger)
        else:
            run_output: Any = operation(logger=logger)

        result: ResultType
        geb_to_close: Any = None

        if (profile_speed or profile_ram) and isinstance(run_output, tuple):
            result = cast(ResultType, run_output[0])
            geb_to_close = run_output[1]
        else:
            result = run_output

        # Capture objects while they are still in scope (and model is not closed if profiling)
        if profile_ram:
            _dump_ram_object_profile(name, date, keep_alive=geb_to_close)

        # Close model if we manually deferred it for profiling
        if geb_to_close is not None:
            # Use getattr to avoid type check complaining about unknown method
            getattr(geb_to_close, "close")()
    finally:
        if speed_profiler:
            speed_profiler.disable()
            _dump_speed_profile(speed_profiler, name, date)
        if ram_profiler_tracker:
            ram_profiler_tracker.__exit__(None, None, None)

            ram_bin: Path = profiling_folder / f"{name}_ram_{date}.bin"
            ram_html: Path = profiling_folder / f"{name}_ram_{date}.html"
            reporter_cmd: list[str | Path] = [
                sys.executable,
                "-m",
                "memray",
                "flamegraph",
                "--temporal",
                "-f",
                "-o",
                ram_html,
                ram_bin,
            ]
            subprocess.run(reporter_cmd, check=True)
            print(f"RAM flamegraph saved to: {ram_html}")

    return result


def _dump_ram_object_profile(name: str, date: str, keep_alive: Any = None) -> None:
    """Dump deep memory usage of the 100 largest objects in memory.

    Args:
        name: Name of the operation.
        date: Date string for naming.
        keep_alive: Optional object to keep in scope during profiling.
    """
    import gc

    import numpy as np
    import pandas as pd
    from geopandas import GeoDataFrame, GeoSeries

    def get_deep_size(obj: Any) -> int:
        """Get deep size of object, including numpy and pandas buffers.

        Args:
            obj: The object to measure.

        Returns:
            The size of the object in bytes.
        """
        # Handle cases where obj might be a proxy or a weakref that died during traversal
        try:
            from geb.store import DynamicArray

            if isinstance(obj, np.ndarray):
                return int(obj.nbytes)
            if isinstance(obj, DynamicArray):
                # DynamicArray._data is the underlying numpy buffer (max_n size)
                return int(obj._data.nbytes)
            if isinstance(obj, (pd.DataFrame, pd.Series, GeoDataFrame, GeoSeries)):
                # memory_usage(deep=True) includes the underlying buffer and string contents
                usage = obj.memory_usage(deep=True)
                return int(usage.sum() if hasattr(usage, "sum") else usage)
            return sys.getsizeof(obj)
        except (ReferenceError, AttributeError):
            # If the object died or is inaccessible, it's essentially 0 bytes now
            return 0

    # Capture objects and their sizes safely
    objects_with_sizes: list[tuple[Any, int]] = []

    # We identify our own internal structures to skip them in the report
    internal_ids = {id(objects_with_sizes), id(get_deep_size)}

    for obj in gc.get_objects():
        try:
            current_id = id(obj)
            # Skip if this is one of our own tracking objects
            if current_id in internal_ids:
                continue

            # We skip some very common small types to speed up and reduce ReferenceError risk
            if isinstance(obj, (int, str, float, bool, type(None))):
                continue
            size = get_deep_size(obj)
            if size > 0:
                objects_with_sizes.append((obj, size))
        except (ReferenceError, AttributeError):
            continue

    # Get total memory and count
    total_objects_count = len(objects_with_sizes)
    total_memory_bytes = sum(size for _, size in objects_with_sizes)

    # Sort and get top 100
    top_objects_with_size = sorted(
        objects_with_sizes, key=lambda x: x[1], reverse=True
    )[:100]

    profiling_folder = Path("profiling")
    txt_path = profiling_folder / f"{name}_ram_objects_{date}.txt"

    with open(txt_path, "w", encoding="utf-8") as f:
        f.write("=" * 80 + "\n")
        f.write(f"RAM OBJECT PROFILE: {name}\n")
        f.write(f"DATE: {date}\n")
        f.write("=" * 80 + "\n\n")

        f.write("SUMMARY STATISTICS:\n")
        f.write(f"  Total tracked objects (non-trivial): {total_objects_count:,}\n")
        f.write(
            f"  Total memory of tracked objects: {total_memory_bytes / 1024 / 1024:.2f} MB\n"
        )
        f.write(
            f"  Average tracked object size: {total_memory_bytes / max(1, total_objects_count):.2f} bytes\n\n"
        )

        f.write("TOP 100 LARGEST OBJECTS:\n")
        f.write("-" * 80 + "\n")
        f.write(
            f"{'Size (MB)':>10} | {'Type':<30} | {'Identity/Repr (truncated)':<35}\n"
        )
        f.write("-" * 80 + "\n")

        for obj, size in top_objects_with_size:
            size_mb = size / 1024 / 1024
            obj_type = str(type(obj))
            try:
                # Replace newlines with spaces to keep it on one line
                obj_repr = repr(obj).replace("\n", " ").replace("\r", " ")
                # Allow very long representations now
                if len(obj_repr) > 1000:
                    obj_repr = obj_repr[:997] + "..."
            except Exception:
                obj_repr = "<REPR FAILED>"

            f.write(f"{size_mb:10.4f} | {obj_type[:30]:<30} | {obj_repr}\n")

            # Try to find what refers to this object (where it is referenced from)
            try:
                referrers = gc.get_referrers(obj)
                # Filter out the objects we just created in this function
                actual_referrers = [
                    ref
                    for ref in referrers
                    if ref is not obj
                    and ref is not objects_with_sizes
                    and not isinstance(ref, (list, dict, tuple))
                    or (isinstance(ref, (dict, list)) and len(ref) > 0)
                ][:3]

                if actual_referrers:
                    f.write(f"{'':>10} | {'  Referenced by:':<30} | ")
                    ref_reprs = []
                    for ref in actual_referrers:
                        ref_type = type(ref).__name__
                        ref_name = ""

                        # Try to find the variable name if the referrer is a dict (like __dict__)
                        if isinstance(ref, dict):
                            for key, val in ref.items():
                                if val is obj:
                                    ref_name = f"['{key}'] "
                                    break
                        elif hasattr(ref, "__dict__"):
                            ref_dict = getattr(ref, "__dict__")
                            for key, val in ref_dict.items():
                                if val is obj:
                                    ref_name = f".{key} "
                                    break

                        try:
                            r_repr = repr(ref).replace("\n", " ")
                            if len(r_repr) > 500:
                                r_repr = r_repr[:497] + "..."
                        except Exception:
                            r_repr = "<REPR FAILED>"
                        ref_reprs.append(f"{ref_type}{ref_name}({r_repr})")
                    f.write(" | ".join(ref_reprs) + "\n")
            except Exception:
                pass

        f.write("-" * 80 + "\n")

    print(f"RAM object profile saved to: {txt_path}")


def get_model_builder_class(custom_model: None | str) -> type:
    """Get model builder class.

    This is usually the GEBModelBuild class, but can be a custom model builder class
    from the geb.build.custom_models module. This class would usually
    specify some custom build methods, but largely re-use the existing GEBModelBuild methods.

    Args:
        custom_model: Name of the custom model to use. If None, the default GEBModelBuild is used.
            custom_models are available in the geb.build.custom_models module.

    Returns:
        Model builder class.

    Raises:
        ValueError: If the custom model is not found in the geb.build.custom_models module.
    """
    if custom_model is None:
        return GEBModelBuild
    else:
        from geb import build as geb_build

        importlib.import_module(
            "." + custom_model.split(".")[0], package="geb.build.custom_models"
        )
        if not hasattr(geb_build, "custom_models"):
            raise ValueError("Custom models module not found")
        return attrgetter(custom_model)(geb_build.custom_models)


def customize_data_catalog(data_catalog: Path, data_root: None | Path = None) -> Path:
    """This functions adds the GEB_DATA_ROOT to the data catalog if it is set as an environment variable.

    This enables reading the data catalog from a different location than the location of the yml-file
    without the need to specify root in the meta of the data catalog.

    Args:
        data_catalog: List of paths to data catalog yml files.
        data_root: Root folder where the data is located. If None, the data catalog is not modified.

    Returns:
        List of paths to data catalog yml files, possibly modified to include the data_root.
    """
    if data_root:
        with open(data_catalog, "r") as stream:
            data_catalog_yml = yaml.load(stream, Loader=yaml.FullLoader)

            if "meta" not in data_catalog_yml:
                data_catalog_yml["meta"] = {}
            data_catalog_yml["meta"]["root"] = str(data_root)

        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".yml") as tmp:
            yaml.dump(data_catalog_yml, tmp, default_flow_style=False)
        return Path(tmp.name)
    else:
        return data_catalog


def get_builder(
    config: Path | dict[str, Any],
    data_catalog: Path,
    logger: logging.Logger,
    custom_model: str | None,
    data_provider: str | None,
    data_root: Path | None,
) -> GEBModelBuild:
    """Get model builder.

    Args:
        config: Path to the model configuration file.
        data_catalog: Path to the data catalog file.
        logger: Logger instance to pass to the model builder.
        custom_model: Name of the custom model to use. If None, the default GEBModelBuild is used.
            custom_models are available in the geb.build.custom_models module.
        data_provider: Data variant to use from data catalog (see hydroMT documentation).
        data_root: Root folder where the data is located. If None, the data catalog is not modified.

    Returns:
        Instance of the model builder.
    """
    config = parse_config(config, schema=Config)
    input_folder = Path(config["general"]["input_folder"])

    data_catalog = customize_data_catalog(data_catalog, data_root=data_root)

    arguments = {
        "root": input_folder,
        "data_catalog": data_catalog,
        "logger": logger,
        "data_provider": data_provider,
    }

    builder_class = get_model_builder_class(custom_model)(**arguments)
    build_method.validate_tree()
    return builder_class


def init_fn(
    config: str | Path,
    build_config: str | Path,
    update_config: str | Path,
    working_directory: str | Path,
    from_example: str,
    basin_id: str | None = None,
    ISO3: str | None = None,
    overwrite: bool = False,
) -> None:
    """Create a new model.

    Args:
        config: Path to the model configuration file to create.
        build_config: Path to the model build configuration file to create.
        update_config: Path to the model update configuration file to create.
        working_directory: Working directory for the model.
        from_example: Name of the example to use as a base for the model.
        basin_id: Basin ID(s) to use for the model. Can be a comma-separated list of integers.
            If not set, the basin ID is taken from the config file.
        ISO3: ISO3 country code to use for the model. Cannot be used together with --basin-id.
        overwrite: If True, overwrite existing config and build config files. Defaults to False.

    Raises:
        ValueError: If both basin_id and ISO3 are set.

    """
    if basin_id is not None and ISO3 is not None:
        raise ValueError("Cannot use --basin-id and --ISO3 together.")

    config: Path = Path(config)
    build_config: Path = Path(build_config)
    update_config: Path = Path(update_config)
    working_directory: Path = Path(working_directory)

    if not working_directory.exists():
        working_directory.mkdir(parents=True, exist_ok=True)

    def init_operation(logger: logging.Logger, **kwargs: Any) -> None:
        logger.info("Initializing model from example '%s'.", from_example)

        if config.exists() and not overwrite:
            raise FileExistsError(
                f"Config file {config} already exists. Please remove it or use a different name, or use --overwrite."
            )

        if build_config.exists() and not overwrite:
            raise FileExistsError(
                f"Build config file {build_config} already exists. Please remove it or use a different name, or use --overwrite."
            )

        if update_config.exists() and not overwrite:
            raise FileExistsError(
                f"Update config file {update_config} already exists. Please remove it or use a different name, or use --overwrite."
            )

        example_folder: Path = GEB_PACKAGE_DIR / "examples" / from_example
        if not example_folder.exists():
            raise FileNotFoundError(
                f"Example folder {example_folder} does not exist. Did you use the right --from-example option?"
            )

        config_dict: dict = yaml.load(
            open(example_folder / CONFIG_DEFAULT, "r"),
            Loader=DetectDuplicateKeysYamlLoader,
        )

        if basin_id is not None:
            # Allow passing a comma-separated list of integers
            if "," in basin_id:
                basin_ids: list[int] = [
                    int(x) for x in basin_id.split(",") if x.strip()
                ]
            else:
                basin_ids: int = int(basin_id)

            config_dict["general"]["region"]["subbasin"] = basin_ids
        elif ISO3 is not None:
            del config_dict["general"]["region"]["subbasin"]
            config_dict["general"]["region"] = {
                "geom": {
                    "source": "GADM_level1",
                    "key": ISO3,
                    "column": "GID_0",
                }
            }

        with open(config, "w") as f:
            # do not sort keys, to keep the order of the config file
            yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False)

        shutil.copy(example_folder / BUILD_DEFAULT, build_config)
        shutil.copy(example_folder / UPDATE_DEFAULT, update_config)

        logger.info(
            "Initialized model files: %s, %s, %s",
            config,
            build_config,
            update_config,
        )

    with WorkingDirectory(working_directory):
        _run_with_optional_profiling(
            PROFILE_SPEED_DEFAULT,
            PROFILE_RAM_DEFAULT,
            init_operation,
            name="init",
            logger=create_logger(name="init"),
        )


def set_fn(
    config: Path,
    working_directory: Path = WORKING_DIRECTORY_DEFAULT,
    **kwargs: Any,
) -> None:
    """Set model configuration values by updating a YAML configuration file.

    This function loads the existing configuration from the specified file,
    updates it with the provided keyword arguments (supporting nested keys
    using dot notation, e.g., 'section.subsection.key'), and saves the
    modified configuration back to the file.

        config: Path to the model configuration file (as a string or Path object).
        working_directory: Working directory for the model.
        **kwargs: Keyword arguments representing keys and values to set in the config.
                  Keys can be nested using dots (e.g., 'model.lr' sets 'lr' under 'model').

    Note:
        If a nested key does not exist, intermediate dictionaries are created automatically.
        The file is overwritten with the updated configuration in YAML format.

    Args:
        config: Path to the model configuration file.
        working_directory: Working directory for the model.
        **kwargs: Keyword arguments to set in the config file.
    """

    def set_operation(logger: logging.Logger, **kwargs: Any) -> None:
        logger.info("Updating model configuration values.")

        config_dict: dict[str, Any] = parse_config(config, schema=Config)
        for key, value in kwargs.items():
            if key.endswith("+"):
                key: str = key[:-1]
                create: bool = True
            else:
                create: bool = False

            keys: list[str] = key.split(".")
            d = config_dict

            for k in keys[:-1]:
                if k not in d or not isinstance(d[k], dict):
                    if create:
                        d[k] = {}
                    else:
                        raise KeyError(
                            f"Key '{k}' not found in config. If you want to create it, use the '+' suffix for the KEY."
                        )
                d: dict[str, Any] = d[k]

            if value == "null":
                value = None
            elif value == "true":
                value = True
            elif value == "false":
                value = False

            if not create and keys[-1] not in d:
                raise KeyError(
                    f"Key '{keys[-1]}' not found in config. If you want to create it, use the '+' suffix for the KEY."
                )
            d[keys[-1]] = value

        with open(config, "w") as f:
            yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False)

        logger.info("Updated configuration file: %s", config)

    with WorkingDirectory(working_directory):
        _run_with_optional_profiling(
            PROFILE_SPEED_DEFAULT,
            PROFILE_RAM_DEFAULT,
            set_operation,
            name="set",
            logger=create_logger(name="set"),
        )


def build_fn(
    data_catalog: Path = DATA_CATALOG_DEFAULT,
    config: Path | dict[str, Any] = CONFIG_DEFAULT,
    build_config: Path | dict[str, Any] = BUILD_DEFAULT,
    working_directory: Path = WORKING_DIRECTORY_DEFAULT,
    data_provider: str = DATA_PROVIDER_DEFAULT,
    data_root: Path = DATA_ROOT_DEFAULT,
    continue_: bool = False,
    profile_speed: bool = PROFILE_SPEED_DEFAULT,
    profile_ram: bool = PROFILE_RAM_DEFAULT,
) -> None:
    """Build model.

    Args:
        data_catalog: Path to the data catalog file.
        config: Path to the model configuration file.
        build_config: Path to the model build configuration file.
        working_directory: Working directory for the model.
        data_provider: Data variant to use from data catalog (see hydroMT documentation).
        data_root: Root folder where the data is located. If None, the data catalog is not modified.
        continue_: Continue previous build if it was interrupted or failed.
        profile_speed: If True, run the build with speed profiling.
        profile_ram: If True, run the build with RAM profiling.
    """
    build_config_input: Path | dict[str, Any] = build_config

    def build_operation(logger: logging.Logger, **kwargs: Any) -> None:
        parsed_build_config = parse_config(build_config_input)
        model = get_builder(
            config,
            data_catalog,
            logger=logger,
            custom_model=parsed_build_config["_custom_model"]
            if "_custom_model" in parsed_build_config
            else None,
            data_provider=data_provider,
            data_root=data_root,
        )
        methods: dict[str, Any] = {
            method: args
            for method, args in parsed_build_config.items()
            if not method.startswith("_")
        }
        model.build(
            methods=methods,
            region=parse_config(config, schema=Config)["general"]["region"],
            continue_=continue_,
        )

    with WorkingDirectory(working_directory):
        _run_with_optional_profiling(
            profile_speed,
            profile_ram,
            build_operation,
            name="build",
            logger=create_logger(name="build"),
        )


def alter_fn(
    data_catalog: Path = DATA_CATALOG_DEFAULT,
    config: Path = CONFIG_DEFAULT,
    build_config: Path | dict[str, Any] = BUILD_DEFAULT,
    working_directory: Path = WORKING_DIRECTORY_DEFAULT,
    from_model: Path = ALTER_FROM_MODEL_DEFAULT,
    data_provider: str = DATA_PROVIDER_DEFAULT,
    data_root: Path = DATA_ROOT_DEFAULT,
    profile_speed: bool = PROFILE_SPEED_DEFAULT,
    profile_ram: bool = PROFILE_RAM_DEFAULT,
) -> None:
    """Create alternative version from base model with only changed files.

    This function is useful to create a new model based on an existing one, but with
    only a few changes. It will copy the base model and overwrite the files that are
    specified in the config and build config files. The rest of the files will be
    linked to the original model to reduce disk space.

    Args:
        data_catalog: Path to the data catalog file.
        config: Path to the model configuration file.
        build_config: Path to the model build configuration file.
        working_directory: Working directory for the model.
        from_model: Folder for the existing model.
        data_provider: Data variant to use from data catalog (see hydroMT documentation).
        data_root: Root folder where the data is located. If None, the data catalog is not modified.
        profile_speed: If True, run the alter flow with speed profiling.
        profile_ram: If True, run the alter flow with RAM profiling.
    """
    from_model: Path = Path(from_model)
    build_config_input: Path | dict[str, Any] = build_config

    def alter_operation(logger: logging.Logger, **kwargs: Any) -> None:
        original_config: Path = from_model / config
        if not original_config.exists():
            raise FileNotFoundError(
                f"Config file {original_config} does not exist in the original model. Cannot create alternative model based on it."
            )

        # if config does not exist, create a new config that inherits from the original model
        if not config.exists():
            original_config: Path = from_model / config

            config_dict: dict[str, str] = {"inherits": str(original_config)}
            with open(config, "w") as f:
                # do not sort keys, to keep the order of the config file
                yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False)

        # if config exists, we just make sure it inherits from the original model
        # if this is set already, we just leave it as is and assume the user knows what they are doing
        else:
            # Read existing config
            with open(config, "r") as f:
                raw_config = yaml.load(f, Loader=DetectDuplicateKeysYamlLoader)
            if "inherits" not in raw_config:
                raw_config["inherits"] = str(from_model / config)
            with open(config, "w") as f:
                yaml.dump(raw_config, f, default_flow_style=False, sort_keys=False)

        config_from_original_model = parse_config(from_model / config, schema=Config)
        input_folder: Path = Path(config_from_original_model["general"]["input_folder"])

        original_input_path: Path = from_model / input_folder

        original_files = read_params(original_input_path / "files.yml")

        for file_class, files in original_files.items():
            for file_name, file_path in files.items():
                if not file_path.startswith("/"):
                    original_files[file_class][file_name] = str(
                        Path("..") / original_input_path / file_path
                    )

        input_folder.mkdir(parents=True, exist_ok=True)
        with open(input_folder / "files.yml", "w") as f:
            yaml.dump(original_files, f, default_flow_style=False)

        parsed_build_config = parse_config(build_config_input)
        model = get_builder(
            config,
            data_catalog,
            logger=logger,
            custom_model=parsed_build_config["_custom_model"]
            if "_custom_model" in parsed_build_config
            else None,
            data_provider=data_provider,
            data_root=data_root,
        )
        methods = {
            method: args
            for method, args in parsed_build_config.items()
            if not method.startswith("_")
        }

        model.update(
            methods=methods,
        )

    with WorkingDirectory(working_directory):
        _run_with_optional_profiling(
            profile_speed,
            profile_ram,
            alter_operation,
            name="alter",
            logger=create_logger(name="alter"),
        )


def update_version_fn(
    data_catalog: Path = DATA_CATALOG_DEFAULT,
    config: Path | dict[str, Any] = CONFIG_DEFAULT,
    working_directory: Path = WORKING_DIRECTORY_DEFAULT,
    data_provider: str = DATA_PROVIDER_DEFAULT,
    data_root: Path = DATA_ROOT_DEFAULT,
    profile_speed: bool = PROFILE_SPEED_DEFAULT,
    profile_ram: bool = PROFILE_RAM_DEFAULT,
    **kwargs: Any,
) -> None:
    """Update the model version file to the current model version.

    This function initializes the GEBModelBuild, which automatically checks and updates
    the version file if it is outdated, printing any necessary update instructions.

    Args:
        data_catalog: Path to the data catalog file.
        config: Path to the model configuration file.
        working_directory: Working directory for the model.
        data_provider: Data variant to use from data catalog.
        data_root: Root folder where the data is located.
        profile_speed: If True, run the update flow with speed profiling.
        profile_ram: If True, run the update flow with RAM profiling.
        **kwargs: Additional keyword arguments.
    """

    def update_version_operation(logger: logging.Logger, **kwargs: Any) -> None:
        parsed_config = parse_config(config, schema=Config)
        input_folder = Path(parsed_config["general"]["input_folder"])
        custom_model = (
            parsed_config["general"]["custom_model"]
            if "custom_model" in parsed_config["general"]
            else None
        )

        data_catalog_path = customize_data_catalog(data_catalog, data_root=data_root)

        builder_class = get_model_builder_class(custom_model)
        builder_class(
            logger=logger,
            root=input_folder,
            data_catalog=str(data_catalog_path),
            data_provider=data_provider,
        )

    with WorkingDirectory(working_directory):
        _run_with_optional_profiling(
            profile_speed,
            profile_ram,
            update_version_operation,
            name="update-version",
            logger=create_logger(name="update-version"),
        )


def update_fn(
    data_catalog: Path = DATA_CATALOG_DEFAULT,
    config: Path | dict[str, Any] = CONFIG_DEFAULT,
    build_config: Path | dict[str, Any] = BUILD_DEFAULT,
    working_directory: Path = WORKING_DIRECTORY_DEFAULT,
    data_provider: str = DATA_PROVIDER_DEFAULT,
    data_root: Path = DATA_ROOT_DEFAULT,
    profile_speed: bool = PROFILE_SPEED_DEFAULT,
    profile_ram: bool = PROFILE_RAM_DEFAULT,
) -> None:
    """Update model.

    Args:
        data_catalog: Path to the data catalog file.
        config: Path to the model configuration file.
        build_config: Path to the model build configuration file or a specific method within the build file using :: syntax, e.g., 'build.yml::setup_economic_data' to only run the setup_economic_data method. If the method ends with a '+', all subsequent methods are run as well.
        working_directory: Working directory for the model.
        data_provider: Data variant to use from data catalog (see hydroMT documentation).
        data_root: Root folder where the data is located. If None, the data catalog is not modified.
        profile_speed: If True, run the update flow with speed profiling.
        profile_ram: If True, run the update flow with RAM profiling.
    """
    build_config_input: Path | dict[str, Any] = build_config

    def update_operation(logger: logging.Logger, **kwargs: Any) -> None:
        if isinstance(build_config_input, Path):
            build_config_list: list[str] = str(build_config_input).split("::")
            build_config_file: Path = Path(build_config_list[0])

            try:
                parsed_build_config: dict[str, Any] = parse_config(build_config_file)
            except FileNotFoundError:
                if ":" in str(build_config_file) and "::" not in str(build_config_file):
                    raise FileNotFoundError(
                        f"Build config file '{str(build_config_file)}' not found. Did you mean '{str(build_config_file).replace(':', '::')}'?"
                    )
                raise

            methods = {
                method: args
                for method, args in parsed_build_config.items()
                if not method.startswith("_")
            }

            if len(build_config_list) > 1:
                assert len(build_config_list) == 2
                build_config_function: str = build_config_list[1]

                # Check if the method is specified with a trailing '+' or '#'.
                # If + we set a flag to keep all subsequent methods
                # If # we set a flag to keep all dependent methods
                if build_config_function.endswith("+"):
                    build_config_function: str = build_config_function[:-1]
                    keys_to_remove: list[str] = []

                    for key in methods.keys():
                        if key == build_config_function:
                            break
                        else:
                            keys_to_remove.append(key)
                    else:
                        raise KeyError(
                            f"Method '{build_config_function}' not found in build config file '{build_config_file}'. "
                            "Available methods: "
                            f"{', '.join(methods.keys())}"
                        )

                    # remove all functions from the methods dict except the one we want to run
                    for key in keys_to_remove:
                        del methods[key]

                elif build_config_function.endswith("#"):
                    build_config_function: str = build_config_function[:-1]
                    dependents: list[str] = build_method.get_dependents(
                        build_config_function
                    )
                    dependents.append(build_config_function)
                    methods = {
                        dependent: methods[dependent]
                        for dependent in dependents
                        if dependent in methods
                    }

                else:
                    methods = {build_config_function: methods[build_config_function]}

        elif isinstance(build_config_input, dict):
            parsed_build_config = build_config_input
            methods = {
                method: args
                for method, args in parsed_build_config.items()
                if not method.startswith("_")
            }

        else:
            raise ValueError("build_config must be a str or dict.")

        model = get_builder(
            config,
            data_catalog,
            logger=logger,
            custom_model=parsed_build_config["_custom_model"]
            if "_custom_model" in parsed_build_config
            else None,
            data_provider=data_provider,
            data_root=data_root,
        )

        model.update(methods=methods)

    with WorkingDirectory(working_directory):
        _run_with_optional_profiling(
            profile_speed,
            profile_ram,
            update_operation,
            name="update",
            logger=create_logger(name="update"),
        )


def share_fn(
    working_directory: Path,
    name: str,
    include_cache: bool,
    include_output: bool,
) -> None:
    """Share model."""
    with WorkingDirectory(working_directory):
        # create a zip file called model.zip with the folders input, and model files
        # in the working directory
        folders: list = ["input"]
        if include_cache:
            folders.append("cache")
        if include_output:
            folders.append("output")
        files: list = [CONFIG_DEFAULT, BUILD_DEFAULT]
        optional_files: list = [UPDATE_DEFAULT, DATA_CATALOG_DEFAULT]
        optional_folders: list = ["data"]
        zip_filename: str = f"{name}.zip"
        with zipfile.ZipFile(zip_filename, "w") as zipf:
            total_files: int = (
                sum(
                    [
                        sum(len(files) for _, _, files in os.walk(folder))
                        for folder in folders
                    ]
                )
                + sum(
                    [
                        sum(len(files) for _, _, files in os.walk(folder))
                        for folder in optional_folders
                        if os.path.exists(folder)
                    ]
                )
                + len(files)
                + len(optional_files)
            )  # Count total number of files
            progress: int = 0  # Initialize progress counter
            for folder in folders:
                for root, _, filenames in os.walk(folder):
                    for filename in filenames:
                        zipf.write(os.path.join(root, filename))
                        progress += 1  # Increment progress counter
                        print(
                            f"Exporting file {progress}/{total_files} to {zip_filename}",
                            end="\r",
                        )  # Print progress
            for folder in optional_folders:
                if os.path.exists(folder):
                    for root, _, filenames in os.walk(folder):
                        for filename in filenames:
                            zipf.write(os.path.join(root, filename))
                            progress += 1
                            print(
                                f"Exporting file {progress}/{total_files} to {zip_filename}",
                                end="\r",
                            )
            for file in files:
                zipf.write(file)
                progress += 1  # Increment progress counter
                print(
                    f"Exporting file {progress}/{total_files} to {zip_filename}",
                    end="\r",
                )  # Print progress
            for file in optional_files:
                if os.path.exists(file):
                    zipf.write(file)
                progress += 1  # Increment progress counter
                print(
                    f"Exporting file {progress}/{total_files} to {zip_filename}",
                    end="\r",
                )  # Print progress
            print(f"Exporting file {progress}/{total_files} to {zip_filename}")
            print("Done!")


def init_multiple_fn(
    config: str | Path,
    build_config: str | Path,
    update_config: str | Path,
    working_directory: str | Path,
    from_example: str,
    geometry_bounds: str,
    target_area_km2: float,
    cluster_prefix: str,
    init_multiple_dir: str,
    region_shapefile: str | None = None,
    skip_merged_geometries: bool = False,
    skip_visualization: bool = False,
    min_bbox_efficiency: float = 0.99,
    ocean_outlets_only: bool = False,
) -> None:
    """Create multiple models from a geometry by clustering downstream subbasins.

    Args:
        config: Path to the base model configuration file.
        build_config: Path to the base model build configuration file.
        update_config: Path to the base model update configuration file.
        working_directory: Working directory for the models.
        from_example: Name of the example to use as a base for the models.
        geometry_bounds: Bounding box as "xmin,ymin,xmax,ymax" to select subbasins.
        target_area_km2: Target cumulative upstream area per cluster (default: Danube basin ~817,000 km2).
        cluster_prefix: Prefix for cluster directory names.
        region_shapefile: Optional path to region shape file. Defaults to geometry bounds if not specified.
        skip_merged_geometries: If True, skip creating dissolved basin polygon file (much faster).
        skip_visualization: If True, skip creating visualization map (faster).
        min_bbox_efficiency: Minimum bbox efficiency (0-1) for cluster merging. Lower values allow more elongated clusters.
        ocean_outlets_only: If True, only include clusters that flow to the ocean (exclude endorheic basins).
        init_multiple_dir: Name of the subdirectory in models/ where the large scale model directories will be created (e.g. 'large_scale' or 'large_scale2').

    Raises:
        FileNotFoundError: If the example folder does not exist, or if the parent
            models/ directory does not exist.
        ValueError: If geometry_bounds format is invalid.
    """
    # set paths
    config: Path = Path(config)
    build_config: Path = Path(build_config)
    update_config: Path = Path(update_config)
    working_directory: Path = Path(working_directory)

    # Initialize data catalog and logger
    data_catalog_instance = NewDataCatalog()
    with WorkingDirectory(working_directory):
        logger = create_logger("init_multiple")

    # Create the models/init_multiple_dir directory structure.
    # models_dir is the parent 'models' folder expected at <cwd>/../models;
    # init_multiple_dir_path is the target subdirectory to create.
    models_dir = Path.cwd().parent / "models"
    if not models_dir.is_dir():
        raise FileNotFoundError(
            f"Models directory not found: {models_dir}\n"
            "Run 'geb init-multiple' from within the GEB repository root, "
            f"or ensure a 'models' directory exists at {models_dir}."
        )
    init_multiple_dir_path: Path = models_dir / init_multiple_dir
    init_multiple_dir_path.mkdir(parents=True, exist_ok=True)

    # create river
    logger.info("Starting multiple model initialization")
    logger.info(f"Using geometry bounds: {geometry_bounds}")
    logger.info(f"Target area: {target_area_km2:,.0f} km²")

    logger.info("Loading river network...")
    river_graph = get_river_graph(data_catalog_instance)

    # Parse geometry bounds and convert to geodataframe
    bounds = [float(x.strip()) for x in geometry_bounds.split(",")]
    if len(bounds) != 4:
        raise ValueError(
            "Invalid geometry_bounds format. Expected 'xmin,ymin,xmax,ymax'."
        )
    xmin, ymin, xmax, ymax = bounds

    bbox_geom = gpd.GeoDataFrame(
        geometry=[box(xmin, ymin, xmax, ymax)], crs="EPSG:4326"
    )

    downstream_subbasins = get_all_downstream_subbasins_in_geom(
        data_catalog_instance, bbox_geom, logger
    )  # get all downstream subbasins in the bounding box geometry

    if not downstream_subbasins:
        raise ValueError("No downstream subbasins found in the specified geometry")

    logger.info(f"Found {len(downstream_subbasins)} downstream subbasins")

    logger.info("Clustering subbasins by following coastline...")
    clusters = cluster_subbasins_following_coastline(
        data_catalog_instance,
        downstream_subbasins,
        target_area_km2=target_area_km2,
        logger=logger,
        river_graph=river_graph,
        min_bbox_efficiency=min_bbox_efficiency,
    )

    logger.info(f"Created {len(clusters)} clusters")

    # Verify example folder exists
    example_folder: Path = GEB_PACKAGE_DIR / "examples" / from_example
    if not example_folder.exists():
        raise FileNotFoundError(
            f"Example folder {example_folder} does not exist. Did you use the right --from-example option?"
        )

    logger.info(f"Creating cluster configurations using example: {from_example}")

    # Create cluster configurations
    cluster_directories = create_multi_basin_configs(
        clusters=clusters,
        working_directory=init_multiple_dir_path,
        cluster_prefix=cluster_prefix,
        data_catalog=data_catalog_instance,
        river_graph=river_graph,
    )

    # save geoparquet and maps to default location
    save_geoparquet = init_multiple_dir_path / f"{cluster_prefix}_outlets.geoparquet"
    save_map = init_multiple_dir_path / f"{cluster_prefix}_clusters_map.png"

    logger.info(f"Saving outlet basins to geoparquet: {save_geoparquet}")
    # Save outlet-only clusters to geoparquet (simplified geometries)
    save_clusters_to_geoparquet(
        clusters=clusters,
        data_catalog=data_catalog_instance,
        output_path=save_geoparquet,
        cluster_prefix=cluster_prefix,
    )

    # Save complete basins as merged geometries (full upstream basins as single polygons)
    # This is slow for large datasets, so allow skipping
    if not skip_merged_geometries:
        merged_basins_path = (
            init_multiple_dir_path / f"{cluster_prefix}_complete_basins.geoparquet"
        )
        logger.info(
            f"Saving complete basins as merged geometries: {merged_basins_path}"
        )
        save_clusters_as_merged_geometries(
            clusters=clusters,
            data_catalog=data_catalog_instance,
            river_graph=river_graph,
            output_path=merged_basins_path,
            cluster_prefix=cluster_prefix,
            buffer_distance_km=5.0,  # 5km buffer to merge nearby polygons (reduces MultiPolygon complexity)
        )
    else:
        logger.info("Skipping merged geometries (--skip-merged-geometries flag set)")

    # Create visualization map (optional)
    if not skip_visualization:
        logger.info(f"Creating visualization map: {save_map}")
        create_cluster_visualization_map(
            clusters=clusters,
            data_catalog=data_catalog_instance,
            river_graph=river_graph,
            output_path=save_map,
            cluster_prefix=cluster_prefix,
        )
    else:
        logger.info("Skipping visualization map (--skip-visualization flag set)")

    logger.info(
        f"Successfully created {len(cluster_directories)} model configurations:"
    )
