"""
Interop adapters for rasterio, xarray, and GeoTIFF export.

All optional dependencies (xarray, geopandas) are imported lazily inside
function bodies so that they are only required when actually used.
"""

import os
from typing import TYPE_CHECKING, List, Optional

import numpy as np

if TYPE_CHECKING:
    import xarray


def frame_to_rasterio(
    frame: np.ndarray,
    crs,
    transform,
    nodata: float = float("nan"),
):
    """
    Wrap a (H, W, C) numpy array in a rasterio MemoryFile.

    The returned MemoryFile can be passed directly to rasterio workflows.

    Args:
        frame: Array of shape (H, W, C).
        crs: rasterio CRS.
        transform: Affine transform.
        nodata: NoData value.

    Returns:
        rasterio.io.MemoryFile (open; caller should close when done).
    """
    import rasterio
    from rasterio.io import MemoryFile

    if frame.ndim == 2:
        frame = frame[:, :, np.newaxis]

    H, W, C = frame.shape
    memfile = MemoryFile()
    with memfile.open(
        driver="GTiff",
        height=H,
        width=W,
        count=C,
        dtype=frame.dtype,
        crs=crs,
        transform=transform,
        nodata=nodata,
    ) as dst:
        # rasterio expects (C, H, W) — transpose from (H, W, C)
        for b in range(C):
            dst.write(frame[:, :, b], indexes=b + 1)

    return memfile


def save_frame_as_tif(
    frame: np.ndarray,
    path: str,
    crs,
    transform,
    out_dtype: str = "float32",
) -> None:
    """
    Save a (H, W, C) frame as a GeoTIFF file.

    Args:
        frame: Array of shape (H, W, C).
        path: Output file path.
        crs: rasterio CRS.
        transform: Affine transform.
        out_dtype: Output data type string.
    """
    import rasterio

    if frame.ndim == 2:
        frame = frame[:, :, np.newaxis]

    H, W, C = frame.shape

    nodata = np.nan if "float" in out_dtype else 0
    out_frame = frame.astype(out_dtype)
    if "int" in out_dtype:
        out_frame = np.nan_to_num(out_frame, nan=0)

    profile = {
        "driver": "GTiff",
        "height": H,
        "width": W,
        "count": C,
        "dtype": out_dtype,
        "crs": crs,
        "transform": transform,
        "compress": "lzw",
        "nodata": nodata,
        "bigtiff": "YES",
    }

    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)

    with rasterio.open(path, "w", **profile) as dst:
        for b in range(C):
            dst.write(out_frame[:, :, b], indexes=b + 1)


def save_all_tifs(
    ds,
    output_dir: str,
    out_dtype: str = "float32",
    progress: bool = True,
    name_prefix: str = "Frame",
) -> None:
    """
    Export all frames from a GNDCDataset as GeoTIFF files.

    Args:
        ds: GNDCDataset instance.
        output_dir: Output directory.
        out_dtype: Output data type.
        progress: Whether to show a progress bar.
        name_prefix: Filename prefix.
    """
    os.makedirs(output_dir, exist_ok=True)

    timestamps = ds._meta.get("timestamps", [])
    T = ds.n_timestamps
    iterator = range(T)

    if progress:
        try:
            from tqdm import tqdm
            iterator = tqdm(iterator, desc="Exporting TIFs")
        except ImportError:
            pass

    for t_idx in iterator:
        frame = ds.read(t=t_idx)
        ts_str = timestamps[t_idx] if t_idx < len(timestamps) else str(t_idx)
        safe_date = str(ts_str).split(" ")[0].replace("-", "")
        fname = f"{name_prefix}_{safe_date}.tif"
        out_path = os.path.join(output_dir, fname)
        save_frame_as_tif(frame, out_path, ds.crs, ds.transform, out_dtype)


def dataset_to_xarray(
    ds,
    times: Optional[List[int]] = None,
) -> "xarray.Dataset":
    """
    Convert a GNDCDataset to an xarray.Dataset.

    Each band becomes a DataArray variable with dimensions (time, y, x).

    Args:
        ds: GNDCDataset instance.
        times: List of time indices. None = all timestamps.

    Returns:
        xarray.Dataset with spatial coordinates and CRS attributes.
    """
    try:
        import xarray as xr
    except ImportError:
        raise ImportError(
            "xarray is required for to_xarray(). "
            "Install it with: pip install xarray"
        )

    T_total = ds.n_timestamps
    if times is None:
        times = list(range(T_total))

    all_timestamps = ds.timestamps
    time_coords = [all_timestamps[i] for i in times]

    # Build spatial coordinates
    t = ds.transform
    if t is not None:
        x_coords = np.array([t.c + (col + 0.5) * t.a for col in range(ds.width)])
        y_coords = np.array([t.f + (row + 0.5) * t.e for row in range(ds.height)])
    else:
        x_coords = np.arange(ds.width, dtype=np.float64)
        y_coords = np.arange(ds.height, dtype=np.float64)

    # Read all requested frames
    n_bands = ds.n_bands
    band_names = ds.band_names or [f"band_{i}" for i in range(n_bands)]

    data_vars = {}
    for b_idx, b_name in enumerate(band_names):
        cube = np.stack([ds.read(t=ti, bands=[b_idx])[:, :, 0] for ti in times], axis=0)
        data_vars[b_name] = xr.DataArray(
            cube,
            dims=["time", "y", "x"],
            coords={"time": time_coords, "y": y_coords, "x": x_coords},
        )

    xds = xr.Dataset(data_vars)

    # Attach CRS as attribute
    if ds.crs is not None:
        xds.attrs["crs"] = str(ds.crs)

    return xds
