import os
import json
import tempfile
from importlib import import_module
from typing import Any, Dict, List, Optional

# 每个条目描述如何调用该算法的加密/解密函数，以及支持哪些参数
_ALGORITHM_REGISTRY: Dict[str, Dict[str, Any]] = {
    "pixel": {
        "label": "像素级随机置换 (Pixel Scramble)",
        "module": "ImageStego.obfuscation.pixel_scramble",
        "encrypt_name": "obfuscate",
        "decrypt_name": "deobfuscate",
        "encrypt_params": ["seed", "use_timestamp"],
        "decrypt_params": ["seed"],
        "defaults": {"seed": None, "use_timestamp": False},
    },
    "block": {
        "label": "块级随机置换 (Block Scramble)",
        "module": "ImageStego.obfuscation.block_scramble",
        "encrypt_name": "obfuscate",
        "decrypt_name": "deobfuscate",
        "encrypt_params": ["block_size", "seed", "use_timestamp"],
        "decrypt_params": ["block_size", "seed"],
        "defaults": {"block_size": 16, "seed": None, "use_timestamp": False},
    },
    "arnold": {
        "label": "Arnold 猫映射 (Arnold Cat Map)",
        "module": "ImageStego.obfuscation.arnold_cat_map",
        "encrypt_name": "obfuscate",
        "decrypt_name": "deobfuscate",
        "encrypt_params": ["iterations", "seed", "use_timestamp"],
        "decrypt_params": ["iterations", "seed"],
        "defaults": {"iterations": 50, "seed": None, "use_timestamp": False},
    },
    "bakers": {
        "label": "Baker 映射 + XOR 加密 (Baker's Map)",
        "module": "ImageStego.obfuscation.bakers_map",
        "encrypt_name": "obfuscate",
        "decrypt_name": "deobfuscate",
        "encrypt_params": ["iterations", "seed", "password", "use_timestamp"],
        "decrypt_params": ["iterations", "seed", "password", "orig_width", "orig_height", "padded_size", "left_pad",
                           "top_pad", "encrypted_flag"],
        "defaults": {"iterations": 20, "seed": None, "password": None, "use_timestamp": False},
        "capture_metadata": True,  # 标记需要从加密输出中捕获元数据
    },
    "henon": {
        "label": "Henon 混沌映射 (Henon Map)",
        "module": "ImageStego.obfuscation.henon_map",
        "encrypt_name": "obfuscate",
        "decrypt_name": "deobfuscate",
        "encrypt_params": ["iterations", "seed", "password"],
        "decrypt_params": ["iterations", "seed", "password"],
        "defaults": {"iterations": 1, "seed": None, "password": None},
    },
}

SUPPORTED_ALGORITHMS = list(_ALGORITHM_REGISTRY.keys())


def _sidecar_path(image_path: str) -> str:
    """根据图片路径生成对应的 sidecar JSON 路径。"""
    stem, _ = os.path.splitext(image_path)
    return stem + ".layers.json"


def _call_encrypt(algorithm: str, input_path: str, output_path: str, params: Dict[str, Any]) -> Dict[str, Any]:
    """
    统一调用加密函数，屏蔽各算法接口差异。

    Returns:
        捕获的元数据字典（如果算法需要），否则返回空字典
    """
    entry = _ALGORITHM_REGISTRY[algorithm]
    module = import_module(entry["module"])
    encrypt_fn = getattr(module, entry["encrypt_name"])
    allowed_keys = entry["encrypt_params"]
    kwargs = {key: params[key] for key in allowed_keys if key in params}
    encrypt_fn(input_path, output_path, **kwargs)

    # 如果算法需要捕获元数据（如 bakers），从输出图片读取
    captured_metadata = {}
    if entry.get("capture_metadata", False):
        from PIL import Image
        try:
            img = Image.open(output_path)
            if algorithm == "bakers":
                # 捕获 bakers 的关键元数据
                captured_metadata["orig_width"] = int(img.info.get("orig_width", 0))
                captured_metadata["orig_height"] = int(img.info.get("orig_height", 0))
                captured_metadata["padded_size"] = int(img.info.get("padded_size", 0))
                captured_metadata["left_pad"] = int(img.info.get("left_pad", 0))
                captured_metadata["top_pad"] = int(img.info.get("top_pad", 0))
                captured_metadata["encrypted_flag"] = img.info.get("bakers_encrypted", "0")
        except Exception:
            pass  # 元数据读取失败不影响主流程

    return captured_metadata


def _call_decrypt(algorithm: str, input_path: str, output_path: str, params: Dict[str, Any]) -> None:
    """统一调用解密函数，屏蔽各算法接口差异。"""
    entry = _ALGORITHM_REGISTRY[algorithm]
    module = import_module(entry["module"])
    decrypt_fn = getattr(module, entry["decrypt_name"])
    allowed_keys = entry["decrypt_params"]
    kwargs = {key: params[key] for key in allowed_keys if key in params}
    decrypt_fn(input_path, output_path, **kwargs)

class ObfuscationPipeline:
    def __init__(self) -> None:
        # 每个元素：{"algorithm": str, "params": dict}
        self._layers: List[Dict[str, Any]] = []

    def add_layer(self, algorithm: str, **params: Any) -> "ObfuscationPipeline":
        if algorithm not in _ALGORITHM_REGISTRY:
            raise ValueError(
                f"不支持的算法 '{algorithm}'，可选：{SUPPORTED_ALGORITHMS}"
            )
        entry = _ALGORITHM_REGISTRY[algorithm]
        merged_params = {**entry["defaults"], **params}
        self._layers.append({"algorithm": algorithm, "params": merged_params})
        return self

    def clear_layers(self) -> None:
        """清空所有已添加的混淆层。"""
        self._layers.clear()

    @property
    def layer_count(self) -> int:
        return len(self._layers)

    def encrypt(self, input_path: str, output_path: str) -> str:
        if not self._layers:
            raise ValueError("没有添加任何混淆层，请先调用 add_layer()")
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件不存在：{input_path}")

        print(f"\n开始多层加密，共 {len(self._layers)} 层")
        print(f"输入：{input_path}")

        current_input = input_path
        temp_files: List[str] = []

        try:
            for index, layer in enumerate(self._layers):
                algorithm = layer["algorithm"]
                params = layer["params"]
                label = _ALGORITHM_REGISTRY[algorithm]["label"]
                is_last_layer = (index == len(self._layers) - 1)

                print(f"\n  第 {index + 1} 层：{label}")

                if is_last_layer:
                    # 最后一层直接写到目标路径
                    layer_output = output_path
                else:
                    # 中间层写到临时文件
                    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
                    tmp.close()
                    layer_output = tmp.name
                    temp_files.append(layer_output)

                captured_metadata = _call_encrypt(algorithm, current_input, layer_output, params)
                # 将捕获的元数据合并到 params 中，供解密时使用
                if captured_metadata:
                    params.update(captured_metadata)
                current_input = layer_output

        finally:
            # 清理中间临时文件
            for temp_file in temp_files:
                if os.path.exists(temp_file):
                    os.remove(temp_file)

        # 写 sidecar JSON
        sidecar = _sidecar_path(output_path)
        with open(sidecar, "w", encoding="utf-8") as json_file:
            json.dump(self._layers, json_file, ensure_ascii=False, indent=2)

        print("\n[OK] ????")
        print(f"  加密图片：{output_path}")
        print(f"  层信息文件：{sidecar}")
        return sidecar

    def decrypt(self, encrypted_path: str, output_path: str, sidecar: Optional[str] = None) -> None:
        """
        从 sidecar JSON 读取层信息，逆序执行解密还原图片。

        Args:
            encrypted_path: 加密后的图片路径
            output_path: 解密后图片输出路径
            sidecar: sidecar JSON 路径，默认自动推断（{encrypted_stem}.layers.json）
        """
        if not os.path.exists(encrypted_path):
            raise FileNotFoundError(f"加密图片不存在：{encrypted_path}")

        sidecar_path = sidecar or _sidecar_path(encrypted_path)
        if not os.path.exists(sidecar_path):
            raise FileNotFoundError(
                f"找不到层信息文件：{sidecar_path}\n"
                "请确保 .layers.json 文件与加密图片在同一目录"
            )

        with open(sidecar_path, "r", encoding="utf-8") as json_file:
            layers: List[Dict[str, Any]] = json.load(json_file)

        if not layers:
            raise ValueError("层信息文件为空，无法解密")

        print(f"\n开始多层解密，共 {len(layers)} 层（逆序执行）")
        print(f"输入：{encrypted_path}")

        # 逆序解密
        reversed_layers = list(reversed(layers))
        current_input = encrypted_path
        temp_files: List[str] = []

        try:
            for index, layer in enumerate(reversed_layers):
                algorithm = layer["algorithm"]
                params = layer["params"]
                label = _ALGORITHM_REGISTRY[algorithm]["label"]
                is_last_layer = (index == len(reversed_layers) - 1)
                original_layer_number = len(layers) - index

                print(f"\n  还原第 {original_layer_number} 层：{label}")

                if is_last_layer:
                    layer_output = output_path
                else:
                    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
                    tmp.close()
                    layer_output = tmp.name
                    temp_files.append(layer_output)

                _call_decrypt(algorithm, current_input, layer_output, params)
                current_input = layer_output

        finally:
            for temp_file in temp_files:
                if os.path.exists(temp_file):
                    os.remove(temp_file)

        print("\n[OK] ????")
        print(f"  还原图片：{output_path}")

    def summary(self) -> str:
        """返回当前 pipeline 的层信息摘要。"""
        if not self._layers:
            return "（当前没有任何混淆层）"
        lines = [f"当前混淆层（共 {len(self._layers)} 层）："]
        for index, layer in enumerate(self._layers):
            algorithm = layer["algorithm"]
            label = _ALGORITHM_REGISTRY[algorithm]["label"]
            params = {key: value for key, value in layer["params"].items() if value is not None}
            lines.append(f"  第 {index + 1} 层：{label}  参数：{params}")
        return "\n".join(lines)
