# -*- coding: utf-8 -*-
import os
import re
import tempfile
import psycopg2
import shutil
from zipfile import ZipFile, ZIP_DEFLATED, BadZipfile
import zlib

# dst = r'Y:\SAUDI\SHP_EXPORT'

# for dirpath, dirnames, filenames in os.walk(r'Y:\SAUDI'):
#     right_filenames = [os.path.join(dirpath, fnm) for fnm in filenames if os.path.splitext(fnm)[1].lower().endswith('.zip')]
#     for zip_filename in right_filenames:
#         print(zip_filename)
#         with zipfile.ZipFile(zip_filename) as z:
#             for zfile in z.namelist():
#                 if 'pixel_shape' in zfile.lower():
#                     # z.extract(file, r'Y:\SAUDI\SHP_EXPORT')
#                     source = z.open(zfile)
#                     target = open(os.path.join(dst, os.path.basename(zfile)), 'wb')
#                     with source, target:
#                         try:
#                             shutil.copyfileobj(source, target)
#                         except (zipfile.BadZipfile, zlib.error):
#                             continue
# target_filename = '053326823050_01_P001_PSH'
# target_dirname = r'Y:\SAUDI\SHP_EXPORT'
#
# db_zip_path = r"Y:\SAUDI\053326819080_01.zip"
# db_img_id = '053326819080_01_P001'


def extract_zip_image(src_zip_path, img_id, dst_path):
    with ZipFile(src_zip_path) as z:
                    for zitem in z.namelist():
                        regexp = '.+' + img_id + '.{4}$'
                        if re.match(
                                regexp,
                                os.path.dirname(zitem), re.IGNORECASE) is not None and not zitem.endswith('/'):
                            print(zitem)
                            source = z.open(zitem)
                            image_dirname = os.path.join(dst_path, os.path.basename(os.path.dirname(zitem)))
                            if not os.path.exists(image_dirname):
                                os.mkdir(image_dirname)
                            target = open(os.path.join(dst_path, image_dirname, os.path.basename(zitem)), 'wb')
                            with source, target:
                                try:
                                    shutil.copyfileobj(source, target)
                                except (BadZipfile, zlib.error):
                                    print('##'*80, zitem)
                                    continue

                    zip_path = image_dirname + '.zip'
                    os.chdir(os.path.dirname(image_dirname))
                    with ZipFile(zip_path,
                                 "w",
                                 compression=ZIP_DEFLATED, allowZip64=True) as zf:
                        for filename in os.listdir(image_dirname):
                            # отбираем для архивации стандартный комплект файлов по маске
                            if re.match(
                                    r'^\d{2}\D{3}\d{8}-.{4}(|_.{4,})-\d{12}_\d{2}_.{4}(|_README|-BROWSE)\.(IMD|RPB|TIF|TIL|XML|TXT|JPG)$',
                                    filename, re.IGNORECASE) is not None:
                                # if not name.lower().endswith(('.zip', '.rar', '052563231040_01_order_shape.shp')):
                                name = os.path.join(image_dirname, filename)
                                zf.write(name, os.path.basename(name))
                    shutil.rmtree(image_dirname)
    return zip_path


conn = psycopg2.connect(
    database='innoter', user='postgres', password='postgres', host='localhost', port='5432'
)
# TODO добавить проверку на инкремент и наличие снимка в БД
cursor = conn.cursor()
# regexp_replace(selection.old_path, '[^\\\\]+$', '')
cursor.execute("""SELECT old_path, substring(filename, length(filename) - 19), innoter_id
FROM geoarchive.imagery
WHERE comment LIKE '%zip%' AND zip_path is NULL""")
image_info_list = cursor.fetchall()
counter = 0
for item in image_info_list:
    counter += 1
    db_zip_path, db_img_id, inn_id = item[0], item[1], item[2]

    downloading_flag_file = os.path.join(r'Y:\!Archive_НЕ_ТРОГАТЬ\БД', inn_id + '.TEMP')
    if os.path.exists(downloading_flag_file):
        print('{} архивируется другим процессом, пропускаем'.format(inn_id))
        continue
    done_flag_file = os.path.join(r'Y:\!Archive_НЕ_ТРОГАТЬ\БД', inn_id + '.DONE')
    if os.path.exists(done_flag_file):
        print('{} готов, пропускаем'.format(inn_id))
        continue
    else:
        open(downloading_flag_file, 'a').close()

        print('{}/{} Работаем с {}'.format(counter, len(image_info_list), db_zip_path))
        dst_archive = extract_zip_image(db_zip_path, db_img_id, r'Y:\!Archive_НЕ_ТРОГАТЬ\БД')
        cursor.execute("""
    UPDATE geoarchive.imagery
    SET zip_path = %s
    WHERE innoter_id LIKE %s""", (dst_archive.replace('Y:\\', '\\\\nas1\\storage\\'), inn_id,))
        conn.commit()
        os.remove(downloading_flag_file)
        open(done_flag_file, 'a').close()

