from __future__ import annotations

import gin


@gin.configurable
class PipelineConfig:
    """Pipeline stage configuration: what to process and where to save."""

    def __init__(
        self,
        input_root: str = "/data/UAV-GeoLoc",
        output_root: str = "/data/UAV-GeoLoc-aug",
        stages: list[str] | None = None,
        save_npy: bool = True,
        save_vis: bool = True,
        save_concat: bool = False,
        save_safetensors: bool = True,
        cleanup_npy: bool = False,
        seg_fix_dark_water: bool = True,
        seg_fix_uniform_water: bool = True,
        seg_reclassify_wetland: bool = False,
        resume: bool = True,
        subset: str | None = None,
        source: str | None = None,
        log_level: str = "INFO",
    ) -> None:
        self.input_root = input_root
        self.output_root = output_root
        self.stages = stages or ["depth", "edges", "segmentation"]
        self.save_npy = save_npy
        self.save_vis = save_vis
        self.save_concat = save_concat
        self.save_safetensors = save_safetensors
        self.cleanup_npy = cleanup_npy
        self.seg_fix_dark_water = seg_fix_dark_water
        self.seg_fix_uniform_water = seg_fix_uniform_water
        self.seg_reclassify_wetland = seg_reclassify_wetland
        self.resume = resume
        self.subset = subset
        self.source = source
        self.log_level = log_level


def get_pipeline_cfg(path2cfg: str) -> PipelineConfig:
    """Load ONLY pipeline config (for isolated testing)."""
    gin.clear_config()
    gin.parse_config_file(f"{path2cfg}pipeline.gin")
    return PipelineConfig()
