from PIL import Image
from typing import Optional
import random
import hashlib
from cryptography.fernet import Fernet, InvalidToken
import base64
import time
from datetime import datetime, timezone, timedelta


def _generate_key_from_password(password: str) -> bytes:
    """从密码生成 Fernet 密钥（使用 SHA-256）"""
    return base64.urlsafe_b64encode(hashlib.sha256(password.encode()).digest())


def embed_message(
    image_path: str,
    message: str,
    output_path: str,
    seed: Optional[int] = None,
    password: Optional[str] = None,
    include_timestamp: bool = True
) -> None:
    """
    使用 LSB 隐写将消息嵌入图片。
    支持可选的随机种子打乱位置 + 密码加密 + 时间戳嵌入。

    Args:
        image_path: 输入图片路径（推荐 PNG 无损格式）
        message: 要隐藏的消息（字符串）
        output_path: 输出图片路径
        seed: 随机种子（用于打乱嵌入顺序，提高隐蔽性）
        password: 加密密码（若提供，则先加密消息再嵌入）
        include_timestamp: 是否在消息前嵌入当前时间戳（Unix 时间，秒级）
    """
    img = Image.open(image_path).convert("RGB")
    pixels = list(img.getdata())
    num_pixels = len(pixels)
    available_bits = num_pixels * 3

    # 获取当前时间戳（秒级）
    timestamp = int(time.time()) if include_timestamp else 0
    timestamp_bytes = timestamp.to_bytes(8, byteorder='big')

    # 准备负载：时间戳 + 消息
    message_bytes = message.encode('utf-8')
    payload = timestamp_bytes + message_bytes

    # 如果有密码，加密负载
    if password:
        key = _generate_key_from_password(password)
        fernet = Fernet(key)
        to_embed = fernet.encrypt(payload)
    else:
        to_embed = payload

    # 添加 32 位长度前缀（to_embed 的长度）
    msg_len = len(to_embed)
    len_binary = format(msg_len, '032b')
    binary_message = len_binary + ''.join(format(b, '08b') for b in to_embed)
    required_bits = len(binary_message)

    if required_bits > available_bits:
        raise ValueError(f"消息太长（需要 {required_bits} 位，可用 {available_bits} 位）")

    # 生成嵌入位置（顺序或打乱）
    positions = list(range(available_bits))
    if seed is not None:
        random.seed(seed)
        random.shuffle(positions)

    # 嵌入位
    bit_index = 0
    new_pixels = [list(p) for p in pixels]  # 转为可变列表
    for pos in positions[:required_bits]:
        pixel_idx = pos // 3
        channel_idx = pos % 3
        bit = int(binary_message[bit_index])
        new_pixels[pixel_idx][channel_idx] = (new_pixels[pixel_idx][channel_idx] & ~1) | bit
        bit_index += 1

    # 转回 tuple 并保存
    new_pixels = [tuple(p) for p in new_pixels]
    img.putdata(new_pixels)
    img.save(output_path)
    print(f"消息嵌入成功，已保存到：{output_path}")


def extract_message(
    image_path: str,
    seed: Optional[int] = None,
    password: Optional[str] = None
) -> dict:
    """
    从图片中提取使用 LSB 隐藏的消息。
    支持随机种子打乱 + 密码解密 + 时间戳解析。

    Args:
        image_path: 隐写图片路径
        seed: 嵌入时使用的随机种子（必须一致）
        password: 加密时使用的密码（必须一致）

    Returns:
        dict: {
            "message": str,  # 原始消息
            "timestamp": int,  # Unix 时间戳（秒）
            "timestamp_human": str  # 人可读格式，如 '2025-02-13 12:00:00 UTC'
        }

    Raises:
        ValueError: 各种解析/解密失败情况
    """
    img = Image.open(image_path).convert("RGB")
    pixels = list(img.getdata())
    num_pixels = len(pixels)
    available_bits = num_pixels * 3

    # 生成位置序列（与嵌入时相同）
    positions = list(range(available_bits))
    if seed is not None:
        random.seed(seed)
        random.shuffle(positions)

    # 位置 → 原始顺序索引 的映射（高效查找）
    pos_to_original_idx = {pos: idx for idx, pos in enumerate(positions)}

    # 提取所有位，按原始顺序排列
    binary_list = ['0'] * available_bits
    for pos in range(available_bits):
        pixel_idx = pos // 3
        channel_idx = pos % 3
        bit = pixels[pixel_idx][channel_idx] & 1
        original_idx = pos_to_original_idx[pos]
        binary_list[original_idx] = str(bit)

    binary = ''.join(binary_list)

    # 读取长度前缀
    len_str = binary[:32]
    try:
        msg_len = int(len_str, 2)
    except ValueError:
        raise ValueError("无法解析消息长度前缀，可能不是本方法嵌入的图片")

    required_bits = 32 + msg_len * 8
    if required_bits > available_bits:
        raise ValueError("图片大小不足以包含声明的消息长度")

    msg_binary = binary[32:required_bits]

    # 检查长度完整性
    if len(msg_binary) != msg_len * 8:
        raise ValueError(f"提取的二进制长度 {len(msg_binary)} 与预期 {msg_len * 8} 不匹配")

    # 二进制转字节
    msg_bytes = bytearray(
        int(msg_binary[i:i+8], 2)
        for i in range(0, len(msg_binary), 8)
    )

    # 如果有密码，尝试解密
    if password:
        key = _generate_key_from_password(password)
        fernet = Fernet(key)
        try:
            decrypted = fernet.decrypt(bytes(msg_bytes))
        except InvalidToken:
            raise ValueError("密码错误或数据损坏")
        except Exception as e:
            raise ValueError(f"解密失败：{type(e).__name__} - {str(e)}")
    else:
        decrypted = bytes(msg_bytes)

    # 解析时间戳 + 消息
    if len(decrypted) < 8:
        raise ValueError("数据太短，无法解析时间戳")

    timestamp_bytes = decrypted[:8]
    message_bytes = decrypted[8:]

    timestamp = int.from_bytes(timestamp_bytes, byteorder='big')
    try:
        message = message_bytes.decode('utf-8')
    except UnicodeDecodeError:
        raise ValueError("消息解码失败，可能数据损坏")
    tz_utc_8 = timezone(timedelta(hours=8))
    human_time = datetime.fromtimestamp(timestamp, tz=tz_utc_8).strftime('%Y-%m-%d %H:%M:%S UTC+8')

    return {
        "message": message,
        "timestamp": timestamp,
        "timestamp_human": human_time
    }


