import numpy as np
from PIL import Image
from typing import Optional, Dict, List
import hashlib
import base64
import time
from datetime import datetime, timezone, timedelta
from cryptography.fernet import Fernet, InvalidToken

def _generate_key_from_password(password: str) -> bytes:
    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:
    img = Image.open(image_path).convert("RGB")
    pixels = np.array(img)
    height, width, channels = pixels.shape
    total_channels = height * width * channels

    # 准备载荷
    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')                     # 32位长度前缀
    binary_message = len_binary + ''.join(format(b, '08b') for b in to_embed)
    required_bits = len(binary_message)

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

    # 1. 先将长度前缀嵌入到前32个像素通道（固定位置）
    new_pixels = pixels.copy().astype(np.uint8)
    fixed_positions = list(range(32))  # 平坦通道索引 0~31
    for idx, flat_pos in enumerate(fixed_positions):
        bit = int(len_binary[idx])
        pixel_idx = flat_pos // channels
        channel_idx = flat_pos % channels
        row = pixel_idx // width
        col = pixel_idx % width
        new_pixels[row, col, channel_idx] = (new_pixels[row, col, channel_idx] & ~1) | bit

    # 2. 剩余位置随机打乱
    remaining_positions = list(range(32, total_channels))
    rng = np.random.RandomState(seed) if seed is not None else np.random.RandomState()
    rng.shuffle(remaining_positions)

    # 3. 嵌入剩余消息比特
    bit_index = 32
    for flat_pos in remaining_positions[:required_bits - 32]:
        bit = int(binary_message[bit_index])
        bit_index += 1
        pixel_idx = flat_pos // channels
        channel_idx = flat_pos % channels
        row = pixel_idx // width
        col = pixel_idx % width
        new_pixels[row, col, channel_idx] = (new_pixels[row, col, channel_idx] & ~1) | bit

    new_img = Image.fromarray(new_pixels)
    new_img.save(output_path)
    print(f"全图随机LSB消息嵌入成功，已保存到：{output_path}")
    print(f"实际嵌入位数: {bit_index}")

def extract_message(
    image_path: str,
    seed: Optional[int] = None,
    password: Optional[str] = None,
) -> Dict[str, any]:
    img = Image.open(image_path).convert("RGB")
    pixels = np.array(img)
    height, width, channels = pixels.shape
    total_channels = height * width * channels

    # 1. 从固定位置（前32个像素通道）读取长度前缀
    fixed_positions = list(range(32))
    len_binary = ''
    for flat_pos in fixed_positions:
        pixel_idx = flat_pos // channels
        channel_idx = flat_pos % channels
        row = pixel_idx // width
        col = pixel_idx % width
        bit = pixels[row, col, channel_idx] & 1
        len_binary += str(bit)

    try:
        msg_len = int(len_binary, 2)
    except ValueError:
        raise ValueError("无法解析消息长度前缀")
    required_bits = 32 + msg_len * 8
    print(f"长度前缀解析成功，msg_len = {msg_len}, 所需总 bit 数: {required_bits}")

    if required_bits > total_channels:
        raise ValueError(f"所需位数 {required_bits} 超过图像总通道数 {total_channels}，数据可能损坏")

    # 2. 剩余位置随机打乱（必须与嵌入时使用相同的种子）
    remaining_positions = list(range(32, total_channels))
    rng = np.random.RandomState(seed) if seed is not None else np.random.RandomState()
    rng.shuffle(remaining_positions)

    # 3. 读取消息比特
    binary_list = list(len_binary)  # 前32位已有
    bit_index = 32
    for flat_pos in remaining_positions[:required_bits - 32]:
        pixel_idx = flat_pos // channels
        channel_idx = flat_pos % channels
        row = pixel_idx // width
        col = pixel_idx % width
        bit = pixels[row, col, channel_idx] & 1
        binary_list.append(str(bit))
        bit_index += 1

    binary = ''.join(binary_list)
    print("实际提取总 bit 数:", len(binary))

    if len(binary) < required_bits:
        raise ValueError(f"提取位数不足，期望 {required_bits}，实际 {len(binary)}")

    # 4. 提取消息字节并解密
    msg_binary = binary[32:required_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("密码错误或数据损坏")
        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
    }