# -*- coding: utf-8 -*-
"""
===============================================================================
传统GeoTIFF转COG

COG (Cloud Optimized GeoTIFF) 是一种优化设计的 GeoTIFF 文件格式, 特别适用于云环境下
的存储和访问. 自带分块存储功能与影像金字塔, 能够在请求时按需加载, 实现了高效的压缩和传输.

-------------------------------------------------------------------------------
Authors: CVEO Team
Last Updated: 2026-01-08
===============================================================================
"""

import os
import sys
import logging
from pathlib import Path
from osgeo import gdal
import numpy as np

# 添加父目录到 sys.path 以导入 utils
BASE_DIR = Path(__file__).parent.parent
sys.path.append(str(BASE_DIR))

from utils.common_utils import setup_logging

gdal.UseExceptions()
# 设置 GDAL 选项以优化性能
gdal.SetConfigOption("GDAL_NUM_THREADS", "ALL_CPUS")
gdal.SetConfigOption("GDAL_CACHEMAX", "1024")


def tif_to_cog(
    input_path: str,
    output_path: str,
    output_type: gdal.GDT_DataType = gdal.GDT_Float32,
    no_data: float = np.nan,
    compress: str = "DEFLATE",
    scaleParams: list = None,
):
    """
    将传统 GeoTIFF 转换为 COG 格式

    Parameters
    ----------
    input_path : str
        输入 GeoTIFF 文件路径
    output_path : str
        输出 COG 文件路径
    output_type : gdal.GDT_DataType, optional
        输出数据类型, 默认 float32, by default gdal.GDT_Float32
    no_data : float, optional
        无效值, 默认 np.nan, by default np.nan
    compress : str, optional
        压缩算法, 可选 "DEFLATE" 或 "LZW", by default "DEFLATE"
    """
    logging.info(f"输入文件: {input_path}")
    ds = gdal.Open(str(input_path))
    if ds is None:
        logging.error(f"无法打开输入文件: {input_path}")
        return
    # 查看原数据是否已设置无效值
    band = ds.GetRasterBand(1)
    original_no_data = band.GetNoDataValue()
    original_data_type = band.DataType
    if np.isnan(no_data):
        if original_no_data is None:
            # 根据输出数据类型设置无效值
            if output_type == gdal.GDT_Int16:
                no_data = -32768
            elif output_type == gdal.GDT_Float32:
                no_data = -9999
            else:
                no_data = np.nan
            logging.warning(f"原数据未设无效值, 现已设为: {no_data}")
        else:
            no_data = original_no_data
            logging.info(f"原数据无效值为: {no_data}")
    translate_options = gdal.TranslateOptions(
        format="COG",  # 输出为 COG 格式, 自动构建金字塔
        outputType=output_type,
        noData=no_data,
        creationOptions=[
            f"COMPRESS={compress}",
            "PREDICTOR=2",  # 差值预测, 利于影像压缩
            "NUM_THREADS=ALL_CPUS",
            "BIGTIFF=IF_SAFER",
            # "ZLEVEL=4",  # 压缩级别, 支持 1-9, 默认为 6, COG 不支持设置
            # "TILED=YES", # 分块参数, COG 格式自带分块, 不支持手动设置
            # "PHOTOMETRIC=RGB", # 可视化参数, COG 格式不支持设置
        ],
        callback=gdal.TermProgress_nocb,  # 进度回调, 不显示进度条
    )
    logging.info(f"开始转换为 COG 格式...")
    gdal.Translate(str(output_path), ds, options=translate_options)
    logging.info(f"结果已保存到: {output_path}")
    return


def main(input_dir, file_name, output_path):
    input_dir = Path(input_dir)
    output_dir = Path(output_path)
    os.makedirs(output_dir, exist_ok=True)
    # 配置日志记录
    log_file = output_dir / "toCOG.log"
    setup_logging(str(log_file))
    input_path = input_dir / file_name
    output_path = output_dir / file_name
    # 关于无效值
    # Float32 类型可以设为 -9999; 8bit 类型可以设为 0
    tif_to_cog(str(input_path), str(output_path), no_data=-9999)


if __name__ == "__main__":
    # 输入目录: 包含分块tif影像的根目录
    input_root = Path(r"D:\CVEOdata\RS_Data\Terrain")
    file_name = "GLO30.DSM.2014.Hubei.30m.tif"
    # 输出目录: 存放最终RGB镶嵌结果的目录
    output_root = Path(r"D:\CVEOdata\RS_Data\Terrain\COG")
    main(input_root, file_name, output_root)