#!/usr/bin/env python3
"""从“郑州情趣王一手科技.exe”提取的 JPEG 图片处理逻辑。

处理步骤与原程序一致：
1. 裁掉 JPEG 底部两行并以 quality=95 重新编码；
2. 把 SOF 段中的图片高度改回原始高度；
3. 删除最后一个 JPEG EOI（FF D9）结束标志；
4. 按指定大小追加随机字节，输出为同名 .file 文件。

注意：输出是刻意构造的非标准 JPEG，部分看图软件可能拒绝打开。
"""

from __future__ import annotations

import argparse
import random
import sys
from pathlib import Path
from typing import Iterable

try:
    from PIL import Image
except ImportError as exc:
    raise SystemExit("缺少 Pillow，请先执行：python3 -m pip install pillow") from exc


JPEG_EXTENSIONS = {".jpg", ".jpeg", ".jpe", ".jfif"}
SOF_MARKERS = {
    0xC0,
    0xC1,
    0xC2,
    0xC3,
    0xC5,
    0xC6,
    0xC7,
    0xC9,
    0xCA,
    0xCB,
    0xCD,
    0xCE,
    0xCF,
}


def find_and_modify_height(data: bytearray, original_height: int) -> bytearray | None:
    """找到第一个 JPEG SOF 段，把高度字段恢复为裁剪前的高度。"""
    index = 0
    length = len(data)

    while index < length - 1:
        if data[index] != 0xFF:
            index += 1
            continue

        marker = data[index + 1]
        if marker in SOF_MARKERS:
            height_offset = index + 5
            if height_offset + 1 >= length:
                return None
            data[height_offset] = (original_height >> 8) & 0xFF
            data[height_offset + 1] = original_height & 0xFF
            return data

        # SOI 和 EOI 标记没有长度字段。
        if marker in (0xD8, 0xD9):
            index += 2
            continue

        if index + 3 < length:
            segment_length = (data[index + 2] << 8) | data[index + 3]
            index += 2 + segment_length
        else:
            index += 2

    return None


def remove_eoi(data: bytearray) -> tuple[bytearray, bool]:
    """从末尾向前查找并删除最后一个 FF D9 结束标志。"""
    for index in range(len(data) - 1, 0, -1):
        if data[index - 1] == 0xFF and data[index] == 0xD9:
            return data[: index - 1] + data[index + 1 :], True
    return data, False


def process_single_image(
    input_file: Path | str,
    output_file: Path | str,
    junk_size_mb: float = 1.0,
) -> Path:
    """按原程序逻辑处理一张图片，并返回输出路径。"""
    input_path = Path(input_file)
    output_path = Path(output_file)
    temp_path = Path(f"{output_path}.tmp")

    if not 0 <= junk_size_mb <= 100:
        raise ValueError("junk_size_mb 必须在 0 到 100 之间")

    output_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        with Image.open(input_path) as image:
            width, height = image.size
            if height <= 2:
                raise ValueError(f"图片高度不足 3 行：{input_path}")

            cropped_image = image.crop((0, 0, width, height - 2))
            try:
                cropped_image.save(temp_path, "JPEG", quality=95)
            finally:
                cropped_image.close()

        image_data = bytearray(temp_path.read_bytes())
        modified_data = find_and_modify_height(image_data, height) or image_data
        modified_data, _ = remove_eoi(modified_data)

        junk_size_bytes = int(junk_size_mb * 1024 * 1024)
        if junk_size_bytes > 0:
            # 保留原程序的逐字节伪随机生成方式。
            modified_data += bytearray(
                random.getrandbits(8) for _ in range(junk_size_bytes)
            )

        output_path.write_bytes(modified_data)
        return output_path
    finally:
        if temp_path.exists():
            temp_path.unlink()


def collect_images(paths: Iterable[Path]) -> list[Path]:
    """收集输入文件，并按原程序支持的扩展名递归扫描目录。"""
    images: list[Path] = []
    for path in paths:
        if path.is_file() and path.suffix.lower() in JPEG_EXTENSIONS:
            images.append(path)
        elif path.is_dir():
            images.extend(
                candidate
                for candidate in path.rglob("*")
                if candidate.is_file()
                and candidate.suffix.lower() in JPEG_EXTENSIONS
            )
        else:
            print(f"跳过无效路径或非 JPEG 文件：{path}", file=sys.stderr)
    return images


def build_output_path(input_path: Path, output_dir: Path | None) -> Path:
    """原程序使用图片主文件名加 .file 作为输出名。"""
    directory = output_dir if output_dir is not None else input_path.parent
    return directory / f"{input_path.stem}.file"


def main() -> int:
    parser = argparse.ArgumentParser(
        description="裁掉 JPEG 底部两行、伪造原高度、删除 EOI 并追加随机数据。"
    )
    parser.add_argument("inputs", nargs="+", type=Path, help="JPEG 文件或目录")
    parser.add_argument("-o", "--output-dir", type=Path, help="统一输出目录")
    parser.add_argument(
        "--junk-mb",
        type=float,
        default=1.0,
        help="每张图片追加的随机数据大小，范围 0–100 MB，默认 1 MB",
    )
    args = parser.parse_args()

    if not 0 <= args.junk_mb <= 100:
        parser.error("--junk-mb 必须在 0 到 100 之间")

    images = collect_images(args.inputs)
    if not images:
        print("没有找到可处理的 JPEG 图片。", file=sys.stderr)
        return 1

    succeeded = 0
    for index, input_path in enumerate(images, start=1):
        output_path = build_output_path(input_path, args.output_dir)
        try:
            process_single_image(input_path, output_path, args.junk_mb)
            succeeded += 1
            print(f"[{index}/{len(images)}] 完成：{input_path} -> {output_path}")
        except Exception as exc:
            print(f"[{index}/{len(images)}] 失败：{input_path} - {exc}", file=sys.stderr)

    print(f"处理完成：总计 {len(images)}，成功 {succeeded}，失败 {len(images) - succeeded}")
    return 0 if succeeded == len(images) else 2


if __name__ == "__main__":
    raise SystemExit(main())