from PIL import Image
import numpy as np
from typing import Optional
import random
import hashlib
from cryptography.fernet import Fernet, InvalidToken
import base64
import time
from datetime import datetime, timezone, timedelta
from scipy.fftpack import dct, idct


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


def _dct2(block):
    """2D DCT 变换"""
    return dct(dct(block.T, norm='ortho').T, norm='ortho')


def _idct2(block):
    """2D IDCT 逆变换"""
    return idct(idct(block.T, norm='ortho').T, norm='ortho')


def _embed_bit_in_dct_block(block, bit, position=(4, 4), strength=10):
    i, j = position
    coeff = block[i, j]

    # 量化嵌入：根据比特值调整系数
    if bit == 1:
        # 确保系数为正且足够大
        if coeff < strength:
            block[i, j] = strength
        elif coeff < 0:
            block[i, j] = strength
    else:
        # 确保系数为负或足够小
        if coeff > -strength:
            block[i, j] = -strength
        elif coeff > 0:
            block[i, j] = -strength

    return block


def _extract_bit_from_dct_block(block, position=(4, 4)):
    i, j = position
    coeff = block[i, j]
    return 1 if coeff > 0 else 0


def embed_message(
        image_path: str,
        message: str,
        output_path: str,
        seed: Optional[int] = None,
        password: Optional[str] = None,
        include_timestamp: bool = True,
        strength: int = 15,
        block_size: int = 8
) -> None:
    # 读取图像并转换为 YCbCr 色彩空间（在 Y 通道嵌入）
    img = Image.open(image_path).convert('RGB')
    img_ycbcr = img.convert('YCbCr')
    y, cb, cr = img_ycbcr.split()
    y_array = np.array(y, dtype=np.float32)

    height, width = y_array.shape

    # 确保图像尺寸是块大小的整数倍
    if height % block_size != 0 or width % block_size != 0:
        new_height = (height // block_size) * block_size
        new_width = (width // block_size) * block_size
        y_array = y_array[:new_height, :new_width]
        height, width = new_height, new_width

    # 计算可用块数
    num_blocks_h = height // block_size
    num_blocks_w = width // block_size
    total_blocks = num_blocks_h * num_blocks_w

    # 准备消息
    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

    # 添加长度前缀
    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 > total_blocks:
        raise ValueError(f"消息太长（需要 {required_bits} 块，可用 {total_blocks} 块）")

    # 生成块序列（顺序或打乱）
    block_indices = list(range(total_blocks))
    if seed is not None:
        random.seed(seed)
        random.shuffle(block_indices)

    # 嵌入消息
    bit_index = 0
    for block_idx in block_indices[:required_bits]:
        # 计算块位置
        block_row = block_idx // num_blocks_w
        block_col = block_idx % num_blocks_w

        # 提取块
        row_start = block_row * block_size
        col_start = block_col * block_size
        block = y_array[row_start:row_start + block_size, col_start:col_start + block_size]

        # DCT 变换
        dct_block = _dct2(block)

        # 嵌入比特
        bit = int(binary_message[bit_index])
        dct_block = _embed_bit_in_dct_block(dct_block, bit, strength=strength)

        # IDCT 逆变换
        modified_block = _idct2(dct_block)

        # 写回
        y_array[row_start:row_start + block_size, col_start:col_start + block_size] = modified_block
        bit_index += 1

    # 裁剪到 [0, 255] 并转换回 uint8
    y_array = np.clip(y_array, 0, 255).astype(np.uint8)

    # 重建图像
    y_modified = Image.fromarray(y_array, mode='L')
    img_modified = Image.merge('YCbCr', (y_modified, cb, cr))
    img_rgb = img_modified.convert('RGB')

    # 保存
    img_rgb.save(output_path, quality=95)  # 使用高质量保存
    print(f"DCT 隐写成功，已保存到：{output_path}")
    print(f"嵌入强度：{strength}，块大小：{block_size}x{block_size}")


def extract_message(
        image_path: str,
        seed: Optional[int] = None,
        password: Optional[str] = None,
        block_size: int = 8
) -> dict:
    # 读取图像
    img = Image.open(image_path).convert('RGB')
    img_ycbcr = img.convert('YCbCr')
    y, _, _ = img_ycbcr.split()
    y_array = np.array(y, dtype=np.float32)

    height, width = y_array.shape

    # 确保尺寸对齐
    if height % block_size != 0 or width % block_size != 0:
        new_height = (height // block_size) * block_size
        new_width = (width // block_size) * block_size
        y_array = y_array[:new_height, :new_width]
        height, width = new_height, new_width

    # 计算块数
    num_blocks_h = height // block_size
    num_blocks_w = width // block_size
    total_blocks = num_blocks_h * num_blocks_w

    # 生成块序列
    block_indices = list(range(total_blocks))
    if seed is not None:
        random.seed(seed)
        random.shuffle(block_indices)

    # 提取长度前缀（前 32 位）
    len_bits = []
    for i in range(32):
        block_idx = block_indices[i]
        block_row = block_idx // num_blocks_w
        block_col = block_idx % num_blocks_w

        row_start = block_row * block_size
        col_start = block_col * block_size
        block = y_array[row_start:row_start + block_size, col_start:col_start + block_size]

        dct_block = _dct2(block)
        bit = _extract_bit_from_dct_block(dct_block)
        len_bits.append(str(bit))

    len_str = ''.join(len_bits)
    try:
        msg_len = int(len_str, 2)
    except ValueError:
        raise ValueError("无法解析消息长度前缀")

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

    # 提取消息位
    msg_bits = []
    for i in range(32, required_bits):
        block_idx = block_indices[i]
        block_row = block_idx // num_blocks_w
        block_col = block_idx % num_blocks_w

        row_start = block_row * block_size
        col_start = block_col * block_size
        block = y_array[row_start:row_start + block_size, col_start:col_start + block_size]

        dct_block = _dct2(block)
        bit = _extract_bit_from_dct_block(dct_block)
        msg_bits.append(str(bit))

    # 二进制转字节
    msg_binary = ''.join(msg_bits)
    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("密码错误或数据损坏")
    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
    }