import base64
import hashlib
import time
from datetime import datetime, timedelta, timezone
from typing import Iterator, List, Optional, Tuple

import numpy as np
from cryptography.fernet import Fernet, InvalidToken
from PIL import Image


def _generate_key_from_password(password: str) -> bytes:
    """从密码生成 Fernet 密钥。"""
    return base64.urlsafe_b64encode(hashlib.sha256(password.encode()).digest())


def _get_range_table() -> List[Tuple[int, int, int]]:
    """
    PVD 范围表，定义不同差异值区间对应的嵌入位数。
    返回: [(下限, 上限, 嵌入位数), ...]
    """
    return [
        (0, 7, 3),
        (8, 15, 3),
        (16, 31, 4),
        (32, 63, 5),
        (64, 127, 6),
        (128, 255, 7),
    ]


def _get_range_info(difference: int) -> Tuple[int, int, int]:
    """根据差异值返回对应的区间信息。"""
    for lower, upper, bits in _get_range_table():
        if lower <= difference <= upper:
            return lower, upper, bits
    return _get_range_table()[0]


def _get_embedding_capacity(difference: int) -> int:
    """根据像素差异值获取可嵌入的位数。"""
    return _get_range_info(difference)[2]


def _embed_in_pair(p1: int, p2: int, data_bits: str) -> Tuple[int, int]:
    """
    在一对像素中嵌入数据。

    Args:
        p1: 第一个像素值
        p2: 第二个像素值
        data_bits: 要嵌入的二进制字符串

    Returns:
        (new_p1, new_p2): 嵌入后的像素值
    """
    difference = abs(p2 - p1)
    lower_bound, _, embed_bits = _get_range_info(difference)

    if len(data_bits) < embed_bits:
        data_bits = data_bits.ljust(embed_bits, "0")
    else:
        data_bits = data_bits[:embed_bits]

    data_value = int(data_bits, 2)
    new_difference = lower_bound + data_value

    # 保持原有的大小规律，同时确保新差值可以被精确提取。
    if p1 <= p2:
        new_p1 = p1
        new_p2 = p1 + new_difference
        if new_p2 > 255:
            new_p2 = 255
            new_p1 = 255 - new_difference
    else:
        new_p2 = p2
        new_p1 = p2 + new_difference
        if new_p1 > 255:
            new_p1 = 255
            new_p2 = 255 - new_difference

    return int(new_p1), int(new_p2)


def _extract_from_pair(p1: int, p2: int, bit_count: Optional[int] = None) -> str:
    """
    从一对像素中提取嵌入的数据。

    Args:
        p1: 第一个像素值
        p2: 第二个像素值
        bit_count: 需要截取的位数，默认返回该像素对的全部可提取位数

    Returns:
        提取的二进制字符串
    """
    difference = abs(p2 - p1)
    lower_bound, _, embed_bits = _get_range_info(difference)
    data_value = difference - lower_bound
    extracted_bits = format(data_value, f"0{embed_bits}b")
    if bit_count is not None:
        return extracted_bits[:bit_count]
    return extracted_bits


def _iter_pixel_pairs(
        pixels: np.ndarray,
        height: int,
        width: int,
        channels: int
) -> Iterator[Tuple[int, int]]:
    """按嵌入时相同的顺序遍历像素对。"""
    for ch in range(channels):
        for row in range(height):
            for col in range(0, width, 2):
                yield int(pixels[row, col, ch]), int(pixels[row, col + 1, ch])


def embed_message(
        image_path: str,
        message: str,
        output_path: str,
        password: Optional[str] = None,
        include_timestamp: bool = True
) -> None:
    """
    使用 PVD (Pixel Value Differencing) 方法将消息嵌入图片。

    Args:
        image_path: 输入图片路径
        message: 要隐藏的消息
        output_path: 输出图片路径
        password: 可选密码，提供时先加密后嵌入
        include_timestamp: 是否在消息前附加当前时间戳
    """
    img = Image.open(image_path).convert("RGB")
    pixels = np.array(img, dtype=np.int32)
    height, width, channels = pixels.shape

    if width % 2 != 0:
        width -= 1
        pixels = pixels[:, :width, :]

    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(byte, "08b") for byte in to_embed)

    total_pairs = height * (width // 2) * channels
    estimated_capacity = total_pairs * 4

    if len(binary_message) > estimated_capacity:
        raise ValueError(
            f"消息太长，需要约 {len(binary_message)} 位，估计可用约 {estimated_capacity} 位"
        )

    new_pixels = pixels.copy()
    bit_index = 0
    total_embedded = 0

    for ch in range(channels):
        for row in range(height):
            for col in range(0, width, 2):
                if bit_index >= len(binary_message):
                    break

                p1 = int(new_pixels[row, col, ch])
                p2 = int(new_pixels[row, col + 1, ch])
                embed_bits = _get_embedding_capacity(abs(p2 - p1))

                remaining_bits = len(binary_message) - bit_index
                current_bits = min(embed_bits, remaining_bits)
                data_bits = binary_message[bit_index:bit_index + current_bits]

                new_p1, new_p2 = _embed_in_pair(p1, p2, data_bits)
                new_pixels[row, col, ch] = new_p1
                new_pixels[row, col + 1, ch] = new_p2

                bit_index += current_bits
                total_embedded += current_bits

            if bit_index >= len(binary_message):
                break

        if bit_index >= len(binary_message):
            break

    if bit_index < len(binary_message):
        raise ValueError(f"图片容量不足，只嵌入了 {bit_index}/{len(binary_message)} 位")

    new_img = Image.fromarray(np.clip(new_pixels, 0, 255).astype(np.uint8))
    new_img.save(output_path)
    print(f"PVD 消息嵌入成功，已保存到：{output_path}")
    print(f"总共嵌入了 {total_embedded} 位数据")


def extract_message(
        image_path: str,
        password: Optional[str] = None
) -> dict:
    """
    从图片中提取使用 PVD 方法隐藏的消息。

    Args:
        image_path: 含密图片路径
        password: 加密时使用的密码

    Returns:
        dict: {
            "message": str,
            "timestamp": int,
            "timestamp_human": str
        }
    """
    img = Image.open(image_path).convert("RGB")
    pixels = np.array(img, dtype=np.int32)
    height, width, channels = pixels.shape

    if width % 2 != 0:
        width -= 1

    pair_iterator = _iter_pixel_pairs(pixels, height, width, channels)

    extracted_bits = ""
    while len(extracted_bits) < 32:
        try:
            p1, p2 = next(pair_iterator)
        except StopIteration as exc:
            raise ValueError("无法解析消息长度前缀，图片中可提取数据不足") from exc

        extracted_bits += _extract_from_pair(p1, p2)

    len_binary = extracted_bits[:32]

    try:
        msg_len = int(len_binary[:32], 2)
    except ValueError as exc:
        raise ValueError("无法解析消息长度前缀，可能不是 PVD 方法嵌入的图片") from exc

    payload_bits_needed = msg_len * 8
    total_bits_needed = 32 + payload_bits_needed

    while len(extracted_bits) < total_bits_needed:
        try:
            p1, p2 = next(pair_iterator)
        except StopIteration as exc:
            actual_payload_bits = max(0, len(extracted_bits) - 32)
            raise ValueError(
                f"提取数据不完整，期望 {payload_bits_needed} 位消息体，实际 {actual_payload_bits} 位"
            ) from exc

        extracted_bits += _extract_from_pair(p1, p2)

    msg_binary = extracted_bits[32:total_bits_needed]
    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 as exc:
            raise ValueError("密码错误或数据损坏") from exc
        except Exception as exc:
            raise ValueError(f"解密失败：{type(exc).__name__} - {str(exc)}") from exc
    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 as exc:
        raise ValueError("消息解码失败，可能数据损坏") from exc

    tz_utc_8 = timezone(timedelta(hours=8))
    timestamp_human = ""
    if timestamp > 0:
        timestamp_human = datetime.fromtimestamp(
            timestamp, tz=tz_utc_8
        ).strftime("%Y-%m-%d %H:%M:%S UTC+8")

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