import os
import numpy as np
import rasterio
import spectral as sp
from osgeo import gdal
from rasmas.converter import convert_dat_to_tiff  # Convert `.dat` to `.tif` if needed

def save_masked_image_envi(masked_image, original_hdr_path, output_base_path):
    """
    Saves a masked hyperspectral image in ENVI format.

    Parameters:
        masked_image (numpy array): The masked hyperspectral image.
        original_hdr_path (str): Path to the original `.hdr` file.
        output_base_path (str): Base path (without extension) for saving.
    """
    original_image = sp.open_image(original_hdr_path)
    metadata = original_image.metadata

    # Update metadata dimensions
    metadata['lines'] = str(masked_image.shape[0])
    metadata['samples'] = str(masked_image.shape[1])
    metadata['bands'] = str(masked_image.shape[2])

    # Convert masked image to original dtype
    masked_image = masked_image.astype(original_image.dtype)

    # Save the masked image in ENVI format
    hdr_path = output_base_path + '.hdr'
    sp.envi.save_image(hdr_path, masked_image, dtype=original_image.dtype, metadata=metadata, interleave='bil')

    # Rename `.img` file to `.dat`
    img_path = output_base_path + '.img'
    dat_path = output_base_path + '.dat'
    if os.path.exists(img_path):
        os.rename(img_path, dat_path)

    print(f"[INFO] Masked image saved to {dat_path} and {hdr_path}")


def apply_soil_mask(input_file, green_band_index, red_band_index, blue_band_index, nir_band_index, output_dir=None, save_as_envi=False):
    """
    Applies a soil mask to a hyperspectral/multispectral image.

    Masking is based on:
    - Green > Red
    - Green > Blue
    - NIR > Green

    Parameters:
        input_file (str): Path to the input hyperspectral/multispectral image (.dat or .tif).
        green_band_index (int): Index of the Green band.
        red_band_index (int): Index of the Red band.
        blue_band_index (int): Index of the Blue band.
        nir_band_index (int): Index of the NIR band.
        output_dir (str, optional): Directory to save the masked image. If None, only returns the masked image.
        save_as_envi (bool, optional): If True, saves the masked image as ENVI (`.dat` + `.hdr`).

    Returns:
        numpy array: Masked hyperspectral image.
    """

    # ✅ Convert `.dat` to `.tif` if necessary
    is_dat_file = input_file.lower().endswith(".dat")
    if is_dat_file:
        print(f"[INFO] Converting {input_file} to TIFF format...")
        tiff_file = convert_dat_to_tiff(input_file)
    else:
        tiff_file = input_file

    # ✅ Read the image using Rasterio (for `.tif`)
    with rasterio.open(tiff_file) as src:
        img_array = src.read()  # Shape: (bands, height, width)
        profile = src.profile  # Save metadata for writing output

    # ✅ Extract the specified bands
    green_band = img_array[green_band_index, :, :]
    red_band = img_array[red_band_index, :, :]
    blue_band = img_array[blue_band_index, :, :]
    nir_band = img_array[nir_band_index, :, :]

    # ✅ Apply the soil mask
    mask = np.logical_and.reduce((green_band > red_band, green_band > blue_band, nir_band > green_band)).astype(np.uint8)

    # ✅ Apply mask to the entire hyperspectral image (Set non-vegetation areas to 0)
    masked_image = img_array * mask  

    # ✅ Save the masked image if an output directory is provided
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
        base_filename = os.path.basename(input_file).replace(".tif", "").replace(".dat", "")

        # ✅ Save as TIFF if `save_as_envi` is False
        if not save_as_envi:
            mask_output_path = os.path.join(output_dir, f"{base_filename}_masked.tif")
            with rasterio.open(mask_output_path, 'w', **profile) as dst:
                dst.write(masked_image)
            print(f"[INFO] Masked image saved at: {mask_output_path}")

        # ✅ Save as ENVI if `save_as_envi=True`
        else:
            hdr_path = input_file.replace(".dat", ".hdr") if is_dat_file else None
            envi_output_base = os.path.join(output_dir, f"{base_filename}_masked")
            save_masked_image_envi(masked_image, hdr_path, envi_output_base)

    return masked_image  # ✅ Now returning the **masked image**, not just the mask


def apply_shadow_mask(input_file, red_band_index, nir_band_index, threshold=0.2, output_dir=None):
    """
    Applies a shadow mask to a hyperspectral/multispectral image using NDVI.

    Parameters:
        input_file (str): Path to the input image (.tif).
        red_band_index (int): Index of the Red band.
        nir_band_index (int): Index of the NIR band.
        threshold (float, optional): NDVI threshold for shadow masking (default is 0.2).
        output_dir (str, optional): If provided, saves the shadow mask to this directory.

    Returns:
        numpy array: Binary shadow mask (1 = shadow, 0 = non-shadow).
    """
    # ✅ Open the image using Rasterio
    with rasterio.open(input_file) as src:
        img_array = src.read()  # Shape: (bands, height, width)
        profile = src.profile  # Save metadata

    # ✅ Extract the Red and NIR bands
    red_band = img_array[red_band_index, :, :].astype(np.float32)
    nir_band = img_array[nir_band_index, :, :].astype(np.float32)

    # ✅ Compute NDVI
    ndvi = (nir_band - red_band) / (nir_band + red_band + 1e-8)  # Avoid division by zero

    # ✅ Create shadow mask (1 = shadow, 0 = non-shadow)
    shadow_mask = (ndvi < threshold).astype(np.uint8)

    # ✅ Save the shadow mask if output_dir is provided
    if output_dir:
        import os
        os.makedirs(output_dir, exist_ok=True)
        mask_output_path = os.path.join(output_dir, f"{os.path.basename(input_file).replace('.tif', '_shadow_mask.tif')}")

        with rasterio.open(mask_output_path, 'w', **profile) as dst:
            dst.write(shadow_mask, 1)  # Save mask as a single-band image

        print(f"[INFO] Shadow mask saved at: {mask_output_path}")

    return shadow_mask  # ✅ Return the binary shadow mask
