"""Content-based constraints for narrowing partial QR code data.

When we know (or assume) something about the QR code's content, we can
set additional data bits and reduce the erasure count for Reed-Solomon.

Supported constraints:
- ASCII: assume all payload bytes are printable ASCII (0x20-0x7E).
  Sets MSB=0 for each payload byte (since printable ASCII < 0x80).
- URL prefix: assume the data starts with "http://" or "https://".
  Sets exact bit values for the prefix bytes.
- Padding: after all data segments end, QR codes fill remaining data
  codewords with a deterministic padding sequence: 4-bit terminator
  (0000), byte alignment (0s), then alternating 0xEC and 0x11.
  This constraint is always applied when segment headers are readable.

ASCII and URL constraints assume byte mode (mode indicator 0100).
"""

from typing import Optional

from .data import PartialCodeword, bits_to_partial_codewords
from .stream import CHAR_COUNT_BITS


def _header_bit_count(version: int, mode: int = 0b0100) -> int:
    """Total header bits (mode indicator + character count) for byte mode."""
    if version <= 9:
        vg = 0
    elif version <= 26:
        vg = 1
    else:
        vg = 2
    return 4 + CHAR_COUNT_BITS[mode][vg]


def apply_content_constraints(
    data_partial_cws: list[PartialCodeword],
    version: int,
    assume_ascii: bool = False,
    url_prefix: str | None = None,
) -> tuple[list[PartialCodeword], dict[str, object]]:
    """Apply content constraints to data partial codewords.

    The input should be the concatenated data partial codewords in
    data-stream order (block 0 data, block 1 data, etc.).

    Returns (constrained_codewords, stats).
    stats["conflict"] is set if a constraint contradicts a known bit
    (meaning the assumption is wrong); the original codewords are returned.
    """
    if not assume_ascii and url_prefix is None:
        return list(data_partial_cws), {"bits_set": 0}

    bits: list[Optional[int]] = []
    for pcw in data_partial_cws:
        bits.extend(pcw.known_bits)

    original_known = sum(1 for b in bits if b is not None)
    header_len = _header_bit_count(version)

    def try_set(idx: int, val: int) -> bool:
        """Set a bit if unknown. Returns False on conflict."""
        if idx >= len(bits):
            return True
        if bits[idx] is None:
            bits[idx] = val
            return True
        return bits[idx] == val

    # Mode indicator = 0100 (byte mode)
    for i, b in enumerate([0, 1, 0, 0]):
        if not try_set(i, b):
            return list(data_partial_cws), {"bits_set": 0, "conflict": "mode"}

    if url_prefix is not None:
        for char_idx, byte_val in enumerate(url_prefix.encode("ascii")):
            for bit_pos in range(8):
                flat_idx = header_len + char_idx * 8 + bit_pos
                bit_val = (byte_val >> (7 - bit_pos)) & 1
                if not try_set(flat_idx, bit_val):
                    return list(data_partial_cws), {
                        "bits_set": 0,
                        "conflict": f"prefix_byte_{char_idx}",
                    }

    # ASCII: MSB=0 for each payload byte
    if assume_ascii or url_prefix is not None:
        start_byte = len(url_prefix) if url_prefix else 0
        char_count = _read_known_bits(bits, 4, header_len - 4)
        if char_count is not None:
            max_payload_bytes = char_count
        else:
            max_payload_bytes = (len(bits) - header_len) // 8
        for byte_idx in range(start_byte, max_payload_bytes):
            try_set(header_len + byte_idx * 8, 0)

    result = bits_to_partial_codewords(bits)

    new_known = sum(1 for b in bits if b is not None)
    orig_resolved = sum(1 for p in data_partial_cws if p.is_fully_known)
    new_resolved = sum(1 for p in result if p.is_fully_known)

    stats: dict[str, object] = {
        "bits_set": new_known - original_known,
        "codewords_resolved": new_resolved - orig_resolved,
    }
    if url_prefix:
        stats["prefix"] = url_prefix

    return result, stats


def _compute_segment_data_bits(mode: int, char_count: int) -> int | None:
    """Return number of data bits for a segment with given mode and char count."""
    if mode == 0b0001:  # Numeric
        full_groups, remainder = divmod(char_count, 3)
        data_bits = full_groups * 10
        if remainder == 2:
            data_bits += 7
        elif remainder == 1:
            data_bits += 4
        return data_bits
    elif mode == 0b0010:  # Alphanumeric
        full_pairs, remainder = divmod(char_count, 2)
        return full_pairs * 11 + remainder * 6
    elif mode == 0b0100:  # Byte
        return char_count * 8
    elif mode == 0b1000:  # Kanji
        return char_count * 13
    return None


def _read_known_bits(bits: list[Optional[int]], pos: int, count: int) -> int | None:
    """Read `count` bits from position `pos`. Returns None if any bit is unknown."""
    if pos + count > len(bits):
        return None
    value = 0
    for i in range(count):
        b = bits[pos + i]
        if b is None:
            return None
        value = (value << 1) | b
    return value


def apply_padding_constraints(
    data_partial_cws: list[PartialCodeword],
    version: int,
) -> tuple[list[PartialCodeword], dict[str, object]]:
    """Apply QR padding constraints to data partial codewords.

    After all encoded data segments end, QR codes use a deterministic
    padding sequence:
    1. Terminator: up to 4 zero bits (0000)
    2. Byte alignment: zero bits to reach the next byte boundary
    3. Pad bytes: alternating 0xEC (11101100) and 0x11 (00010001)

    This function parses known segment headers (mode indicator +
    character count) to determine where padding starts, then fills
    in the known padding bits.

    Returns (constrained_codewords, stats).
    """
    bits: list[Optional[int]] = []
    for pcw in data_partial_cws:
        bits.extend(pcw.known_bits)

    total_bits = len(bits)
    original_known = sum(1 for b in bits if b is not None)

    if version <= 9:
        vg = 0
    elif version <= 26:
        vg = 1
    else:
        vg = 2

    # Parse segment headers to find where data ends
    pos = 0
    parsed_segments = 0
    while pos + 4 <= total_bits:
        mode_val = _read_known_bits(bits, pos, 4)
        if mode_val is None:
            if parsed_segments == 0:
                return list(data_partial_cws), {"padding_bits_set": 0}
            # We've parsed at least one segment.  The unknown mode bits
            # might be a terminator (0000) or another segment.  Optimistically
            # try setting them to 0000 — if any conflict, we can't determine
            # padding, but if they accept zeros that's consistent with a
            # terminator.
            can_terminate = True
            for i in range(4):
                if pos + i < total_bits:
                    b = bits[pos + i]
                    if b is not None and b != 0:
                        can_terminate = False
                        break
            if can_terminate:
                # Treat as terminator — break out so we start padding at pos
                break
            return list(data_partial_cws), {"padding_bits_set": 0}

        if mode_val == 0b0000:
            break

        if mode_val not in CHAR_COUNT_BITS:
            return list(data_partial_cws), {"padding_bits_set": 0}

        cc_bit_len = CHAR_COUNT_BITS[mode_val][vg]
        char_count = _read_known_bits(bits, pos + 4, cc_bit_len)
        if char_count is None:
            return list(data_partial_cws), {"padding_bits_set": 0}

        data_bits = _compute_segment_data_bits(mode_val, char_count)
        if data_bits is None:
            return list(data_partial_cws), {"padding_bits_set": 0}

        pos = pos + 4 + cc_bit_len + data_bits
        parsed_segments += 1

    if pos >= total_bits:
        return list(data_partial_cws), {"padding_bits_set": 0}

    def try_set(idx: int, val: int) -> bool:
        if idx >= total_bits:
            return True
        if bits[idx] is None:
            bits[idx] = val
            return True
        return bits[idx] == val

    # 1. Terminator: up to 4 zero bits
    terminator_len = min(4, total_bits - pos)
    for i in range(terminator_len):
        if not try_set(pos + i, 0):
            return list(data_partial_cws), {"padding_bits_set": 0, "conflict": "terminator"}
    pos += terminator_len

    # 2. Byte alignment: zeros to next byte boundary
    byte_boundary = ((pos + 7) // 8) * 8
    while pos < byte_boundary and pos < total_bits:
        if not try_set(pos, 0):
            return list(data_partial_cws), {"padding_bits_set": 0, "conflict": "byte_align"}
        pos += 1

    # 3. Padding bytes: alternating 0xEC and 0x11
    PAD_BYTES = [0xEC, 0x11]
    pad_idx = 0
    while pos + 8 <= total_bits:
        pad_val = PAD_BYTES[pad_idx % 2]
        for bit_pos in range(8):
            bit_val = (pad_val >> (7 - bit_pos)) & 1
            if not try_set(pos + bit_pos, bit_val):
                return list(data_partial_cws), {
                    "padding_bits_set": 0,
                    "conflict": f"pad_byte_{pad_idx}",
                }
        pos += 8
        pad_idx += 1

    result = bits_to_partial_codewords(bits)

    new_known = sum(1 for b in bits if b is not None)
    orig_resolved = sum(1 for p in data_partial_cws if p.is_fully_known)
    new_resolved = sum(1 for p in result if p.is_fully_known)

    return result, {
        "padding_bits_set": new_known - original_known,
        "padding_codewords_resolved": new_resolved - orig_resolved,
    }
