"""QR code data decoding: bit extraction, block assembly, RS correction, payload parsing."""

from dataclasses import dataclass, field
from typing import Optional

from ...matrix import Cell, Matrix
from ..mask import _get_function_pattern_mask, version_from_size


@dataclass
class BlockInfo:
    """Describes the block structure for a (version, EC level) combination."""

    ec_codewords_per_block: int
    group1_blocks: int
    group1_data_codewords: int
    group2_blocks: int  # 0 if no second group
    group2_data_codewords: int  # 0 if no second group

    @property
    def total_data_codewords(self) -> int:
        return (
            self.group1_blocks * self.group1_data_codewords
            + self.group2_blocks * self.group2_data_codewords
        )

    @property
    def total_blocks(self) -> int:
        return self.group1_blocks + self.group2_blocks

    @property
    def total_codewords(self) -> int:
        return self.total_data_codewords + self.total_blocks * self.ec_codewords_per_block


# Block structure table (ISO 18004 / Thonky's EC table)
BLOCK_TABLE: dict[tuple[int, str], BlockInfo] = {
    (1, "L"): BlockInfo(7, 1, 19, 0, 0),
    (1, "M"): BlockInfo(10, 1, 16, 0, 0),
    (1, "Q"): BlockInfo(13, 1, 13, 0, 0),
    (1, "H"): BlockInfo(17, 1, 9, 0, 0),
    (2, "L"): BlockInfo(10, 1, 34, 0, 0),
    (2, "M"): BlockInfo(16, 1, 28, 0, 0),
    (2, "Q"): BlockInfo(22, 1, 22, 0, 0),
    (2, "H"): BlockInfo(28, 1, 16, 0, 0),
    (3, "L"): BlockInfo(15, 1, 55, 0, 0),
    (3, "M"): BlockInfo(26, 1, 44, 0, 0),
    (3, "Q"): BlockInfo(18, 2, 17, 0, 0),
    (3, "H"): BlockInfo(22, 2, 13, 0, 0),
    (4, "L"): BlockInfo(20, 1, 80, 0, 0),
    (4, "M"): BlockInfo(18, 2, 32, 0, 0),
    (4, "Q"): BlockInfo(26, 2, 24, 0, 0),
    (4, "H"): BlockInfo(16, 4, 9, 0, 0),
    (5, "L"): BlockInfo(26, 1, 108, 0, 0),
    (5, "M"): BlockInfo(24, 2, 43, 0, 0),
    (5, "Q"): BlockInfo(18, 2, 15, 2, 16),
    (5, "H"): BlockInfo(22, 2, 11, 2, 12),
    (6, "L"): BlockInfo(18, 2, 68, 0, 0),
    (6, "M"): BlockInfo(16, 4, 27, 0, 0),
    (6, "Q"): BlockInfo(24, 4, 19, 0, 0),
    (6, "H"): BlockInfo(28, 4, 15, 0, 0),
    (7, "L"): BlockInfo(20, 2, 78, 0, 0),
    (7, "M"): BlockInfo(18, 4, 31, 0, 0),
    (7, "Q"): BlockInfo(18, 2, 14, 4, 15),
    (7, "H"): BlockInfo(26, 4, 13, 1, 14),
    (8, "L"): BlockInfo(24, 2, 97, 0, 0),
    (8, "M"): BlockInfo(22, 2, 38, 2, 39),
    (8, "Q"): BlockInfo(22, 4, 18, 2, 19),
    (8, "H"): BlockInfo(26, 4, 14, 2, 15),
    (9, "L"): BlockInfo(30, 2, 116, 0, 0),
    (9, "M"): BlockInfo(22, 3, 36, 2, 37),
    (9, "Q"): BlockInfo(20, 4, 16, 4, 17),
    (9, "H"): BlockInfo(24, 4, 12, 4, 13),
    (10, "L"): BlockInfo(18, 2, 68, 2, 69),
    (10, "M"): BlockInfo(26, 4, 43, 1, 44),
    (10, "Q"): BlockInfo(24, 6, 19, 2, 20),
    (10, "H"): BlockInfo(28, 6, 15, 2, 16),
    (11, "L"): BlockInfo(20, 4, 81, 0, 0),
    (11, "M"): BlockInfo(30, 1, 50, 4, 51),
    (11, "Q"): BlockInfo(28, 4, 22, 4, 23),
    (11, "H"): BlockInfo(24, 3, 12, 8, 13),
    (12, "L"): BlockInfo(24, 2, 92, 2, 93),
    (12, "M"): BlockInfo(22, 6, 36, 2, 37),
    (12, "Q"): BlockInfo(26, 4, 20, 6, 21),
    (12, "H"): BlockInfo(28, 7, 14, 4, 15),
    (13, "L"): BlockInfo(26, 4, 107, 0, 0),
    (13, "M"): BlockInfo(22, 8, 37, 1, 38),
    (13, "Q"): BlockInfo(24, 8, 20, 4, 21),
    (13, "H"): BlockInfo(22, 12, 11, 4, 12),
    (14, "L"): BlockInfo(30, 3, 115, 1, 116),
    (14, "M"): BlockInfo(24, 4, 40, 5, 41),
    (14, "Q"): BlockInfo(20, 11, 16, 5, 17),
    (14, "H"): BlockInfo(24, 11, 12, 5, 13),
    (15, "L"): BlockInfo(22, 5, 87, 1, 88),
    (15, "M"): BlockInfo(24, 5, 41, 5, 42),
    (15, "Q"): BlockInfo(30, 5, 24, 7, 25),
    (15, "H"): BlockInfo(24, 11, 12, 7, 13),
    (16, "L"): BlockInfo(24, 5, 98, 1, 99),
    (16, "M"): BlockInfo(28, 7, 45, 3, 46),
    (16, "Q"): BlockInfo(24, 15, 19, 2, 20),
    (16, "H"): BlockInfo(30, 3, 15, 13, 16),
    (17, "L"): BlockInfo(28, 1, 107, 5, 108),
    (17, "M"): BlockInfo(28, 10, 46, 1, 47),
    (17, "Q"): BlockInfo(28, 1, 22, 15, 23),
    (17, "H"): BlockInfo(28, 2, 14, 17, 15),
    (18, "L"): BlockInfo(30, 5, 120, 1, 121),
    (18, "M"): BlockInfo(26, 9, 43, 4, 44),
    (18, "Q"): BlockInfo(28, 17, 22, 1, 23),
    (18, "H"): BlockInfo(28, 2, 14, 19, 15),
    (19, "L"): BlockInfo(28, 3, 113, 4, 114),
    (19, "M"): BlockInfo(26, 3, 44, 11, 45),
    (19, "Q"): BlockInfo(26, 17, 21, 4, 22),
    (19, "H"): BlockInfo(26, 9, 13, 16, 14),
    (20, "L"): BlockInfo(28, 3, 107, 5, 108),
    (20, "M"): BlockInfo(26, 3, 41, 13, 42),
    (20, "Q"): BlockInfo(30, 15, 24, 5, 25),
    (20, "H"): BlockInfo(28, 15, 15, 10, 16),
    (21, "L"): BlockInfo(28, 4, 116, 4, 117),
    (21, "M"): BlockInfo(26, 17, 42, 0, 0),
    (21, "Q"): BlockInfo(28, 17, 22, 6, 23),
    (21, "H"): BlockInfo(30, 19, 16, 6, 17),
    (22, "L"): BlockInfo(28, 2, 111, 7, 112),
    (22, "M"): BlockInfo(28, 17, 46, 0, 0),
    (22, "Q"): BlockInfo(30, 7, 24, 16, 25),
    (22, "H"): BlockInfo(24, 34, 13, 0, 0),
    (23, "L"): BlockInfo(30, 4, 121, 5, 122),
    (23, "M"): BlockInfo(28, 4, 47, 14, 48),
    (23, "Q"): BlockInfo(30, 11, 24, 14, 25),
    (23, "H"): BlockInfo(30, 16, 15, 14, 16),
    (24, "L"): BlockInfo(30, 6, 117, 4, 118),
    (24, "M"): BlockInfo(28, 6, 45, 14, 46),
    (24, "Q"): BlockInfo(30, 11, 24, 16, 25),
    (24, "H"): BlockInfo(30, 30, 16, 2, 17),
    (25, "L"): BlockInfo(26, 8, 106, 4, 107),
    (25, "M"): BlockInfo(28, 8, 47, 13, 48),
    (25, "Q"): BlockInfo(30, 7, 24, 22, 25),
    (25, "H"): BlockInfo(30, 22, 15, 13, 16),
    (26, "L"): BlockInfo(28, 10, 114, 2, 115),
    (26, "M"): BlockInfo(28, 19, 46, 4, 47),
    (26, "Q"): BlockInfo(28, 28, 22, 6, 23),
    (26, "H"): BlockInfo(30, 33, 16, 4, 17),
    (27, "L"): BlockInfo(30, 8, 122, 4, 123),
    (27, "M"): BlockInfo(28, 22, 45, 3, 46),
    (27, "Q"): BlockInfo(30, 8, 23, 26, 24),
    (27, "H"): BlockInfo(30, 12, 15, 28, 16),
    (28, "L"): BlockInfo(30, 3, 117, 10, 118),
    (28, "M"): BlockInfo(28, 3, 45, 23, 46),
    (28, "Q"): BlockInfo(30, 4, 24, 31, 25),
    (28, "H"): BlockInfo(30, 11, 15, 31, 16),
    (29, "L"): BlockInfo(30, 7, 116, 7, 117),
    (29, "M"): BlockInfo(28, 21, 45, 7, 46),
    (29, "Q"): BlockInfo(30, 1, 23, 37, 24),
    (29, "H"): BlockInfo(30, 19, 15, 26, 16),
    (30, "L"): BlockInfo(30, 5, 115, 10, 116),
    (30, "M"): BlockInfo(28, 19, 47, 10, 48),
    (30, "Q"): BlockInfo(30, 15, 24, 25, 25),
    (30, "H"): BlockInfo(30, 23, 15, 25, 16),
    (31, "L"): BlockInfo(30, 13, 115, 3, 116),
    (31, "M"): BlockInfo(28, 2, 46, 29, 47),
    (31, "Q"): BlockInfo(30, 42, 24, 1, 25),
    (31, "H"): BlockInfo(30, 23, 15, 28, 16),
    (32, "L"): BlockInfo(30, 17, 115, 0, 0),
    (32, "M"): BlockInfo(28, 10, 46, 23, 47),
    (32, "Q"): BlockInfo(30, 10, 24, 35, 25),
    (32, "H"): BlockInfo(30, 19, 15, 35, 16),
    (33, "L"): BlockInfo(30, 17, 115, 1, 116),
    (33, "M"): BlockInfo(28, 14, 46, 21, 47),
    (33, "Q"): BlockInfo(30, 29, 24, 19, 25),
    (33, "H"): BlockInfo(30, 11, 15, 46, 16),
    (34, "L"): BlockInfo(30, 13, 115, 6, 116),
    (34, "M"): BlockInfo(28, 14, 46, 23, 47),
    (34, "Q"): BlockInfo(30, 44, 24, 7, 25),
    (34, "H"): BlockInfo(30, 59, 16, 1, 17),
    (35, "L"): BlockInfo(30, 12, 121, 7, 122),
    (35, "M"): BlockInfo(28, 12, 47, 26, 48),
    (35, "Q"): BlockInfo(30, 39, 24, 14, 25),
    (35, "H"): BlockInfo(30, 22, 15, 41, 16),
    (36, "L"): BlockInfo(30, 6, 121, 14, 122),
    (36, "M"): BlockInfo(28, 6, 47, 34, 48),
    (36, "Q"): BlockInfo(30, 46, 24, 10, 25),
    (36, "H"): BlockInfo(30, 2, 15, 64, 16),
    (37, "L"): BlockInfo(30, 17, 122, 4, 123),
    (37, "M"): BlockInfo(28, 29, 46, 14, 47),
    (37, "Q"): BlockInfo(30, 49, 24, 10, 25),
    (37, "H"): BlockInfo(30, 24, 15, 46, 16),
    (38, "L"): BlockInfo(30, 4, 122, 18, 123),
    (38, "M"): BlockInfo(28, 13, 46, 32, 47),
    (38, "Q"): BlockInfo(30, 48, 24, 14, 25),
    (38, "H"): BlockInfo(30, 42, 15, 32, 16),
    (39, "L"): BlockInfo(30, 20, 117, 4, 118),
    (39, "M"): BlockInfo(28, 40, 47, 7, 48),
    (39, "Q"): BlockInfo(30, 43, 24, 22, 25),
    (39, "H"): BlockInfo(30, 10, 15, 67, 16),
    (40, "L"): BlockInfo(30, 19, 118, 6, 119),
    (40, "M"): BlockInfo(28, 18, 47, 31, 48),
    (40, "Q"): BlockInfo(30, 34, 24, 34, 25),
    (40, "H"): BlockInfo(30, 20, 15, 61, 16),
}


def data_coordinates(size: int) -> list[tuple[int, int]]:
    """Generate the sequence of (row, col) coordinates for data module placement.

    QR codes place data in 2-module-wide columns, starting from the bottom-right
    corner, moving upward. Direction alternates between up and down for each
    column pair. Column 6 (vertical timing pattern) is skipped entirely.

    Returns a list of (row, col) tuples in the order bits are placed.
    """
    coords: list[tuple[int, int]] = []
    col = size - 1
    upward = True

    while col >= 0:
        if col == 6:
            col -= 1
            continue

        rows = range(size - 1, -1, -1) if upward else range(size)
        for row in rows:
            for c in (col, col - 1):
                if 0 <= c < size:
                    coords.append((row, c))

        col -= 2
        upward = not upward

    return coords


def extract_data_bits(matrix: Matrix, version: int) -> list[Optional[int]]:
    """Extract data bits from the matrix in zigzag placement order.

    Function pattern modules are skipped. Each returned element is:
    - 1 for FILLED
    - 0 for EMPTY
    - None for UNKNOWN

    The result length equals the total number of data+EC bits for this QR version.
    """
    size = matrix.height
    func_mask = _get_function_pattern_mask(version, size)
    coords = data_coordinates(size)

    bits: list[Optional[int]] = []
    for row, col in coords:
        if func_mask[row][col]:
            continue
        cell = matrix.cells[row][col]
        if cell == Cell.FILLED:
            bits.append(1)
        elif cell == Cell.EMPTY:
            bits.append(0)
        else:
            bits.append(None)

    return bits


def write_data_bits(matrix: Matrix, version: int, bits: list[Optional[int]]) -> None:
    """Write data bits into the matrix in zigzag placement order (in-place).

    Inverse of extract_data_bits: places bits into data module positions,
    skipping function pattern modules.  None bits are written as UNKNOWN.
    """
    size = matrix.height
    func_mask = _get_function_pattern_mask(version, size)
    coords = data_coordinates(size)

    bit_idx = 0
    for row, col in coords:
        if func_mask[row][col]:
            continue
        if bit_idx < len(bits):
            b = bits[bit_idx]
            if b is None:
                matrix.cells[row][col] = Cell.UNKNOWN
            else:
                matrix.cells[row][col] = Cell.FILLED if b else Cell.EMPTY
            bit_idx += 1


def bits_to_codewords(bits: list[Optional[int]]) -> list[Optional[int]]:
    """Group a flat bit list into 8-bit codewords.

    Each codeword is an int 0-255, or None if any of its 8 bits is None.
    """
    codewords: list[Optional[int]] = []
    for i in range(0, len(bits), 8):
        byte_bits = bits[i : i + 8]
        if len(byte_bits) < 8:
            break
        if any(b is None for b in byte_bits):
            codewords.append(None)
        else:
            value = 0
            for b in byte_bits:
                value = (value << 1) | b  # type: ignore[operator]
            codewords.append(value)
    return codewords


@dataclass
class PartialCodeword:
    """A codeword that may be partially known at the bit level.

    known_bits: 8-element list, each is 0, 1, or None.
    value: resolved int if all bits known, else None.
    candidates: list of possible byte values consistent with known bits.
    """

    known_bits: list[Optional[int]]
    value: Optional[int] = None
    candidates: list[int] = field(default_factory=list)

    @property
    def n_unknown_bits(self) -> int:
        return sum(1 for b in self.known_bits if b is None)

    @property
    def is_fully_known(self) -> bool:
        return self.value is not None

    @property
    def is_fully_unknown(self) -> bool:
        return all(b is None for b in self.known_bits)


def bits_to_partial_codewords(bits: list[Optional[int]]) -> list[PartialCodeword]:
    """Group bits into PartialCodewords, preserving bit-level knowledge."""
    result: list[PartialCodeword] = []
    for i in range(0, len(bits), 8):
        byte_bits = bits[i : i + 8]
        if len(byte_bits) < 8:
            break
        known = list(byte_bits)
        unknown_positions = [j for j, b in enumerate(known) if b is None]

        if not unknown_positions:
            value = 0
            for b in known:
                value = (value << 1) | b  # type: ignore[operator]
            result.append(PartialCodeword(known_bits=known, value=value, candidates=[value]))
        elif len(unknown_positions) == 8:
            result.append(
                PartialCodeword(known_bits=known, value=None, candidates=list(range(256)))
            )
        else:
            base = 0
            for b in known:
                base = (base << 1) | (b if b is not None else 0)

            candidates: list[int] = []
            for combo in range(1 << len(unknown_positions)):
                val = base
                for bit_idx, pos in enumerate(unknown_positions):
                    if (combo >> (len(unknown_positions) - 1 - bit_idx)) & 1:
                        val |= 1 << (7 - pos)
                candidates.append(val)

            result.append(PartialCodeword(known_bits=known, value=None, candidates=candidates))

    return result


def deinterleave_blocks(
    codewords: list[Optional[int]], block_info: BlockInfo
) -> list[list[Optional[int]]]:
    """De-interleave the codeword stream back into individual RS blocks.

    QR codes interleave data codewords round-robin across all blocks,
    then interleave EC codewords round-robin across all blocks.

    Returns a list of blocks, where each block is [data_cw..., ec_cw...].
    """
    n_blocks = block_info.total_blocks

    data_counts: list[int] = []
    for _ in range(block_info.group1_blocks):
        data_counts.append(block_info.group1_data_codewords)
    for _ in range(block_info.group2_blocks):
        data_counts.append(block_info.group2_data_codewords)

    ec_count = block_info.ec_codewords_per_block
    max_data = max(data_counts)

    # --- De-interleave data codewords ---
    data_per_block: list[list[Optional[int]]] = [[] for _ in range(n_blocks)]
    idx = 0
    for col_i in range(max_data):
        for block_i in range(n_blocks):
            if col_i < data_counts[block_i]:
                data_per_block[block_i].append(codewords[idx] if idx < len(codewords) else None)
                idx += 1

    # --- De-interleave EC codewords ---
    ec_per_block: list[list[Optional[int]]] = [[] for _ in range(n_blocks)]
    for col_i in range(ec_count):
        for block_i in range(n_blocks):
            ec_per_block[block_i].append(codewords[idx] if idx < len(codewords) else None)
            idx += 1

    blocks: list[list[Optional[int]]] = []
    for block_i in range(n_blocks):
        blocks.append(data_per_block[block_i] + ec_per_block[block_i])

    return blocks
