"""QR code data stream decoder.

Parses the data codeword bitstream into mode segments and character data.
Handles partial data gracefully — unknown bits yield None characters.
"""

from dataclasses import dataclass, field
from typing import Optional

# Character count indicator bit lengths per mode and version group
# Version group: 0 = versions 1-9, 1 = versions 10-26, 2 = versions 27-40
CHAR_COUNT_BITS: dict[int, tuple[int, int, int]] = {
    0b0001: (10, 12, 14),  # Numeric
    0b0010: (9, 11, 13),  # Alphanumeric
    0b0100: (8, 16, 16),  # Byte
    0b1000: (8, 10, 12),  # Kanji
    0b0111: (0, 0, 0),  # ECI (special handling)
}

MODE_NAMES: dict[int, str] = {
    0b0001: "Numeric",
    0b0010: "Alphanumeric",
    0b0100: "Byte",
    0b1000: "Kanji",
    0b0111: "ECI",
    0b0000: "Terminator",
}

# Alphanumeric character table (ISO 18004 Table 5)
ALPHANUMERIC_CHARS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ $%*+-./:"


@dataclass
class DataSegment:
    """A decoded segment from the QR data stream."""

    mode: int
    mode_name: str
    char_count: Optional[int]
    data: str  # decoded characters, '?' for unknown


@dataclass
class DecodedData:
    """Result of decoding the QR data stream."""

    segments: list[DataSegment] = field(default_factory=list)
    raw_text: str = ""  # concatenation of all segment data


def _version_group(version: int) -> int:
    """Return version group index for character count lookup."""
    if version <= 9:
        return 0
    elif version <= 26:
        return 1
    else:
        return 2


def _read_bits(bits: list[Optional[int]], pos: int, count: int) -> tuple[Optional[int], int]:
    """Read `count` bits from position `pos`.

    Returns (value, new_pos). Value is None if any bit is unknown.
    """
    if pos + count > len(bits):
        return None, pos + count

    value = 0
    all_known = True
    for i in range(count):
        b = bits[pos + i]
        if b is None:
            all_known = False
            value = value << 1
        else:
            value = (value << 1) | b

    return value if all_known else None, pos + count


def _read_bits_partial(
    bits: list[Optional[int]], pos: int, count: int
) -> tuple[list[Optional[int]], int]:
    """Read `count` bits from position `pos`, preserving None for unknowns."""
    result: list[Optional[int]] = []
    for i in range(count):
        if pos + i < len(bits):
            result.append(bits[pos + i])
        else:
            result.append(None)
    return result, pos + count


def _decode_numeric(bits: list[Optional[int]], pos: int, char_count: int) -> tuple[str, int]:
    """Decode numeric mode data: 3 digits per 10 bits, remainder 1 digit/4 bits or 2 digits/7 bits."""
    result = []
    remaining = char_count

    while remaining >= 3:
        val, pos = _read_bits(bits, pos, 10)
        if val is not None and 0 <= val <= 999:
            result.append(f"{val:03d}")
        else:
            result.append("???")
        remaining -= 3

    if remaining == 2:
        val, pos = _read_bits(bits, pos, 7)
        if val is not None and 0 <= val <= 99:
            result.append(f"{val:02d}")
        else:
            result.append("??")
        remaining = 0

    if remaining == 1:
        val, pos = _read_bits(bits, pos, 4)
        if val is not None and 0 <= val <= 9:
            result.append(str(val))
        else:
            result.append("?")

    return "".join(result), pos


def _decode_alphanumeric(bits: list[Optional[int]], pos: int, char_count: int) -> tuple[str, int]:
    """Decode alphanumeric mode: 2 chars per 11 bits, remainder 1 char/6 bits."""
    result = []
    remaining = char_count

    while remaining >= 2:
        val, pos = _read_bits(bits, pos, 11)
        if val is not None:
            c1_idx = val // 45
            c2_idx = val % 45
            c1 = ALPHANUMERIC_CHARS[c1_idx] if c1_idx < len(ALPHANUMERIC_CHARS) else "?"
            c2 = ALPHANUMERIC_CHARS[c2_idx] if c2_idx < len(ALPHANUMERIC_CHARS) else "?"
            result.append(c1 + c2)
        else:
            result.append("??")
        remaining -= 2

    if remaining == 1:
        val, pos = _read_bits(bits, pos, 6)
        if val is not None and val < len(ALPHANUMERIC_CHARS):
            result.append(ALPHANUMERIC_CHARS[val])
        else:
            result.append("?")

    return "".join(result), pos


def _decode_byte(bits: list[Optional[int]], pos: int, char_count: int) -> tuple[str, int]:
    """Decode byte mode: each character is 8 bits (ISO 8859-1 / UTF-8)."""
    result = []
    for _ in range(char_count):
        val, pos = _read_bits(bits, pos, 8)
        if val is not None:
            if 32 <= val < 127:
                result.append(chr(val))
            else:
                result.append(f"\\x{val:02x}")
        else:
            result.append("?")
    return "".join(result), pos


def decode_data_stream(
    data_codewords: list[Optional[int]],
    version: int,
    data_bits: list[Optional[int]] | None = None,
) -> DecodedData:
    """Parse the QR data stream from (potentially partial) data codewords.

    Converts codewords to bits, then reads mode indicators, character counts,
    and character data. Unknown codewords produce unknown bits.

    If data_bits is provided, it is used directly instead of deriving bits
    from codewords. This preserves partial bit knowledge (e.g. from content
    constraints) that would be lost when a codeword value is None.
    """
    bits: list[Optional[int]]
    if data_bits is not None:
        bits = list(data_bits)
    else:
        bits = []
        for cw in data_codewords:
            if cw is None:
                bits.extend([None] * 8)
            else:
                for shift in range(7, -1, -1):
                    bits.append((cw >> shift) & 1)

    vg = _version_group(version)
    pos = 0
    segments: list[DataSegment] = []

    while pos + 4 <= len(bits):
        mode_val, new_pos = _read_bits(bits, pos, 4)

        if mode_val is None:
            # Mode bits unknown — skip to next byte boundary and try again
            pos = ((new_pos + 7) // 8) * 8
            continue

        if mode_val == 0b0000:
            break

        mode_name = MODE_NAMES.get(mode_val, f"Unknown({mode_val:#06b})")

        if mode_val in CHAR_COUNT_BITS:
            cc_bits = CHAR_COUNT_BITS[mode_val][vg]
        else:
            # Unknown mode
            break

        char_count, new_pos = _read_bits(bits, new_pos, cc_bits)

        if char_count is None:
            # Character count unknown — for byte mode, decode as many
            # bytes as possible, showing partial content (e.g. URL prefix).
            if mode_val == 0b0100:
                max_count = (len(bits) - new_pos) // 8
                data_str, new_pos = _decode_byte(bits, new_pos, max_count)
                trimmed = data_str.rstrip("?")
                if trimmed:
                    data_str = trimmed + "..."
                else:
                    data_str = "[count unknown]"
                segments.append(
                    DataSegment(
                        mode=mode_val,
                        mode_name=mode_name,
                        char_count=None,
                        data=data_str,
                    )
                )
            else:
                segments.append(
                    DataSegment(
                        mode=mode_val,
                        mode_name=mode_name,
                        char_count=None,
                        data="[count unknown]",
                    )
                )
            pos = new_pos
            break  # Can't continue without knowing how many chars to read

        if mode_val == 0b0001:
            data_str, new_pos = _decode_numeric(bits, new_pos, char_count)
        elif mode_val == 0b0010:
            data_str, new_pos = _decode_alphanumeric(bits, new_pos, char_count)
        elif mode_val == 0b0100:
            data_str, new_pos = _decode_byte(bits, new_pos, char_count)
        else:
            data_str = f"[{mode_name} x{char_count}]"
            # Skip the data bits (rough estimate)
            new_pos += char_count * 8

        segments.append(
            DataSegment(
                mode=mode_val,
                mode_name=mode_name,
                char_count=char_count,
                data=data_str,
            )
        )
        pos = new_pos

    raw_text = "".join(seg.data for seg in segments)
    return DecodedData(segments=segments, raw_text=raw_text)
