# -*- coding: utf-8 -*-
"""
掩膜转换工具模块

提供 shapefile / ndarray / dat / tif 等多种格式掩膜之间的相互转换，
以及水体掩膜的预处理逻辑。
"""
import os
from pathlib import Path
from typing import Optional, Union

import numpy as np

try:
    from osgeo import gdal, ogr
    GDAL_AVAILABLE = True
except ImportError:
    GDAL_AVAILABLE = False


def prepare_water_mask_for_algorithm(
    water_mask: Optional[Union[str, np.ndarray]],
    image_shape: Union[tuple, np.ndarray],
    geotransform: tuple,
    projection: str,
    img_path: str,
    water_mask_dir: Optional[str] = None,
    callback=None
) -> Optional[np.ndarray]:
    """
    准备水域掩膜供算法使用

    支持格式：
    - None：自动使用预先生成的 dat 格式掩膜
    - numpy.ndarray：直接返回（确保是 0/1 格式）
    - .dat / .tif 等栅格文件：读取并返回
    - .shp 文件：先栅格化，再读取返回

    Args:
        water_mask: 掩膜来源
        image_shape: 影像形状 (height, width) 或 (height, width, channels)
        geotransform: GDAL 地理变换参数
        projection: 投影信息
        img_path: 影像路径（用于 shp 栅格化）
        water_mask_dir: 水体掩膜目录（用于缓存栅格化的 shp 结果）
        callback: 进度回调函数（可选）

    Returns:
        numpy数组（dtype=uint8，0=非水域，1=水域）或 None
    """
    img_height, img_width = image_shape[0], image_shape[1]

    if water_mask is None:
        return None

    # numpy 数组直接返回
    if isinstance(water_mask, np.ndarray):
        if water_mask.shape[:2] != (img_height, img_width):
            raise ValueError(f"掩膜尺寸 {water_mask.shape[:2]} 与图像尺寸 {(img_height, img_width)} 不匹配")
        return (water_mask > 0).astype(np.uint8)

    # 字符串路径
    if isinstance(water_mask, str):
        ext = Path(water_mask).suffix.lower()

        # shapefile 格式
        if ext == '.shp':
            return _convert_shp_to_mask(
                shp_path=water_mask,
                img_path=img_path,
                image_shape=image_shape,
                geotransform=geotransform,
                projection=projection,
                water_mask_dir=water_mask_dir,
                callback=callback
            )

        # 栅格文件格式
        return _load_raster_mask(water_mask, img_height, img_width)

    raise ValueError(f"不支持的掩膜类型: {type(water_mask)}")


def _convert_shp_to_mask(shp_path: str, img_path: str,
                          image_shape: tuple,
                          geotransform: tuple,
                          projection: str,
                          water_mask_dir: Optional[str] = None,
                          callback=None) -> np.ndarray:
    """将 shapefile 栅格化为掩膜数组"""
    from src.utils.extract_water_area import rasterize_shp

    safe_shp_path = os.path.abspath(shp_path).replace('\\', '/')
    shp_name = Path(safe_shp_path).stem

    if water_mask_dir:
        temp_mask_path = str(Path(water_mask_dir) / f"water_mask_{shp_name}.dat")
    else:
        temp_mask_path = f"/tmp/water_mask_{shp_name}.dat"

    # 缓存：已栅格化则直接读取
    if Path(temp_mask_path).exists():
        print(f"使用已存在的栅格化掩膜: {temp_mask_path}")
        return _load_raster_mask(temp_mask_path, image_shape[0], image_shape[1])

    # 需要栅格化
    if img_path is None:
        raise ValueError("当 water_mask 为 shp 格式时，需要提供 img_path 参数用于栅格化")

    print(f"正在将 SHP 栅格化: {safe_shp_path}")
    rasterize_shp(safe_shp_path, temp_mask_path, img_path)

    return _load_raster_mask(temp_mask_path, image_shape[0], image_shape[1])


def _load_raster_mask(mask_path: str, img_height: int, img_width: int) -> np.ndarray:
    """从栅格文件加载掩膜"""
    if not GDAL_AVAILABLE:
        raise ImportError("GDAL未安装，无法读取掩膜文件")

    mask_dataset = gdal.Open(mask_path, gdal.GA_ReadOnly)
    if mask_dataset is None:
        raise ValueError(f"无法打开掩膜文件: {mask_path}")

    try:
        mask_array = mask_dataset.GetRasterBand(1).ReadAsArray()
    finally:
        mask_dataset = None

    if mask_array.shape != (img_height, img_width):
        raise ValueError(f"掩膜尺寸 {mask_array.shape} 与图像尺寸 {(img_height, img_width)} 不匹配")

    return (mask_array > 0).astype(np.uint8)


def ensure_water_mask_dat(img_path: str,
                           existing_dat_path: Optional[str] = None,
                           output_dir: Optional[str] = None) -> str:
    """
    确保存在 dat 格式的水体掩膜文件（用于步骤3/4中的算法）

    如果 existing_dat_path 存在且是 .dat 文件，直接返回。
    如果存在同名 .dat 文件，直接返回。
    否则从 img_path 生成并保存到 output_dir。

    Args:
        img_path: 用于生成掩膜的遥感影像路径
        existing_dat_path: 已有的 dat 格式掩膜路径（可选）
        output_dir: 输出目录（可选）

    Returns:
        dat 格式掩膜文件路径
    """
    if existing_dat_path and Path(existing_dat_path).suffix.lower() == '.dat':
        if Path(existing_dat_path).exists():
            return existing_dat_path

    img_name = Path(img_path).stem
    if output_dir is None:
        output_dir = str(Path(img_path).parent)

    dat_path = str(Path(output_dir) / f"{img_name}_water_mask.dat")

    if Path(dat_path).exists():
        return dat_path

    # 如果已有其他格式的掩膜，转换为 dat
    for ext in ['.tif', '.img', '.tiff']:
        alt_path = str(Path(output_dir) / f"{img_name}_water_mask{ext}")
        if Path(alt_path).exists():
            return _convert_to_dat(alt_path, dat_path)

    return dat_path  # 返回目标路径，让调用方决定是否需要生成


def _convert_to_dat(src_path: str, dest_path: str) -> str:
    """将其他栅格格式转换为 ENVI dat 格式"""
    if not GDAL_AVAILABLE:
        raise ImportError("GDAL未安装，无法转换格式")

    src_ds = gdal.Open(src_path, gdal.GA_ReadOnly)
    if src_ds is None:
        raise ValueError(f"无法打开源掩膜文件: {src_path}")

    try:
        geotransform = src_ds.GetGeoTransform()
        projection = src_ds.GetProjection()
        band = src_ds.GetRasterBand(1)
        array = band.ReadAsArray()

        driver = gdal.GetDriverByName('ENVI')
        if driver is None:
            driver = gdal.GetDriverByName('GTiff')

        dest_ds = driver.Create(dest_path, src_ds.RasterXSize, src_ds.RasterYSize, 1, gdal.GDT_Byte)
        if dest_ds is None:
            raise ValueError(f"无法创建输出文件: {dest_path}")

        try:
            dest_ds.SetGeoTransform(geotransform)
            dest_ds.SetProjection(projection)
            dest_band = dest_ds.GetRasterBand(1)
            dest_band.WriteArray((array > 0).astype(np.uint8))
            dest_band.FlushCache()
        finally:
            dest_ds = None

        return dest_path
    finally:
        src_ds = None