"""
Модуль DCT стеганографии с квантованием (JPEG-style)
"""

import numpy as np
from PIL import Image
from scipy.fftpack import dct, idct
from .utils import text_to_bits, bits_to_text

BLOCK_SIZE = 8

# Стандартная таблица квантования JPEG для качества 50%
QUANT_TABLE = np.array([
    [16, 11, 10, 16, 24, 40, 51, 61],
    [12, 12, 14, 19, 26, 58, 60, 55],
    [14, 13, 16, 24, 40, 57, 69, 56],
    [14, 17, 22, 29, 51, 87, 80, 62],
    [18, 22, 37, 56, 68, 109, 103, 77],
    [24, 35, 55, 64, 81, 104, 113, 92],
    [49, 64, 78, 87, 103, 121, 120, 101],
    [72, 92, 95, 98, 112, 100, 103, 99]
])


def _get_zigzag_order() -> list:
    """Zigzag порядок для матрицы 8x8."""
    zigzag = []
    for s in range(15):
        if s % 2 == 0:
            for i in range(min(s, 7), max(0, s - 7) - 1, -1):
                j = s - i
                if 0 <= i < 8 and 0 <= j < 8:
                    zigzag.append(i * 8 + j)
        else:
            for i in range(max(0, s - 7), min(s, 7) + 1):
                j = s - i
                if 0 <= i < 8 and 0 <= j < 8:
                    zigzag.append(i * 8 + j)
    return zigzag


ZIGZAG_ORDER = _get_zigzag_order()
MID_FREQ_INDICES = ZIGZAG_ORDER[1:40]


def _dct_quantize(block: np.ndarray) -> np.ndarray:
    """DCT + квантование."""
    dct_block = dct(dct(block, axis=0, norm='ortho'), axis=1, norm='ortho')
    quantized = np.round(dct_block / QUANT_TABLE)
    return quantized.astype(np.int32)


def _idct_dequantize(quantized: np.ndarray) -> np.ndarray:
    """Обратное квантование + IDCT."""
    dct_block = quantized * QUANT_TABLE
    block = idct(idct(dct_block, axis=0, norm='ortho'), axis=1, norm='ortho')
    block = np.round(block)
    return np.clip(block, 0, 255).astype(np.uint8)


def _extract_bits_from_block(block_quant: np.ndarray, max_bits: int) -> list:
    """Извлекает биты из квантованных коэффициентов."""
    extracted = []
    coeff_flat = block_quant.flatten()
    
    for idx in MID_FREQ_INDICES:
        if len(extracted) >= max_bits:
            break
        if idx >= len(coeff_flat):
            break
        bit = int(coeff_flat[idx]) & 1
        extracted.append(bit)
    return extracted


def _embed_bits_in_block(block_quant: np.ndarray, bits: list, bit_index: int) -> tuple:
    """Внедряет биты в квантованные коэффициенты."""
    modified = block_quant.copy()
    current_idx = bit_index
    coeff_flat = modified.flatten()
    
    for idx in MID_FREQ_INDICES:
        if current_idx >= len(bits):
            break
        if idx >= len(coeff_flat):
            break
        new_val = (int(coeff_flat[idx]) & 0xFE) | bits[current_idx]
        coeff_flat[idx] = new_val
        current_idx += 1
    
    return coeff_flat.reshape(BLOCK_SIZE, BLOCK_SIZE), current_idx


def encode_dct(image_path: str, message: str, output_path: str) -> bool:
    """Скрывает сообщение в изображении методом DCT."""
    img = Image.open(image_path).convert('L')
    pixels = np.array(img, dtype=np.float64)
    height, width = pixels.shape

    # Кодируем сообщение в байты
    message_bytes = message.encode('utf-8')
    msg_length = len(message_bytes)
    
    # Формат: [длина: 32 бита] + [сообщение]
    # 32 бита = 4 байта, достаточно для сообщений до 4 ГБ
    length_bits = format(msg_length, '032b')
    message_bits = text_to_bits(message)
    
    # Собираем всё вместе
    all_bits = length_bits + message_bits
    bit_list = [int(b) for b in all_bits]
    total_bits = len(bit_list)

    # Проверяем вместимость
    blocks_per_row = width // BLOCK_SIZE
    blocks_per_col = height // BLOCK_SIZE
    max_bits = blocks_per_row * blocks_per_col * len(MID_FREQ_INDICES)

    if total_bits > max_bits:
        print(f"Ошибка: сообщение слишком большое.")
        print(f"Доступно: {max_bits}, требуется: {total_bits}")
        return False

    modified_pixels = pixels.copy()
    bit_index = 0

    for i in range(0, height - BLOCK_SIZE + 1, BLOCK_SIZE):
        for j in range(0, width - BLOCK_SIZE + 1, BLOCK_SIZE):
            if bit_index >= total_bits:
                break
            
            block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE]
            quant_block = _dct_quantize(block)
            quant_block, bit_index = _embed_bits_in_block(quant_block, bit_list, bit_index)
            new_block = _idct_dequantize(quant_block)
            modified_pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE] = new_block
            
        if bit_index >= total_bits:
            break

    result_img = Image.fromarray(modified_pixels.astype(np.uint8), mode='L')
    result_img.save(output_path)
    print(f"Успешно! Спрятано {total_bits} бит ({total_bits // 8} байт)")
    return True


def decode_dct(image_path: str) -> str:
    """Извлекает сообщение из изображения методом DCT."""
    img = Image.open(image_path).convert('L')
    pixels = np.array(img, dtype=np.float64)
    height, width = pixels.shape

    # Извлекаем все биты
    all_bits = []
    
    for i in range(0, height - BLOCK_SIZE + 1, BLOCK_SIZE):
        for j in range(0, width - BLOCK_SIZE + 1, BLOCK_SIZE):
            block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE]
            quant_block = _dct_quantize(block)
            bits_from_block = _extract_bits_from_block(quant_block, len(MID_FREQ_INDICES))
            all_bits.extend(bits_from_block)

    # Читаем длину сообщения (первые 32 бита)
    if len(all_bits) < 32:
        return ""
    
    length_bits = all_bits[:32]
    length_str = ''.join(str(b) for b in length_bits)
    msg_length = int(length_str, 2)
    
    # Проверяем, что длина разумная
    if msg_length > len(all_bits) // 8:
        return ""
    
    # Читаем сообщение
    message_bits = all_bits[32:32 + msg_length * 8]
    
    # Дополняем до кратности 8 (хотя уже должно быть кратно)
    remainder = len(message_bits) % 8
    if remainder != 0:
        message_bits.extend([0] * (8 - remainder))
    
    if len(message_bits) == 0:
        return ""
    
    bits_string = ''.join(str(b) for b in message_bits)
    try:
        return bits_to_text(bits_string)
    except Exception as e:
        print(f"Ошибка при декодировании: {e}")
        return ""