from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np


@dataclass
class RasterData:
    array: np.ndarray  # shape: (bands, height, width)
    band_names: list[str]
    crs: Any | None
    transform: Any | None
    profile: dict[str, Any]
    nodata: float | None
    source: str
    sensor: str | None = None


def load_tabular(path: Path):
    import pandas as pd

    if path.suffix.lower() in {".csv"}:
        return pd.read_csv(path)
    if path.suffix.lower() in {".xlsx", ".xls"}:
        return pd.read_excel(path)
    raise ValueError("Unsupported tabular format. Use CSV or XLSX.")


def load_npy(path: Path) -> RasterData:
    import numpy as np

    array = np.load(path, allow_pickle=False)
    if array.ndim == 2:
        array = array[None, :, :]
    elif array.ndim == 3:
        # Heuristic: prefer (bands, height, width)
        if array.shape[0] <= 20 and array.shape[0] < array.shape[1] and array.shape[0] < array.shape[2]:
            pass
        elif array.shape[-1] <= 20:
            array = np.moveaxis(array, -1, 0)
        else:
            # fallback: assume bands first
            pass
    else:
        raise ValueError("Unsupported NPY shape. Expect 2D or 3D array.")

    band_names = [f"band_{idx+1}" for idx in range(array.shape[0])]
    return RasterData(
        array=array.astype("float32"),
        band_names=band_names,
        crs=None,
        transform=None,
        profile={},
        nodata=None,
        source=str(path),
        sensor=None,
    )

def load_geotiff(path: Path) -> RasterData:
    import rasterio

    with rasterio.open(path) as src:
        array = src.read().astype("float32")
        band_names = [f"band_{idx+1}" for idx in range(array.shape[0])]
        return RasterData(
            array=array,
            band_names=band_names,
            crs=src.crs,
            transform=src.transform,
            profile=src.profile,
            nodata=src.nodata,
            source=str(path),
            sensor=None,
        )


def load_netcdf(path: Path, variable: str | None = None) -> RasterData:
    import xarray as xr

    ds = xr.open_dataset(path)
    data_vars = list(ds.data_vars)
    if not data_vars:
        raise ValueError("No data variables found in NetCDF.")

    arrays = []
    band_names = []

    def _reduce_to_2d(values: np.ndarray) -> np.ndarray:
        arr = values
        while arr.ndim > 2:
            arr = arr[0]
        if arr.ndim != 2:
            raise ValueError("Could not reduce NetCDF variable to 2D.")
        return arr

    if variable is None:
        for var in data_vars:
            data = ds[var]
            values = _reduce_to_2d(data.values)
            arrays.append(values)
            band_names.append(var)
    else:
        data = ds[variable]
        values = data.values
        if values.ndim == 2:
            arrays.append(values)
            band_names.append(variable)
        elif values.ndim >= 3:
            reduced = _reduce_to_2d(values)
            arrays.append(reduced)
            band_names.append(variable)
        else:
            raise ValueError("Unsupported NetCDF variable dimensionality.")

    array = np.stack(arrays, axis=0)
    crs = getattr(data, "rio", None).crs if hasattr(data, "rio") else None
    transform = getattr(data, "rio", None).transform() if hasattr(data, "rio") else None

    return RasterData(
        array=array.astype("float32"),
        band_names=band_names,
        crs=crs,
        transform=transform,
        profile={},
        nodata=None,
        source=str(path),
        sensor="ERA5",
    )


def load_envi(hdr_path: Path) -> RasterData:
    from spectral import open_image

    img = open_image(str(hdr_path))
    array = img.load().astype("float32")
    if array.ndim == 3:
        array = np.moveaxis(array, -1, 0)
    band_names = [f"band_{idx+1}" for idx in range(array.shape[0])]
    return RasterData(
        array=array,
        band_names=band_names,
        crs=None,
        transform=None,
        profile={},
        nodata=None,
        source=str(hdr_path),
        sensor="ENMAP",
    )


def scan_sentinel2_folder(folder: Path) -> dict[str, Path]:
    band_files = {}
    for path in folder.rglob("*.jp2"):
        name = path.name
        if "_B" in name:
            parts = name.split("_B")
            if len(parts) > 1:
                band_code = parts[1].split("_")[0]
                band_files[f"B{band_code}"] = path
    return band_files


def scan_sentinel2_tiffs(folder: Path) -> dict[str, Path]:
    band_files = {}
    for path in folder.rglob("*.tif*"):
        name = path.name
        if "_B" in name:
            parts = name.split("_B")
            if len(parts) > 1:
                band_code = parts[1].split("_")[0].replace("A", "A")
                band_files[f"B{band_code}"] = path
    return band_files


def load_sentinel2_band_tiffs(band_paths: dict[str, Path]) -> RasterData:
    import rasterio
    from rasterio.enums import Resampling

    if not band_paths:
        raise ValueError("No Sentinel-2 band files provided.")

    # choose base with largest pixel count
    def _pixel_count(p: Path) -> int:
        with rasterio.open(p) as ds:
            return ds.width * ds.height

    sorted_bands = sorted(band_paths.items())
    base_band, base_path = max(sorted_bands, key=lambda item: _pixel_count(item[1]))

    arrays = []
    names = []
    crs = None
    transform = None
    profile = {}
    nodata = None

    with rasterio.open(base_path) as base_ds:
        base_h, base_w = base_ds.height, base_ds.width
        crs = base_ds.crs
        transform = base_ds.transform
        profile = base_ds.profile
        nodata = base_ds.nodata

        for band_name, path in sorted_bands:
            with rasterio.open(path) as src:
                data = src.read(
                    1,
                    out_shape=(base_h, base_w),
                    resampling=Resampling.bilinear,
                ).astype("float32")
                arrays.append(data)
                names.append(band_name)

    array = np.stack(arrays, axis=0)
    return RasterData(
        array=array,
        band_names=names,
        crs=crs,
        transform=transform,
        profile=profile,
        nodata=nodata,
        source=str(base_path.parent),
        sensor="Sentinel-2",
    )


def _extract_s2_band_name(path: Path) -> str:
    name = path.stem
    if "_B" in name:
        parts = name.split("_B")
        if len(parts) > 1:
            band = parts[1].split("_")[0]
            return f"B{band}"
    return name


def load_sentinel2_bands(band_paths: list[Path]) -> RasterData:
    import rasterio

    arrays = []
    names = []
    crs = None
    transform = None
    profile = {}
    nodata = None

    for path in band_paths:
        with rasterio.open(path) as src:
            arrays.append(src.read(1).astype("float32"))
            names.append(_extract_s2_band_name(path))
            crs = src.crs
            transform = src.transform
            profile = src.profile
            nodata = src.nodata

    array = np.stack(arrays, axis=0)
    return RasterData(
        array=array,
        band_names=names,
        crs=crs,
        transform=transform,
        profile=profile,
        nodata=nodata,
        source=str(band_paths[0].parent),
        sensor="Sentinel-2",
    )


def save_geotiff(path: Path, array: np.ndarray, reference: RasterData) -> None:
    import rasterio

    bands, height, width = array.shape
    profile = reference.profile.copy()
    profile.update(
        {
            "count": bands,
            "height": height,
            "width": width,
            "dtype": array.dtype,
            "transform": reference.transform,
            "crs": reference.crs,
        }
    )
    path.parent.mkdir(parents=True, exist_ok=True)
    with rasterio.open(path, "w", **profile) as dst:
        dst.write(array)
