import binascii
import json
import os
import random
import re
import sys
import time
from PIL import Image
from type.PNG import PNG
# from type.JPG import JPG
# from type.GIF import GIF

types = [PNG] # 暂时先只处理PNG，其他的都没写好呢[PNG, JPG, GIF]


def check_image_type(filename: str, data: bytearray):
    file_type1 = None
    for model in types:
        if re.match(f'.*\.{model.extension}$', filename.lower()):
            file_type1 = model
            break

    file_type2 = None
    weight = 0
    for model in types:
        weight2 = model.test(data)
        if weight2 > weight:
            file_type2 = model
            weight = weight2
    
    if file_type1 != None:
        if file_type1 == file_type2:
            print(f'看起来文件格式是{file_type1.extension}')
        else:
            print(f'文件名中格式是{file_type1.extension}，但是解析文件头格式发现{weight}%是{file_type2.extension}')
    else:
        if file_type2 != None:
            file_type1 = file_type2
            print(f'检测文件头认为文件格式{weight}%是{file_type2.extension}')
        else:
            print('未检测到任何图片格式特征，该文件可能不是图片文件')

    for i in range(len(types)):
        print(f'    {i}、{types[i].extension}')
    index = input('要指定文件格式吗？或者直接回车。请输入文件格式序号：')
    if index.isdigit():
        index = int(index)
        file_type1 = types[index]

    return file_type1



# 总函数
def main(filename: str):
    # 读取图片数据
    data = None
    with open(filename, 'rb') as f:
        data = bytearray(f.read())
        
    # 确定文件格式
    model = check_image_type(filename, data)

    # 检测并修复
    changed = False
    if model != None:
        info = model.get_info(data)
        # print(json.dumps(info))
        changed1, data, tail = model.check_tail(data)
        changed2, data, head = model.repair_sig(data)
        changed3, data, width, height = model.repair_size(data)
        changed4, data, data_append = model.check_data(data)
        checked_channel = False # 是否检测过RGB通道，如果没检测过，则下面还会尝试检测一遍
        try:
            # 尝试分析通道信息
            changed5, data_r, data_g, data_b = model.check_channels(filename, width, height)
            checked_channel = True
        except Exception as e:
            pass
            #print('[x] 尝试打开图片失败，失败原因：' + str(e))

        changed = changed1 or changed2 or changed3
    else:
        print('[*] 未进行任何处理')

    if changed:
        # 创建结果目录
        basename = os.path.basename(filename)
        dirname = filename.replace('.', '_') + str(int(time.time()))
        dirpath = os.path.join(os.path.dirname(os.path.abspath(filename)), dirname)
        os.mkdir(dirpath)

        # 保存解析结果
        imagepath = os.path.join(dirpath, basename)
        f = open(imagepath, 'wb')
        f.write(data)
        f.close()
        if changed1 and tail:
            f = open(os.path.join(dirpath, 'tail'), 'wb')
            f.write(tail)
            f.close()
        if changed2 and head:
            f = open(os.path.join(dirpath, 'head'), 'wb')
            f.write(head)
            f.close()
        for i in range(len(data_append)):
            f = open(os.path.join(dirpath, 'data_append_' + str(i)), 'wb')
            f.write(data_append[i])
            f.close()
        print(f'[*] 分析结果已保存到 {dirpath}')
        if not checked_channel:
            # 尝试分析通道信息
            changed5, data_r, data_g, data_b = model.check_channels(imagepath, width, height)
        try:
            # 尝试打开图片看一看
            im = Image.open(imagepath)
            im.show()
        except Exception as e:
            print('[x] 尝试打开图片失败，失败原因：' + str(e))
    else:
        print('[*] 未进行任何处理')

    print('[*] 分析完成！')

# 开始执行
if len(sys.argv) > 1:
    filename = sys.argv[1]

main(filename)