"""
#!/usr/bin/python3.9
# -*- coding: utf-8 -*-
@Project :pythonCodeSnippet
@File    :downloadTilesTools.py
@IDE     :PyCharm
@Author  :chenxw
@Date    :2023/11/2 13:52
@Descr:
"""
import json
import os
import shutil
import uuid
from datetime import datetime
from math import floor

import cv2
import numpy as np
import requests as req
from PIL import Image
from loguru import logger
from vgis_gis.geoserverTools import GeoserverOperatoer

# 在线地图服务瓦片下载帮助类
class downloadTilesHelper:

    # 构造函数
    def __init__(self, tile_pic_suffix, merge_pic_name, resolution, tile_size, image_service_info, image_service_type,
                 tile_pic_path,
                 vcontat_img_path,
                 out_image_path):
        # 瓦片图片后缀
        self.tile_pic_suffix = tile_pic_suffix
        # 合成图片名称
        self.merge_pic_name = merge_pic_name
        # 分辨率，每像素对应的地图单位值，值越小，精度越高，合成图片越大
        self.resolution = resolution
        # 下载瓦片像素大小
        self.tile_size = tile_size
        # 要下载瓦片的影像服务信息
        self.image_service_info = image_service_info
        # 要下载瓦片的影像服务类型，如wms,wmts,tms,xyz等
        self.image_service_type = image_service_type
        # 对于非wms服务，需要重新计算分辨率和瓦片像素大小
        if image_service_type != "wms" and "geoserver" in image_service_info["service_type"]:
            self.resolution, self.tile_size, self.tile_coords = self.get_details_of_zoomlevel_in_geoserver(
                self.image_service_info["zoom_level"], self.image_service_info["gridset_name"])

        # 瓦片存储路径
        self.tile_pic_path = tile_pic_path if tile_pic_path is not None else os.path.join(os.getcwd(), ".cache",
                                                                                          "allTiles")
        # 每列瓦片拼接的竖条图片存放路径
        self.vcontat_img_path = vcontat_img_path if vcontat_img_path is not None else os.path.join(os.getcwd(),
                                                                                                   ".cache", "vMTiles")
        # 所有瓦片合成大图存放路径
        self.out_image_path = out_image_path
        self.retry_download_curent = 0
        self.retry_download_count = 5
        self.delete_cache()
        self.mkdir(self.tile_pic_path)
        self.mkdir(self.vcontat_img_path)
        self.mkdir(self.out_image_path)

    # 获取geoserver切片里指定分级对应的分辨率
    def get_details_of_zoomlevel_in_geoserver(self, zoom_level, grid_set_name):
        geoserverOperatoer = GeoserverOperatoer(self.image_service_info["service_url"].split("/gwc/")[0],
                                                self.image_service_info["geoserver_username"],
                                                self.image_service_info["geoserver_password"])
        gridset_info = geoserverOperatoer.get_grid_set_info(grid_set_name)
        gridset_json = json.loads(gridset_info)
        return gridset_json["gridSet"]["resolutions"][zoom_level], gridset_json["gridSet"]["tileWidth"], \
            gridset_json["gridSet"]["extent"]["coords"]

    # 根据坐标范围计算WMTS/TMS瓦片的行列号
    def compute_col_row_by_coords(self, xmin, xmax, ymin, ymax):
        tileMinCol, tileMaxCol, tileMinRow, tileMaxRow = None, None, None, None
        image_service_type = self.image_service_type
        # 为补偿浮点计算的不准确性
        epsilon = 1e-6
        # 左上角为原点
        if image_service_type == "wmts" or image_service_type == "xyz":
            if self.image_service_info["service_type"] == "geoserver_wmts":
                # 原点（左上角）
                tileMatrixMinx = self.tile_coords[0]
                tileMatrixMaxY = self.tile_coords[3]
                tileMinCol = round(floor(abs(xmin - tileMatrixMinx) / self.resolution / self.tile_size + epsilon))
                tileMaxCol = round(floor(abs(xmax - tileMatrixMinx) / self.resolution / self.tile_size - epsilon))
                tileMinRow = round(floor(abs(tileMatrixMaxY - ymax) / self.resolution / self.tile_size + epsilon))
                tileMaxRow = round(floor(abs((tileMatrixMaxY - ymin)) / self.resolution / self.tile_size - epsilon))
        # 左下角为原点
        elif image_service_type == "tms":
            if self.image_service_info["service_type"] == "geoserver_tms":
                tileMatrixMinx = self.tile_coords[0]
                tileMatrixMiny = self.tile_coords[1]
                tileMinCol = round(floor(abs(xmin - tileMatrixMinx) / self.resolution / self.tile_size + epsilon))
                tileMaxCol = round(floor(abs(xmax - tileMatrixMinx) / self.resolution / self.tile_size - epsilon))
                tileMinRow = round(floor(abs(ymax - tileMatrixMiny) / self.resolution / self.tile_size + epsilon))
                tileMaxRow = round(floor(abs(ymin - tileMatrixMiny) / self.resolution / self.tile_size - epsilon))
        matrixWidth = pow(2, self.image_service_info["zoom_level"] + 1)
        matrixHeight = pow(2, self.image_service_info["zoom_level"])
        if tileMinCol < 0:
            tileMinCol = 0
        if tileMaxCol >= matrixWidth:
            tileMaxCol = matrixWidth - 1
        if tileMinRow < 0:
            tileMinRow = 0
        if tileMaxRow >= matrixHeight:
            tileMaxRow = matrixHeight - 1
        return tileMinCol, tileMaxCol, tileMinRow, tileMaxRow

    # 下载wmts/tms/xyz瓦片
    def loop_download_tiles(self, xmin, xmax, ymin, ymax):
        tileMinCol, tileMaxCol, tileMinRow, tileMaxRow = self.compute_col_row_by_coords(xmin, xmax, ymin, ymax)
        if self.image_service_type == "wmts":
            for col in range(tileMinCol, tileMaxCol + 1):
                for row in range(tileMinRow, tileMaxRow + 1):
                    tile_url = self.build_wmts_tile_pic_url(col, row, self.image_service_info["zoom_level"])
                    # 下载瓦片
                    self.retry_download_curent = 0
                    self.download_each_tile(tile_url, "row=" + str(row) + "&col=" + str(col), col)
            # 同时计算左上角瓦片的地图单位
            # 瓦片原点在左上角
            left_top_tile_bounds = self.get_bounds_of_geoserver_wmts_tile(self.image_service_info["zoom_level"],
                                                                          tileMinRow, tileMinCol)

        elif self.image_service_type == "tms":
            pass
            # 同时计算左上角瓦片的地图单位
            # 瓦片原点在左下角
            left_top_tile_bounds = self.get_bounds_of_geoserver_wmts_tile(self.image_service_info["zoom_level"],
                                                                          tileMaxRow,
                                                                          tileMinCol)
        elif self.image_service_type == "xyz":
            pass
            # 同时计算左上角瓦片的地图单位
            # 瓦片原点在左上角
            left_top_tile_bounds = self.get_bounds_of_geoserver_wmts_tile(self.image_service_info["zoom_level"],
                                                                          tileMinRow,
                                                                          tileMinCol)
        return left_top_tile_bounds

    # 获取瓦片的地图坐标范围
    def get_bounds_of_geoserver_wmts_tile(self, level, row, col):
        tile_url = self.build_wmts_tile_pic_url(col, row, level)
        payload = {}
        headers = {}
        response = req.request("GET", tile_url, headers=headers, data=payload)
        headers = response.headers
        tile_bounds = headers["geowebcache-tile-bounds"]
        return tile_bounds.split(",")

    # 按照指定范围下载瓦片并拼接成大图
    def fetch_tiles(self, xmin, xmax, ymin, ymax):
        try:
            # 对影像服务的指定区域进行瓦片下载（针对下载不成功的进行重试），拼接成jpg/png，并生成jgw/pgw
            # 如果已经下载则跳过
            # 目录及文件命名形式：
            left_x = xmin
            top_y = ymax

            # 先判断缓存里有没有合成的图片数据，没有再去重新下载
            if not os.path.exists(os.path.join(self.out_image_path, self.merge_pic_name)):
                if self.image_service_type == "wms":
                    # 按照每个瓦片的大小作为步长进行下载
                    # 沿着X方向,逐步长滑动
                    while xmin < xmax:
                        self.loop_download_wms_tiles_follow_y(xmin, ymin, ymax)
                        xmin += (self.resolution * self.tile_size)
                elif self.image_service_type == "wmts" or self.image_service_type == "tms" or self.image_service_type == "xyz":
                    left_top_tile_bounds = self.loop_download_tiles(xmin, xmax, ymin, ymax)

                logger.info('瓦片下载完成')

                # 拼接瓦片图片
                logger.info('正在拼接瓦片为图片')
                try:
                    self.merge_tiles_to_image()
                    # 根据坐标范围反算瓦片的行列号，最后拼成的瓦片大图，比输入的坐标范围要大些，因为需要重新计算左上角点的地图坐标
                    if self.image_service_type != "wms":
                        left_x, top_y = float(left_top_tile_bounds[0]), float(left_top_tile_bounds[3])
                    # 生成world文件
                    self.build_merge_image_world_file_info(left_x, top_y)
                except Exception as tm_exp:
                    logger.error("瓦片拼接失败，请检查瓦片：{}".format(str(tm_exp)))
                    logger.error(tm_exp.__traceback__.tb_frame.f_globals["__file__"])  # 发生异常所在的文件
                    logger.error(tm_exp.__traceback__.tb_lineno)  # 发生异常所在的行数
            else:
                logger.warning("之前已拼接过图片")
            self.delete_cache()
        except Exception as exp:
            logger.error("下载瓦片失败：{}".format(str(exp)))
            logger.error(exp.__traceback__.tb_frame.f_globals["__file__"])  # 发生异常所在的文件
            logger.error(exp.__traceback__.tb_lineno)  # 发生异常所在的行数

    def delete_cache(self):
        # 清空下载的瓦片缓存（中间数据）
        if os.path.exists(self.tile_pic_path):
            shutil.rmtree(self.tile_pic_path)
        # 清空拼接的竖条瓦片图片（中间数据）
        if os.path.exists(self.vcontat_img_path):
            shutil.rmtree(self.vcontat_img_path)

    # 将瓦片合并成大图
    def merge_tiles_to_image(self):
        #  瓦片数据目录组织，是按照每个下载步长
        tiles_stride_dirs = os.listdir(self.tile_pic_path)
        # 先生成所有的竖条
        for tile_stride_dir in tiles_stride_dirs:
            self.vconcat_tileImage(tile_stride_dir)
        # 对所有竖条进行拼接成最终大图
        final, width, height = self.hconcat_vconcatImages(self.vcontat_img_path)
        logger.info('瓦片拼接完成')
        final.save(os.path.join(self.out_image_path, self.merge_pic_name))

    # 生成每一列的图片,即竖条
    def vconcat_tileImage(self, tile_stride_dir):
        tile_img_dir = os.path.join(self.tile_pic_path, tile_stride_dir)  # 原始文件目录
        names = os.listdir(tile_img_dir)
        #
        if self.image_service_type == "wms":
            # 竖着的瓦片要从上往下粘贴，对应的图片名称，Y的值越大越往上，所以做了反转
            # names.reverse()
            names = sorted(names, reverse=True)
            # sorted_names = sorted(names, key=self.custom_sort, reverse=True)
        elif self.image_service_type == "wmts":
            # wmts的原点在左下方，因此COL值越小越往上，不做反转
            names = sorted(names)

        images = []
        for name in names:
            img_path = os.path.join(tile_img_dir, name)
            # 拼接图片不稳定，如果图片中只有少量，就会报错
            # error: (-215:Assertion failed) src[i].dims <= 2 && src[i].cols == src[0].cols && src[i].type() == src[0].type() in function 'cv::vconcat'
            # 不完整的瓦片多了一个通道，需要删除，导致拼接的图片shape不一致,正确图片shape为(256,256,3) 错误瓦片shape为(256,256,4)
            # 修改by chenxw 将参数里的-1改成了1，有些不完整的瓦片通道变成了四个，需要这里只读前三个通道，和完整瓦片的三通道保持一致
            image = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 1)
            images.append(image)
        img = cv2.vconcat(images)
        compress_rate = 1
        height, width = img.shape[:2]
        # 双三次插值
        img_resize = cv2.resize(img, (int(width * compress_rate), int(height * compress_rate)),
                                interpolation=cv2.INTER_AREA)
        # 保存每列瓦片竖条合成图
        vconcat_pic_path = os.path.join(self.vcontat_img_path, '{}.{}'.format(tile_stride_dir, self.tile_pic_suffix))
        # 解决中文路径保存问题
        cv2.imencode('.{}'.format(self.tile_pic_suffix), img_resize)[1].tofile(vconcat_pic_path)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    # 对所有每一列竖条图片进行合并生成大图
    def hconcat_vconcatImages(self, img_dir):
        images = os.listdir(img_dir)
        # 每个竖条图片名称为x值，或者col，从小到大，依次对应大图的从左到右
        images = sorted(images)

        # 只需要图片,过滤掉目录里非图片
        def file_filter(f):
            if os.path.splitext(f)[1] in ['.jpg', '.jpeg', '.png', '.bmp']:
                return True
            else:
                return False

        images = list(filter(file_filter, images))
        img = Image.open(os.path.join(img_dir, images[0]))
        width = img.size[0] * len(images)
        height = img.size[1]
        final = Image.new('RGB', (width, height))

        for index, image in enumerate(images):
            img = Image.open(os.path.join(img_dir, image))
            final.paste(img, ((img.size[0] * index), 0))

        return final, width, height

    # 生成jgw/pgw --world file世界投影文件
    def build_merge_image_world_file_info(self, left_x, top_y):
        world_file_suffix = "jgw"
        if self.tile_pic_suffix == "jpg" or self.tile_pic_suffix == "jpeg":
            world_file_suffix = "jgw"
        elif self.tile_pic_suffix == "png":
            world_file_suffix = "pgw"
        file_pre_path, file_name = os.path.split(self.merge_pic_name)
        shotname, suffix = os.path.splitext(file_name)
        filename = os.path.join(self.out_image_path, "{}.{}".format(shotname, world_file_suffix))
        with open(filename, 'w') as file:
            # 在X方向的比例关系;每像素对应的经度值
            file.write("{}\n".format(self.resolution))
            # 影像数据的旋转参数，一般忽略，记录为0
            file.write("0.0000000000\n")
            # 影像数据的旋转参数，一般忽略，记录为0
            file.write("0.0000000000\n")
            # 在Y方向的比例关系;;每像素对应的纬度值，为负值, 这是由于空间坐标系与影像图片像素坐标系在Y方向上相反
            file.write("-{}\n".format(self.resolution))
            # 影像图片左上角点对应的空间坐标的X坐标
            file.write(str(left_x) + "\n")
            # 影像图片左上角点对应的空间坐标的Y坐标
            file.write(str(top_y) + "\n")

    # 沿着Y方向循环下载wms瓦片
    # 方向为从上往下
    def loop_download_wms_tiles_follow_y(self, xmin, ymin, ymax):
        # 沿着Y方向
        while ymax > ymin:
            bbox = str(xmin) + "%2C" + str(ymax - (self.resolution * self.tile_size)) + "%2C" + str(
                xmin + (self.resolution * self.tile_size)) + "%2C" + str(ymax)
            tile_url = self.build_wms_tile_pic_url(bbox)
            # 下载瓦片
            self.retry_download_curent = 0
            self.download_each_tile(tile_url, "y=" + str(ymax) + "&x=" + str(xmin), xmin)
            ymax -= (self.resolution * self.tile_size)

    # 获取wmts瓦片下载地址
    def build_wmts_tile_pic_url(self, col, row, level):
        url = None
        if self.image_service_info["service_type"] == "geoserver_wmts":
            url = "{}layer={}&style=&tilematrixset={}&Service=WMTS&Request=GetTile&Version=1.0.0&Format=image/{}&TileMatrix={}:{}&TileCol={}&TileRow={}".format(
                self.image_service_info["service_url"], self.image_service_info["layer_name"],
                self.image_service_info["gridset_name"], self.tile_pic_suffix, self.image_service_info["gridset_name"],
                level, col, row)
        else:
            pass
        return url

    # 构造下载的wms地图服务影像图片下载URL
    def build_wms_tile_pic_url(self, bbox):
        # geoserver发布的wms服务
        if self.image_service_info["service_type"] == "geoserver_wms":
            # url里的中文编码问题
            # layer_name = urllib.parse.quote(self.image_service_info["layer_name"].encode("gb2312"))
            data = {
                'LAYERS': self.image_service_info["layer_name"],
                'SERVICE': 'WMS',
                'REQUEST': 'GetMap',
                'SRS': 'EPSG:4326',
                'WIDTH': '{}'.format(self.tile_size),
                'HEIGHT': '{}'.format(self.tile_size),
                'STYLES': '',
                'BBOX': bbox,
                'FORMAT': "image/{}".format(self.tile_pic_suffix)
            }
            # 拼接url参数
            lt = []
            for k, v in data.items():
                lt.append(k + '=' + str(v))
            # 拼接url
            url = self.image_service_info["service_url"] + '&'.join(lt)
        # 卫星中心一张图2米WMS服务
        if self.image_service_info["service_type"] == "satellite_center_wms":
            data = {
                'SERVICE': 'WMS',
                'REQUEST': 'GetMap',
                # 'SRS': 'EPSG%3A4326',
                'SRS': 'WGS84',
                'ACCOUNT': 'zkzy',
                'PASSWD': 'fee528666a542aeb9f086093883c453d',
                'WIDTH': '{}'.format(self.tile_size),
                'HEIGHT': '{}'.format(self.tile_size),
                'FORMAT': 'image/{}'.format(self.tile_pic_suffix),
                'STYLES': '',
                'BBOX': bbox
            }
            # 拼接url参数
            lt = []
            for k, v in data.items():
                lt.append(k + '=' + str(v))
            # 拼接url
            url = self.image_service_info["service_url"] + "/OneMap" + self.image_service_info[
                "layer_name"] + "/wms?" + '&'.join(lt)

        # logger.info("瓦片地址：" + url)
        return url

    # 下载每个小瓦片
    def download_each_tile(self, url, pic_name, pic_dir_name):
        if self.retry_download_curent <= 5:
            try:
                if self.retry_download_curent > 0:
                    logger.info("self.retry_download_curent:{}".format(self.retry_download_curent))
                self.retry_download_curent += 1
                # 读取远程图片
                path = os.path.join(self.tile_pic_path, str(pic_dir_name))
                if not os.path.exists(path):
                    # 创建文件夹
                    self.mkdir(path)
                    logger.info("创建文件夹:" + path + "--" + datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                tilefile = os.path.join(path, '{}.{}'.format(pic_name, self.tile_pic_suffix))
                # # 如果瓦片存在，但大于100K，说明有可能是损坏的瓦片，需要重新下载-修改chenxw
                # if os.path.exists(tilefile):
                #     fileSize = os.path.getsize(tilefile) / 1024
                #     if fileSize > 100:
                #         os.remove(tilefile)
                # 如果瓦片图片不存在才下载
                if not os.path.exists(tilefile):
                    # url中有中文问题
                    # pic = urllib.request.urlopen(url).read()
                    pic = req.get(url)
                    # 创建瓦片图片
                    with open(tilefile, 'wb') as f:
                        f.write(pic.content)
                # 如果已下载瓦片不完整，需要重新下载
                else:
                    # print("检查文件夹下的瓦片:" + path)
                    try:
                        img_cv = cv2.imdecode(np.fromfile(tilefile, dtype=np.uint8), -1)
                        if img_cv is not None:
                            if img_cv.shape[2] == 4:
                                logger.info(
                                    "下载瓦片不完整，需要重新下载：" + tilefile + "--" + datetime.now().strftime(
                                        '%Y-%m-%d %H:%M:%S'))
                                os.remove(tilefile)
                                self.download_each_tile(url, pic_name, pic_dir_name)
                            else:
                                pass
                                # logger.info("瓦片已存在且完整不再下载---" + datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                    except BaseException as ex:
                        logger.error(
                            "读取瓦片文件失败---" + tilefile + "---" + datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                        os.remove(tilefile)
                        # 瓦片下载失败，重试
                        self.download_each_tile(url, pic_name, pic_dir_name)
            except BaseException as ex:
                logger.error("下载失败" + str(ex) + "---" + datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                logger.error(url)
                if str(ex) != "HTTP Error 404: Not Found":
                    logger.error("正在重新下载---" + datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                    # 瓦片下载失败，重试
                    self.download_each_tile(url, pic_name, pic_dir_name)

    # 创建文件夹
    def mkdir(self, path):
        folder = os.path.exists(path)
        # 判断是否存在文件夹如果不存在则创建为文件夹
        if not folder:
            # 创建文件时如果路径不存在会创建这个路径
            os.makedirs(path)

    # 删除文件夹内文件
    def rmdir(self, img_dir):
        images = os.listdir(img_dir)
        for name in images:
            img_path = os.path.join(img_dir, name)
            os.remove(img_path)