import pystac
import xarray as xr
import zarr
from packaging.version import Version
from typing import Any

# === Option 1: Manual UI Discovery ===
# Replace these values with what you find in the STAC Browser UI.
# Note: JSON uses lowercase true/false — capitalize to True/False for Python.

zarr_url = "s3://mdmf/gdp/gridMET.zarr/"

open_kwargs = {
    "chunks": {},
    "consolidated": True,  # JSON: "consolidated": true → Python: True
    "engine": "zarr",
}

storage_options = {
    "anon": True,  # JSON: "anon": true → Python: True
    "client_kwargs": {
        "endpoint_url": "https://usgs.osn.mghpcc.org/"
    },
}

def connect_catalog(catalog_url: str) -> pystac.Catalog:
    """Connect to a STAC catalog and return the catalog object.

    Args:
        catalog_url: URL to the STAC catalog's root JSON document.

    Returns:
        A pystac.Catalog instance connected to the remote catalog.
    """
    return pystac.Catalog.from_file(catalog_url)


def list_collections(catalog: pystac.Catalog) -> list[str]:
    """List all child collections in the catalog.

    Prints each collection's ID and its cite-as link (if available).

    Args:
        catalog: A pystac.Catalog instance.

    Returns:
        A list of collection ID strings.
    """
    collection_ids = []
    for child in catalog.get_children():
        collection_ids.append(child.id)
        cite_link = None
        for link in child.links:
            if link.rel == "cite-as":
                cite_link = link.target
                break
        cite_str = cite_link if cite_link else "Not available"
        print(f"  {child.id} \u2014 cite-as: {cite_str}")
    return collection_ids


def select_collection(catalog: pystac.Catalog, collection_id: str):
    """Navigate into a specified collection.

    If the collection contains assets, prints them (ID, title, description).
    If it contains child collections, prints those instead.

    Args:
        catalog: A pystac.Catalog instance.
        collection_id: The ID of the collection to select.

    Returns:
        The selected pystac collection object.

    Raises:
        ValueError: If collection_id is not found among the catalog's children.
    """
    available_ids = []
    target = None
    for child in catalog.get_children():
        available_ids.append(child.id)
        if child.id == collection_id:
            target = child

    if target is None:
        raise ValueError(
            f"Collection '{collection_id}' not found. "
            f"Available collections: {available_ids}"
        )

    if target.assets:
        print(f"Collection '{collection_id}' assets:")
        for asset_key, asset in target.assets.items():
            title = asset.title or "No title"
            description = asset.description or "No description"
            print(f"  {asset_key}: {title} \u2014 {description}")
    else:
        print(f"Collection '{collection_id}' contains sub-collections:")
        for sub_child in target.get_children():
            print(f"  {sub_child.id}")

    return target


def get_asset_info(collection, asset_id: str) -> dict[str, Any]:
    """Retrieve xarray-compatible metadata from a collection's asset.

    Args:
        collection: A pystac collection object (returned by select_collection).
        asset_id: The ID (key) of the asset to retrieve metadata from.

    Returns:
        A dict with keys: url, storage_options, open_kwargs.

    Raises:
        KeyError: If asset_id doesn't exist or xarray metadata fields are missing.
    """
    assets = collection.assets
    if asset_id not in assets:
        raise KeyError(
            f"Asset '{asset_id}' not found. Available: {list(assets.keys())}"
        )

    asset = assets[asset_id]
    extra = asset.extra_fields

    if "xarray:storage_options" not in extra:
        raise KeyError(
            f"Asset '{asset_id}' is missing 'xarray:storage_options' field."
        )
    if "xarray:open_kwargs" not in extra:
        raise KeyError(
            f"Asset '{asset_id}' is missing 'xarray:open_kwargs' field."
        )

    return {
        "url": asset.href,
        "storage_options": extra["xarray:storage_options"],
        "open_kwargs": extra["xarray:open_kwargs"],
    }


def open_zarr_dataset(
    url: str,
    storage_options: dict,
    open_kwargs: dict,
) -> xr.Dataset:
    """Open a zarr dataset with automatic zarr v2/v3 compatibility handling.

    Detects the installed zarr package version and passes `zarr_format=2`
    when zarr >= 3.0.0, since zarr v3 defaults to format 3 but the WMA
    catalog stores data in zarr format 2.

    Args:
        url: S3 or HTTP URL to the zarr store.
        storage_options: Dict of fsspec storage options (e.g., anon, endpoint_url).
        open_kwargs: Dict of keyword arguments for xarray.open_dataset
                     (e.g., engine, chunks, consolidated).

    Returns:
        An xarray.Dataset opened from the zarr store.
    """
    kwargs = {"storage_options": storage_options, **open_kwargs}
    if Version(zarr.__version__) >= Version("3.0.0"):
        kwargs["zarr_format"] = 2
    return xr.open_dataset(url, **kwargs)

# Get asset metadata — replace "zarr-s3-osn" with your choice from the list above
asset_info = get_asset_info(collection, "zarr-s3-osn")