"""Top-level QR code data decoder.

Orchestrates: bit extraction -> codeword assembly -> de-interleaving ->
content constraints (optional) -> RS error correction -> data stream parsing.
"""

from dataclasses import dataclass, field
from typing import Optional

from ...matrix import Cell, Matrix
from ..formatinfo import EC_REVERSE_MAP, FORMAT_MASK, _bch_encode
from ..mask import ALIGNMENT_PATTERN_POSITIONS, apply_mask, version_from_size
from ..version import _bch_encode_version
from .constraints import apply_content_constraints, apply_padding_constraints
from .data import (
    BLOCK_TABLE,
    BlockInfo,
    PartialCodeword,
    bits_to_partial_codewords,
    extract_data_bits,
    write_data_bits,
)
from .reed_solomon import rs_correct_block, rs_correct_block_aggressive
from .stream import DecodedData, decode_data_stream


@dataclass
class DecodeResult:
    """Full result of QR data decoding."""

    version: int
    ec_level: str
    block_info: BlockInfo
    total_data_codewords: int
    known_data_codewords: int
    total_ec_codewords: int
    known_ec_codewords: int
    blocks_corrected: int  # number of blocks successfully RS-corrected
    decoded: DecodedData
    data_codewords: list[Optional[int]] = field(default_factory=list)
    ec_codewords: list[Optional[int]] = field(default_factory=list)
    constraint_stats: dict[str, object] = field(default_factory=dict)

    @property
    def fully_corrected(self) -> bool:
        """True if all blocks were RS-corrected (no remaining unknowns)."""
        return all(c is not None for c in self.data_codewords) and all(
            c is not None for c in self.ec_codewords
        )


def _deinterleave_partial(
    partial_cws: list[PartialCodeword], block_info: BlockInfo
) -> list[list[PartialCodeword]]:
    """De-interleave partial codewords into RS blocks (same logic as deinterleave_blocks)."""
    n_blocks = block_info.total_blocks

    data_counts: list[int] = _data_counts(block_info)

    ec_count = block_info.ec_codewords_per_block
    max_data = max(data_counts)

    dummy = PartialCodeword(known_bits=[None] * 8, value=None, candidates=list(range(256)))

    data_per_block: list[list[PartialCodeword]] = [[] 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(
                    partial_cws[idx] if idx < len(partial_cws) else dummy
                )
                idx += 1

    ec_per_block: list[list[PartialCodeword]] = [[] 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(partial_cws[idx] if idx < len(partial_cws) else dummy)
            idx += 1

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

    return blocks


def _data_counts(block_info: BlockInfo) -> list[int]:
    """Per-block data codeword counts."""
    return [block_info.group1_data_codewords] * block_info.group1_blocks + [
        block_info.group2_data_codewords
    ] * block_info.group2_blocks


def _extract_data_partial_cws(
    partial_blocks: list[list[PartialCodeword]],
    block_info: BlockInfo,
) -> list[PartialCodeword]:
    """Extract data partial codewords from blocks in data-stream order."""
    counts = _data_counts(block_info)
    result: list[PartialCodeword] = []
    for i, block in enumerate(partial_blocks):
        result.extend(block[: counts[i]])
    return result


def _replace_data_in_blocks(
    partial_blocks: list[list[PartialCodeword]],
    constrained_data: list[PartialCodeword],
    block_info: BlockInfo,
) -> list[list[PartialCodeword]]:
    """Replace data codewords in blocks with constrained versions."""
    counts = _data_counts(block_info)
    new_blocks: list[list[PartialCodeword]] = []
    idx = 0
    for i, block in enumerate(partial_blocks):
        n = counts[i]
        new_blocks.append(list(constrained_data[idx : idx + n]) + block[n:])
        idx += n
    return new_blocks


def _rs_correct_all(
    partial_blocks: list[list[PartialCodeword]],
    block_info: BlockInfo,
    n_ec: int,
    max_brute_bits: int = 10,
    parallel: bool = False,
) -> tuple[list[list[Optional[int]]], int]:
    """Run RS correction on all blocks. Returns (corrected_blocks, blocks_corrected)."""
    counts = _data_counts(block_info)
    corrected_blocks: list[list[Optional[int]]] = []
    blocks_corrected = 0

    for i, pblock in enumerate(partial_blocks):
        n_data = counts[i]
        block: list[Optional[int]] = [pcw.value for pcw in pblock]

        corrected = rs_correct_block(block, n_data, n_ec)
        orig_unknowns = sum(1 for c in block if c is None)
        new_unknowns = sum(1 for c in corrected if c is None)

        if new_unknowns > 0:
            aggressive = rs_correct_block_aggressive(
                block, pblock, n_data, n_ec, max_brute_bits=max_brute_bits, parallel=parallel
            )
            agg_unknowns = sum(1 for c in aggressive if c is None)
            if agg_unknowns < new_unknowns:
                corrected = aggressive
                new_unknowns = agg_unknowns

        if new_unknowns < orig_unknowns:
            blocks_corrected += 1

        corrected_blocks.append(corrected)

    return corrected_blocks, blocks_corrected


def _build_result(
    corrected_blocks: list[list[Optional[int]]],
    blocks_corrected: int,
    version: int,
    ec_level: str,
    block_info: BlockInfo,
    n_ec: int,
    constraint_stats: dict[str, object] | None = None,
    partial_blocks: list[list[PartialCodeword]] | None = None,
) -> DecodeResult:
    """Assemble a DecodeResult from corrected blocks.

    If partial_blocks is provided, constrained bit-level knowledge is
    preserved for data codewords that RS couldn't fully resolve.
    """
    counts = _data_counts(block_info)

    data_codewords: list[Optional[int]] = []
    ec_codewords: list[Optional[int]] = []
    for i, block in enumerate(corrected_blocks):
        n = counts[i]
        data_codewords.extend(block[:n])
        ec_codewords.extend(block[n:])

    data_bits: list[Optional[int]] | None = None
    if partial_blocks is not None:
        data_pcws = _extract_data_partial_cws(partial_blocks, block_info)
        data_bits = []
        for i, cw in enumerate(data_codewords):
            if cw is not None:
                for shift in range(7, -1, -1):
                    data_bits.append((cw >> shift) & 1)
            elif i < len(data_pcws):
                data_bits.extend(data_pcws[i].known_bits)
            else:
                data_bits.extend([None] * 8)

    decoded = decode_data_stream(data_codewords, version, data_bits=data_bits)

    return DecodeResult(
        version=version,
        ec_level=ec_level,
        block_info=block_info,
        total_data_codewords=block_info.total_data_codewords,
        known_data_codewords=sum(1 for c in data_codewords if c is not None),
        total_ec_codewords=block_info.total_blocks * n_ec,
        known_ec_codewords=sum(1 for c in ec_codewords if c is not None),
        blocks_corrected=blocks_corrected,
        decoded=decoded,
        data_codewords=data_codewords,
        ec_codewords=ec_codewords,
        constraint_stats=constraint_stats or {},
    )


def _reinterleave(
    data_codewords: list[Optional[int]],
    ec_codewords: list[Optional[int]],
    block_info: BlockInfo,
) -> list[Optional[int]]:
    """Re-interleave data and EC codewords into the QR codeword stream order.

    Reverses the de-interleaving: arranges codewords round-robin across blocks,
    data first then EC, matching the order they appear in the QR data area.
    """
    counts = _data_counts(block_info)
    n_blocks = block_info.total_blocks
    n_ec = block_info.ec_codewords_per_block
    max_data = max(counts)

    data_blocks: list[list[Optional[int]]] = []
    idx = 0
    for n in counts:
        data_blocks.append(list(data_codewords[idx : idx + n]))
        idx += n

    ec_blocks: list[list[Optional[int]]] = []
    idx = 0
    for _ in range(n_blocks):
        ec_blocks.append(list(ec_codewords[idx : idx + n_ec]))
        idx += n_ec

    interleaved: list[Optional[int]] = []
    for col_i in range(max_data):
        for block_i in range(n_blocks):
            if col_i < counts[block_i]:
                interleaved.append(data_blocks[block_i][col_i])

    for col_i in range(n_ec):
        for block_i in range(n_blocks):
            interleaved.append(ec_blocks[block_i][col_i])

    return interleaved


def _draw_finder_pattern(matrix: Matrix, top: int, left: int) -> None:
    """Draw a 7x7 finder pattern with top-left corner at (top, left)."""
    size = matrix.height
    for r in range(7):
        for c in range(7):
            rr, cc = top + r, left + c
            if 0 <= rr < size and 0 <= cc < size:
                # Outer border, inner 3x3 block → filled; middle ring → empty
                if r in (0, 6) or c in (0, 6) or (2 <= r <= 4 and 2 <= c <= 4):
                    matrix.cells[rr][cc] = Cell.FILLED
                else:
                    matrix.cells[rr][cc] = Cell.EMPTY


def _draw_separator(matrix: Matrix, top: int, left: int, width: int, height: int) -> None:
    """Draw an all-EMPTY rectangle (separator around finder patterns)."""
    size = matrix.height
    for r in range(height):
        for c in range(width):
            rr, cc = top + r, left + c
            if 0 <= rr < size and 0 <= cc < size:
                matrix.cells[rr][cc] = Cell.EMPTY


def _draw_alignment_pattern(matrix: Matrix, center_row: int, center_col: int) -> None:
    """Draw a 5x5 alignment pattern centered at (center_row, center_col)."""
    for dr in range(-2, 3):
        for dc in range(-2, 3):
            r, c = center_row + dr, center_col + dc
            if abs(dr) == 2 or abs(dc) == 2 or (dr == 0 and dc == 0):
                matrix.cells[r][c] = Cell.FILLED
            else:
                matrix.cells[r][c] = Cell.EMPTY


def _draw_function_patterns(
    matrix: Matrix, version: int, ec_level: str, mask_pattern: int
) -> None:
    """Draw all function patterns onto the matrix (in-place).

    Draws finder patterns, separators, timing patterns, alignment patterns,
    dark module, format information, and version information.
    """
    size = matrix.height

    _draw_finder_pattern(matrix, 0, 0)
    _draw_finder_pattern(matrix, 0, size - 7)
    _draw_finder_pattern(matrix, size - 7, 0)

    _draw_separator(matrix, 0, 7, 1, 8)
    _draw_separator(matrix, 7, 0, 8, 1)
    _draw_separator(matrix, 0, size - 8, 1, 8)
    _draw_separator(matrix, 7, size - 8, 8, 1)
    _draw_separator(matrix, size - 8, 7, 1, 8)
    _draw_separator(matrix, size - 8, 0, 8, 1)

    for i in range(8, size - 8):
        val = Cell.FILLED if i % 2 == 0 else Cell.EMPTY
        matrix.cells[6][i] = val
        matrix.cells[i][6] = val

    dark_row = 4 * version + 9
    if dark_row < size:
        matrix.cells[dark_row][8] = Cell.FILLED

    if version in ALIGNMENT_PATTERN_POSITIONS:
        positions = ALIGNMENT_PATTERN_POSITIONS[version]
        for row_center in positions:
            for col_center in positions:
                if row_center <= 8 and col_center <= 8:
                    continue
                if row_center <= 8 and col_center >= size - 8:
                    continue
                if row_center >= size - 8 and col_center <= 8:
                    continue
                _draw_alignment_pattern(matrix, row_center, col_center)

    ec_bits = EC_REVERSE_MAP[ec_level]
    data5 = (ec_bits << 3) | mask_pattern
    format_bits = _bch_encode(data5) ^ FORMAT_MASK

    # Format info location 1: around top-left finder
    # Bits 0-7: row 8, columns 0-5, 7, 8 (skip col 6)
    tl_cols_row8 = [0, 1, 2, 3, 4, 5, 7, 8]
    for i, c in enumerate(tl_cols_row8):
        bit = (format_bits >> (14 - i)) & 1
        matrix.cells[8][c] = Cell.FILLED if bit else Cell.EMPTY
    # Bits 8-14: column 8, rows 7, 5-0 (skip row 6)
    tl_rows_col8 = [7, 5, 4, 3, 2, 1, 0]
    for i, r in enumerate(tl_rows_col8):
        bit = (format_bits >> (14 - 8 - i)) & 1
        matrix.cells[r][8] = Cell.FILLED if bit else Cell.EMPTY

    # Format info location 2: bottom-left + top-right
    for i in range(7):
        bit = (format_bits >> (14 - i)) & 1
        matrix.cells[size - 1 - i][8] = Cell.FILLED if bit else Cell.EMPTY
    for i in range(8):
        bit = (format_bits >> (14 - 7 - i)) & 1
        matrix.cells[8][size - 8 + i] = Cell.FILLED if bit else Cell.EMPTY

    if version >= 7:
        version_bits = _bch_encode_version(version)
        for i in range(18):
            bit = (version_bits >> i) & 1
            col = i // 3
            row = size - 11 + (i % 3)
            matrix.cells[row][col] = Cell.FILLED if bit else Cell.EMPTY
        for i in range(18):
            bit = (version_bits >> i) & 1
            row = i // 3
            col = size - 11 + (i % 3)
            matrix.cells[row][col] = Cell.FILLED if bit else Cell.EMPTY


def reconstruct_matrix(
    original: Matrix,
    result: DecodeResult,
    mask_pattern: int,
) -> Matrix:
    """Reconstruct the QR matrix by writing recovered data back.

    Takes the original (pre-mask) matrix and writes all corrected codewords
    into the data area, then re-applies the mask pattern to produce the
    visual QR code.  Also draws all function patterns (finder, alignment,
    timing, format info, version info, dark module).

    Unknown codewords are written as UNKNOWN (gray) modules.
    """
    version = result.version
    size = original.height

    interleaved = _reinterleave(result.data_codewords, result.ec_codewords, result.block_info)

    bits: list[Optional[int]] = []
    for cw in interleaved:
        if cw is None:
            bits.extend([None] * 8)
        else:
            for shift in range(7, -1, -1):
                bits.append((cw >> shift) & 1)

    reconstructed = Matrix(size, size)
    for row in range(size):
        for col in range(size):
            reconstructed.cells[row][col] = original.cells[row][col]
    write_data_bits(reconstructed, version, bits)

    reconstructed = apply_mask(reconstructed, mask_pattern, version)

    _draw_function_patterns(reconstructed, version, result.ec_level, mask_pattern)

    return reconstructed


def matrix_to_image(
    matrix: Matrix,
    output_path: str,
    module_size: int = 10,
    border_modules: int = 4,
) -> None:
    """Render a QR matrix to a PNG image using Pillow.

    Args:
        matrix: The reconstructed QR matrix.
        output_path: Destination file path (e.g. 'output.png').
        module_size: Pixel size of each QR module.
        border_modules: Number of quiet-zone modules around the code.
    """
    from PIL import Image

    size = matrix.height
    img_size = (size + 2 * border_modules) * module_size
    img = Image.new("RGB", (img_size, img_size), (255, 255, 255))
    pixels = img.load()
    assert pixels is not None

    for row in range(size):
        for col in range(size):
            cell = matrix.cells[row][col]
            if cell == Cell.FILLED:
                color = (0, 0, 0)
            elif cell == Cell.EMPTY:
                color = (255, 255, 255)
            else:
                color = (128, 128, 128)  # UNKNOWN → gray

            px_row = (row + border_modules) * module_size
            px_col = (col + border_modules) * module_size
            for dr in range(module_size):
                for dc in range(module_size):
                    pixels[px_col + dc, px_row + dr] = color

    img.save(output_path)


def decode_qr_data(
    matrix: Matrix,
    version: int | None,
    ec_level: str,
    assume_ascii: bool = False,
    assume_url: bool = False,
    max_brute_bits: int = 10,
    parallel: bool = False,
) -> DecodeResult:
    """Decode QR code data from a (potentially partial) matrix.

    The matrix should already have the mask pattern applied (unmasked).

    Args:
        matrix: The unmasked QR matrix.
        version: QR version (1-40), or None to guess from matrix size.
        ec_level: Error correction level ("L", "M", "Q", or "H").
        assume_ascii: Assume payload is printable ASCII (sets MSB=0 per byte).
        assume_url: Assume data starts with http:// or https://.
        max_brute_bits: Max unknown bits to brute-force per RS block.
            Each extra bit doubles the work (2^N combinations).
            Default 10 (~1024 combos). Set higher for deep recovery.
        parallel: Distribute brute-force across all CPU cores.

    Returns:
        DecodeResult with decoded segments and statistics.
    """

    size = matrix.height
    if version is None:
        version = version_from_size(size)

    key = (version, ec_level)
    if key not in BLOCK_TABLE:
        raise ValueError(f"Unknown version/EC combination: {version}-{ec_level}")

    block_info = BLOCK_TABLE[key]
    n_ec = block_info.ec_codewords_per_block

    # Step 1: Extract data bits and build partial codewords
    bits = extract_data_bits(matrix, version)
    partial_cws = bits_to_partial_codewords(bits)
    partial_blocks = _deinterleave_partial(partial_cws, block_info)

    configs: list[dict[str, object]] = []
    if assume_url:
        configs.append({"ascii": True, "prefix": "https://"})
        configs.append({"ascii": True, "prefix": "http://"})
        configs.append({"ascii": True, "prefix": None})  # ASCII-only fallback
    elif assume_ascii:
        configs.append({"ascii": True, "prefix": None})
    configs.append({"ascii": False, "prefix": None})  # unconstrained fallback

    best_result: DecodeResult | None = None

    for cfg in configs:
        attempt_blocks = [list(b) for b in partial_blocks]

        cstats: dict[str, object] = {}
        if cfg["ascii"] or cfg["prefix"]:
            data_pcws = _extract_data_partial_cws(attempt_blocks, block_info)
            constrained, cstats = apply_content_constraints(
                data_pcws,
                version,
                assume_ascii=bool(cfg["ascii"]),
                url_prefix=str(cfg["prefix"]) if cfg["prefix"] else None,
            )
            if cstats.get("conflict"):
                continue
            attempt_blocks = _replace_data_in_blocks(attempt_blocks, constrained, block_info)

        # Apply padding constraints (always — deterministic after all data segments end)
        data_pcws_for_pad = _extract_data_partial_cws(attempt_blocks, block_info)
        padded, pad_stats = apply_padding_constraints(data_pcws_for_pad, version)
        if not pad_stats.get("conflict"):
            attempt_blocks = _replace_data_in_blocks(attempt_blocks, padded, block_info)
            cstats.update(pad_stats)

        corrected, n_corrected = _rs_correct_all(
            attempt_blocks, block_info, n_ec, max_brute_bits=max_brute_bits, parallel=parallel
        )
        result = _build_result(
            corrected,
            n_corrected,
            version,
            ec_level,
            block_info,
            n_ec,
            constraint_stats=cstats,
            partial_blocks=attempt_blocks,
        )

        if best_result is None or result.blocks_corrected > best_result.blocks_corrected:
            best_result = result

        if result.blocks_corrected == block_info.total_blocks:
            break

    assert best_result is not None
    return best_result
