#!/usr/bin/env python3

from pathlib import Path
import argparse
import json


VALID_VARIABLES = {"pr", "tasmax", "tasmin"}


def sorted_dict(obj):
    """
    Recursively sort dictionary keys for stable output.
    """
    if isinstance(obj, dict):
        return {key: sorted_dict(obj[key]) for key in sorted(obj)}
    if isinstance(obj, list):
        return sorted(obj)
    return obj


def add_variable(node, variable):
    node.setdefault("variables", [])

    if variable not in node["variables"]:
        node["variables"].append(variable)
        node["variables"].sort()


def discover_options(root):
    """
    Expected structures:

      cmip5/downscaling/model/variable
      cmip6/downscaling/model/time/variable

    Output structure:

      cmip5/downscaling/model/variables
      cmip6/downscaling/model/time/variables
    """
    data = {}

    for cmip_path in sorted(root.iterdir()):
        if not cmip_path.is_dir():
            continue

        cmip = cmip_path.name

        if cmip not in {"cmip5", "cmip6"}:
            continue

        data.setdefault(cmip, {})

        for downscaling_path in sorted(cmip_path.iterdir()):
            if not downscaling_path.is_dir():
                continue

            downscaling = downscaling_path.name
            data[cmip].setdefault(downscaling, {})

            for model_path in sorted(downscaling_path.iterdir()):
                if not model_path.is_dir():
                    continue

                model = model_path.name

                if cmip == "cmip5":
                    data[cmip][downscaling].setdefault(model, {})

                    for variable_path in sorted(model_path.iterdir()):
                        if not variable_path.is_dir():
                            continue

                        variable = variable_path.name

                        if variable not in VALID_VARIABLES:
                            continue

                        add_variable(data[cmip][downscaling][model], variable)

                elif cmip == "cmip6":
                    data[cmip][downscaling].setdefault(model, {})

                    for time_path in sorted(model_path.iterdir()):
                        if not time_path.is_dir():
                            continue

                        time = time_path.name
                        data[cmip][downscaling][model].setdefault(time, {})

                        for variable_path in sorted(time_path.iterdir()):
                            if not variable_path.is_dir():
                                continue

                            variable = variable_path.name

                            if variable not in VALID_VARIABLES:
                                continue

                            add_variable(data[cmip][downscaling][model][time], variable)

    return sorted_dict(data)


def write_javascript(output_file, data):
    js_text = f"""// Auto-generated. Do not edit manually.

export const cmipOptions = {json.dumps(data, indent=2)};
"""

    output_file.write_text(js_text)


def main():
    parser = argparse.ArgumentParser(
        description="Create JavaScript dropdown options from CMIP directory structure."
    )

    parser.add_argument(
        "input_dir",
        help="Root directory containing cmip5/ and cmip6/ directories.",
    )

    parser.add_argument(
        "-o",
        "--output",
        default="cmip-options.js",
        help="Output JavaScript file.",
    )

    args = parser.parse_args()

    root = Path(args.input_dir).resolve()
    output_file = Path(args.output).resolve()

    if not root.exists():
        raise FileNotFoundError(f"Input directory does not exist: {root}")

    data = discover_options(root)
    write_javascript(output_file, data)

    print(f"Wrote {output_file}")


if __name__ == "__main__":
    main()