import os
import json
import random
from Lock import PasswordLock


def is_prime(n):
    """判断一个数字是否是素数"""
    if n < 2:
        return False
    for i in range(2, int(n ** 0.5) + 1):
        if n % i == 0:
            return False
    return True


def satisfies_prime_unique_condition(digits):
    """检查是否满足[-1, -1]条件：每位密码为素数且不重复"""
    return all(is_prime(d) for d in digits) and len(set(digits)) == 3


def crack_method1(conditions, stored_hash):
    """从高位到低位回溯（第一位→第二位→第三位）"""
    tries = 1
    found_password = None
    found = False
    lock = PasswordLock()

    # 根据线索生成每位的合法数字范围
    digit_options = [[] for _ in range(3)]

    for i in range(3):
        for d in range(10):
            valid = True
            for cond in conditions:
                if cond == [-1, -1]:
                    if not is_prime(d):
                        valid = False
                        break
                elif len(cond) == 2:
                    a, t = cond
                    a = abs(a)
                    if a - 1 == i and ((t == 0 and d % 2 != 0) or (t == 1 and d % 2 != 1)):
                        valid = False
                        break
                elif len(cond) == 3:
                    if cond[i] != -1 and d != cond[i]:
                        valid = False
                        break
            if valid:
                digit_options[i].append(d)
    
    for i in range(3):
        random.shuffle(digit_options[i])

    def backtrack(index, current_digits):
        nonlocal tries, found, found_password
        if found:
            return

        if index == 3:
            password = "".join(map(str, current_digits))
            calculated_hash = lock.hash_password(password)

            
            if calculated_hash == stored_hash:
                found_password = password
                found = True
                return

            tries += 1
            return

        for digit in digit_options[index]:
            current_digits.append(digit)
            backtrack(index + 1, current_digits)
            current_digits.pop()
            if found:
                return

    backtrack(0, [])
    return [found_password] if found_password else [], tries


def crack_method2(conditions, stored_hash):
    """按条件优先级排序的回溯方法：先处理约束最多的位置"""
    tries = 1
    found_password = None
    found = False
    lock = PasswordLock()

    digit_options = [[] for _ in range(3)]
    for i in range(3):
        for d in range(10):
            valid = True
            for cond in conditions:
                if cond == [-1, -1]:
                    if not is_prime(d):
                        valid = False
                        break
                elif len(cond) == 2:
                    a, t = cond
                    a = abs(a)
                    if a - 1 == i and ((t == 0 and d % 2 != 0) or (t == 1 and d % 2 != 1)):
                        valid = False
                        break
                elif len(cond) == 3:
                    if cond[i] != -1 and d != cond[i]:
                        valid = False
                        break
            if valid:
                digit_options[i].append(d)
                
    for i in range(3):
        random.shuffle(digit_options[i])

    sorted_indices = sorted(range(3), key=lambda x: len(digit_options[x]))

    def backtrack(index, current_digits):
        nonlocal tries, found, found_password
        if found:
            return

        if index == 3:
            reordered = [0, 0, 0]
            for i, pos in enumerate(sorted_indices):
                reordered[pos] = current_digits[i]
            password = "".join(map(str, reordered))
            calculated_hash = lock.hash_password(password)

            if calculated_hash == stored_hash:
                found_password = password
                found = True
                return

            tries += 1
            return

        current_index = sorted_indices[index]
        for digit in digit_options[current_index]:
            current_digits.append(digit)
            backtrack(index + 1, current_digits)
            current_digits.pop()
            if found:
                return

    backtrack(0, [])
    return [found_password] if found_password else [], tries


def crack_method3(conditions, stored_hash):
    """双向回溯：从左右两侧同时构建密码"""
    tries = 1
    found_password = None
    found = False
    lock = PasswordLock()

    digit_options = [[] for _ in range(3)]

    for i in range(3):
        for d in range(10):
            valid = True
            for cond in conditions:
                if cond == [-1, -1]:
                    if not is_prime(d):
                        valid = False
                        break
                elif len(cond) == 2:
                    a, t = cond
                    a = abs(a)
                    if a - 1 == i and ((t == 0 and d % 2 != 0) or (t == 1 and d % 2 != 1)):
                        valid = False
                        break
                elif len(cond) == 3:
                    if cond[i] != -1 and d != cond[i]:
                        valid = False
                        break
            if valid:
                digit_options[i].append(d)

    for i in range(3):
        random.shuffle(digit_options[i])

    def backtrack(left_pos, right_pos, current_digits):
        nonlocal tries, found, found_password

        if found:
            return

        if len(current_digits) == 3:
            password = "".join(map(str, current_digits))
            calculated_hash = lock.hash_password(password)

            if calculated_hash == stored_hash:
                found_password = password
                found = True
                return

            tries += 1
            return

        # 根据当前方向选择位置
        if left_pos <= right_pos:
            for digit in digit_options[left_pos]:
                current_digits.append(digit)
                backtrack(left_pos + 1, right_pos, current_digits)
                current_digits.pop()
                if found:
                    return
        else:
            for digit in digit_options[right_pos]:
                current_digits.insert(0, digit)
                backtrack(left_pos, right_pos - 1, current_digits)
                current_digits.pop(0)
                if found:
                    return

    backtrack(0, 2, [])
    return [found_password] if found_password else [], tries

def format_json(data):
    """自定义 JSON 格式化函数"""
    lines = ['{']

    # 格式化 C 线索部分
    if "C" in data:
        lines.append('  "C": [')
        for i, item in enumerate(data["C"]):
            lines.append(f"    {json.dumps(item, ensure_ascii=False)}{',' if i < len(data['C']) - 1 else ''}")
        lines.append('  ],')

    # 格式化 L 哈希值
    if "L" in data:
        lines.append(f'  "L": {json.dumps(data["L"], ensure_ascii=False)},')

    # 格式化 password
    if "password" in data:
        lines.append(f'  "password": {json.dumps(data["password"], ensure_ascii=False)},')

    # 格式化 results
    if "results" in data:
        lines.append('  "results": {')
        result_items = data["results"]
        method_names = list(result_items.keys())
        for i, method in enumerate(method_names):
            line = f'    "{method}": {json.dumps(result_items[method], ensure_ascii=False)}'
            if i < len(result_items) - 1:
                line += ","
            lines.append(line)
        lines.append('  }')

    lines.append('}')
    return "\n".join(lines)


def main():
    # 输入和输出路径
    input_dir = r"pwd"
    output_dir = r"pwd_output"

    # 如果输出目录不存在，则创建
    os.makedirs(output_dir, exist_ok=True)

    # 遍历输入目录下所有 .json 文件
    for filename in os.listdir(input_dir):
        if filename.endswith(".json"):
            input_file_path = os.path.join(input_dir, filename)
            output_file_path = os.path.join(output_dir, filename)

            with open(input_file_path, 'r', encoding='utf-8') as f:
                try:
                    # 读取原始 JSON 文本内容（字符串）
                    original_content = f.read()
                    # 解析为字典用于处理
                    sample_data = json.loads(original_content)
                except json.JSONDecodeError:
                    print(f"跳过无效的 JSON 文件: {filename}")
                    continue

            conditions = sample_data.get("C", [])
            stored_hash = sample_data.get("L", "")

            print(f"\n正在处理文件: {filename}")

            # 调用三种方法破解密码
            pwd1, tries1 = crack_method1(conditions, stored_hash)
            pwd2, tries2 = crack_method2(conditions, stored_hash)
            pwd3, tries3 = crack_method3(conditions, stored_hash)

            # 找到第一个非空结果作为最终密码
            matched_passwords = []
            if pwd1:
                matched_passwords.append(pwd1[0])
            if pwd2:
                matched_passwords.append(pwd2[0])
            if pwd3:
                matched_passwords.append(pwd3[0])

            first_match = matched_passwords[0] if matched_passwords else ""

            # 打印每个方法的尝试次数
            print(f"方法1尝试次数: {tries1}")
            print(f"方法2尝试次数: {tries2}")
            print(f"方法3尝试次数: {tries3}")

            # 计算并打印平均尝试次数
            avg_tries = (tries1 + tries2 + tries3) / 3
            print(f"平均尝试次数: {avg_tries:.1f}")

            # 更新 JSON 内容中的结果部分
            result_update = {
                "password": first_match,
                "results": {
                    "method1": {"tries": tries1, "password": list(map(int, first_match)) if first_match else []},
                    "method2": {"tries": tries2, "password": list(map(int, first_match)) if first_match else []},
                    "method3": {"tries": tries3, "password": list(map(int, first_match)) if first_match else []}
                }
            }

            # 加载原始内容为 dict 并更新关键字段
            data = json.loads(original_content)
            data.update(result_update)

            # 使用自定义格式化函数生成 JSON 字符串
            formatted_json = format_json(data)

            # 写入文件
            with open(output_file_path, 'w', encoding='utf-8') as f:
                f.write(formatted_json)

            print(f"结果已保存至: {output_file_path}")

            # 新增：比较写入内容是否与原始内容一致（忽略 results 字段中的 tries）
            with open(output_file_path, 'r', encoding='utf-8') as f:
                written_content = f.read()

            # 解析原始和写入后的内容为 dict
            original_dict = json.loads(original_content)
            written_dict = json.loads(written_content)

            # 删除 results 字段用于比较
            original_dict.pop("results", None)
            written_dict.pop("results", None)

            # 重新序列化为字符串进行比较
            if json.dumps(original_dict, sort_keys=True) == json.dumps(written_dict, sort_keys=True):
                print("输出内容（除 results 外）与原始内容一致")
            else:
                print("输出内容（除 results 外）与原始内容不一致")


if __name__ == "__main__":
    main()