"""Reed-Solomon error and erasure correction for QR codes.

QR codes use RS codes over GF(256) with generator roots at alpha^0..alpha^(n_ec-1).
(fcr=0: first consecutive root is alpha^0.)

For partial QR codes, unknown (None) codewords are treated as erasures.
RS can correct up to n_ec erasures, or n_ec/2 errors, or a mix where
2*errors + erasures <= n_ec.
"""

from .galois import GF_EXP, gf_div, gf_inverse, gf_mul, gf_poly_eval, gf_poly_mul


def rs_generator_poly(n_ec: int) -> list[int]:
    """Build the RS generator polynomial g(x) = prod_{i=0}^{n_ec-1} (x - alpha^i)."""
    g = [1]
    for i in range(n_ec):
        g = gf_poly_mul(g, [1, GF_EXP[i]])
    return g


def rs_encode(data: list[int], n_ec: int) -> list[int]:
    """Systematic RS encoding: returns data + ec codewords."""
    gen = rs_generator_poly(n_ec)
    msg_out = list(data) + [0] * n_ec
    for i in range(len(data)):
        coeff = msg_out[i]
        if coeff != 0:
            for j in range(1, len(gen)):
                msg_out[i + j] ^= gf_mul(gen[j], coeff)
    return list(data) + msg_out[len(data) :]


def rs_calc_syndromes(msg: list[int], n_ec: int) -> list[int]:
    """Syndromes S_j = msg(alpha^j) for j = 0..n_ec-1."""
    return [gf_poly_eval(msg, GF_EXP[i]) for i in range(n_ec)]


def rs_correct_erasures(msg: list[int], erasure_positions: list[int], n_ec: int) -> list[int]:
    """Correct erasures at known positions using Forney's algorithm.

    All arithmetic uses ascending-power polynomials internally:
      poly[k] is the coefficient of x^k.
    """
    msg_len = len(msg)
    n_erasures = len(erasure_positions)

    if n_erasures == 0:
        return list(msg)
    if n_erasures > n_ec:
        raise ValueError(f"Too many erasures: {n_erasures} > {n_ec}")

    syndromes = rs_calc_syndromes(msg, n_ec)

    # X_i = alpha^(msg_len - 1 - pos) for each erasure position
    X = [GF_EXP[msg_len - 1 - p] for p in erasure_positions]

    # Erasure locator sigma(x) = prod(1 - X_i * x) in ascending power form
    sigma = [1]
    for xi in X:
        new_sigma = [0] * (len(sigma) + 1)
        for j, s in enumerate(sigma):
            new_sigma[j] ^= s
            new_sigma[j + 1] ^= gf_mul(s, xi)
        sigma = new_sigma

    # Error evaluator Omega(x) = S(x) * sigma(x) mod x^n_erasures
    # S(x) = S_0 + S_1*x + S_2*x^2 + ...
    omega = [0] * n_erasures
    for i in range(n_erasures):
        for j in range(min(i + 1, len(sigma))):
            omega[i] ^= gf_mul(syndromes[i - j], sigma[j])

    # Formal derivative sigma'(x) in ascending powers
    # d/dx(x^k) = x^{k-1} if k odd, 0 if k even (in char 2)
    sigma_prime = []
    for k in range(1, len(sigma)):
        if k % 2 == 1:
            sigma_prime.append(sigma[k])
        else:
            sigma_prime.append(0)
    if not sigma_prime:
        sigma_prime = [0]

    def eval_asc(poly: list[int], x: int) -> int:
        """Evaluate ascending-power polynomial at x."""
        result = 0
        x_pow = 1
        for c in poly:
            result ^= gf_mul(c, x_pow)
            x_pow = gf_mul(x_pow, x)
        return result

    # Forney formula: e_i = X_i * Omega(X_i^{-1}) / sigma'(X_i^{-1})
    result = list(msg)
    for idx, pos in enumerate(erasure_positions):
        xi = X[idx]
        xi_inv = gf_inverse(xi)

        omega_val = eval_asc(omega, xi_inv)
        sigma_prime_val = eval_asc(sigma_prime, xi_inv)

        if sigma_prime_val == 0:
            raise ValueError(f"sigma'(X^-1) = 0 at erasure position {pos}")

        magnitude = gf_mul(xi, gf_div(omega_val, sigma_prime_val))
        result[pos] ^= magnitude

    return result


def rs_correct_block(block: list[int | None], n_data: int, n_ec: int) -> list[int | None]:
    """Attempt Reed-Solomon erasure correction on a single block.

    block: list of n_data + n_ec codewords, where None = erasure (unknown).
    Returns corrected block with as many Nones resolved as possible.
    If correction fails, returns the block unchanged.

    If pure erasure correction fails (syndromes non-zero) and there is
    enough EC capacity (n_erasures + 2 <= n_ec), tries treating each
    known position as an additional erasure to find a single-error
    correction.
    """
    msg_len = n_data + n_ec

    erasure_positions: list[int] = []
    msg: list[int] = []
    for i, cw in enumerate(block):
        if cw is None:
            erasure_positions.append(i)
            msg.append(0)
        else:
            msg.append(cw)

    n_erasures = len(erasure_positions)

    if n_erasures > n_ec:
        return list(block)

    if n_erasures == 0:
        syndromes = rs_calc_syndromes(msg, n_ec)
        if all(s == 0 for s in syndromes):
            return list(block)
        return list(block)

    try:
        corrected = rs_correct_erasures(msg, erasure_positions, n_ec)

        check = rs_calc_syndromes(corrected, n_ec)
        if all(s == 0 for s in check):
            result: list[int | None] = list(corrected)
            return result

        # Pure erasure correction failed — there may be errors in "known"
        # positions.  If we have room for one more erasure (the error),
        # try each known position as an extra erasure.
        if n_erasures + 2 <= n_ec:
            erasure_set = set(erasure_positions)
            for suspect in range(msg_len):
                if suspect in erasure_set:
                    continue
                trial_erasures = erasure_positions + [suspect]
                trial_msg = list(msg)
                trial_msg[suspect] = 0
                try:
                    trial_corrected = rs_correct_erasures(trial_msg, trial_erasures, n_ec)
                    trial_check = rs_calc_syndromes(trial_corrected, n_ec)
                    if all(s == 0 for s in trial_check):
                        return list(trial_corrected)
                except (ValueError, ZeroDivisionError):
                    continue

        return list(block)

    except (ValueError, ZeroDivisionError):
        return list(block)


def _check_combos_chunk(
    args: tuple[
        list[int | None],  # block_template
        list[int],  # partial_positions (brute-forced)
        list[list[int]],  # candidate_lists
        list[int],  # combo_sizes (len of each candidate list)
        int,  # start index
        int,  # end index
        int,  # n_data
        int,  # n_ec
        dict[int, frozenset[int]],  # validation_map: pos → allowed values
    ],
) -> list[list[int]]:
    """Check a range of combo indices for valid RS corrections.

    Used by both the sequential path (single chunk covering all combos) and
    the parallel path (multiple chunks split across workers).

    After RS correction succeeds and syndromes are zero, also checks that
    values assigned to non-brute-forced partial positions are consistent
    with their known bits (via validation_map).

    Returns a list of valid (fully-resolved) corrected blocks found in this
    chunk.  Returns up to 2 results — that's enough for the caller to detect
    ambiguity and abort early.
    """
    (
        block_template,
        partial_positions,
        candidate_lists,
        combo_sizes,
        start,
        end,
        n_data,
        n_ec,
        validation_map,
    ) = args

    n_partial = len(combo_sizes)
    valid: list[list[int]] = []

    for flat_idx in range(start, end):
        # Decode flat index → per-list indices (mixed-radix)
        remaining = flat_idx
        combo_indices = [0] * n_partial
        for i in range(n_partial - 1, -1, -1):
            remaining, combo_indices[i] = divmod(remaining, combo_sizes[i])

        trial: list[int | None] = list(block_template)
        for i, pos in enumerate(partial_positions):
            trial[pos] = candidate_lists[i][combo_indices[i]]

        corrected = rs_correct_block(trial, n_data, n_ec)

        if all(c is not None for c in corrected):
            check = rs_calc_syndromes([c for c in corrected], n_ec)  # type: ignore
            if all(s == 0 for s in check):
                # Validate: RS-assigned values for non-brute-forced partials
                # must be consistent with their known bits.
                if validation_map and not all(
                    corrected[pos] in allowed for pos, allowed in validation_map.items()
                ):
                    continue
                valid.append([c for c in corrected if c is not None])
                if len(valid) > 1:
                    return valid

    return valid


def _collect_valid_results(
    block: list[int | None],
    all_valid: list[list[int]],
) -> list[int | None]:
    """Return the unique valid result, or the original block if 0 or 2+ found."""
    if len(all_valid) == 1:
        out: list[int | None] = list(all_valid[0])
        return out
    return list(block)


def rs_correct_block_aggressive(
    block: list[int | None],
    partial_codewords: list,  # list[PartialCodeword] — avoiding circular import
    n_data: int,
    n_ec: int,
    max_brute_bits: int = 10,
    parallel: bool = False,
) -> list[int | None]:
    """Aggressive RS correction using bit-level partial knowledge.

    For blocks with too many codeword-level erasures but where some codewords
    are only missing a few bits, enumerate the possible values for those
    partially-known codewords and check which combinations produce valid
    RS syndromes.

    Strategy:
    1. Separate codewords into: fully known, partially known (few unknown bits),
       and fully unknown.
    2. Sort partials by number of unknown bits (fewest first).
    3. Select the minimum number of partials to brute-force so that the
       remaining partials + fully_unknown ≤ n_ec (RS capacity).
    4. For each candidate assignment of the selected partials, run RS
       erasure correction on the remaining unknowns.
    5. Pick the unique solution that produces zero syndromes.

    max_brute_bits: maximum total unknown bits to enumerate (2^10 = 1024 combos).
                    Each combo runs a full RS correction, so keep this low.
    parallel: if True, distribute combos across all CPU cores using multiprocessing.
    """
    from .data import PartialCodeword  # local import to avoid circular

    msg_len = n_data + n_ec

    fully_known: list[int] = []
    partial_positions: list[int] = []
    fully_unknown: list[int] = []

    for i, pcw in enumerate(partial_codewords):
        if pcw.is_fully_known:
            fully_known.append(i)
        elif pcw.is_fully_unknown:
            fully_unknown.append(i)
        else:
            partial_positions.append(i)

    if len(fully_unknown) > n_ec:
        return list(block)

    total_erasures = len(partial_positions) + len(fully_unknown)

    if total_erasures <= n_ec:
        return rs_correct_block(block, n_data, n_ec)

    if not partial_positions:
        return rs_correct_block(block, n_data, n_ec)

    # Sort partials by unknown bit count (fewest first = cheapest to enumerate)
    partial_positions.sort(key=lambda i: partial_codewords[i].n_unknown_bits)

    # Select minimum partials to brute-force: remaining partials + fully_unknown <= n_ec
    min_brute = total_erasures - n_ec
    assert min_brute > 0

    brute_positions: list[int] = []
    brute_bits = 0
    for i, pos in enumerate(partial_positions):
        if len(brute_positions) >= min_brute and brute_bits > 0:
            break
        ub = partial_codewords[pos].n_unknown_bits
        if brute_bits + ub > max_brute_bits:
            if len(brute_positions) >= min_brute:
                break
            continue
        brute_positions.append(pos)
        brute_bits += ub

    if len(brute_positions) < min_brute:
        return list(block)

    # For brute-forced partials, we'll substitute candidate values per combo.

    candidate_lists: list[list[int]] = [partial_codewords[i].candidates for i in brute_positions]
    combo_sizes = [len(c) for c in candidate_lists]
    total_combos = 1
    for s in combo_sizes:
        total_combos *= s

    # Validation map: non-brute-forced partial codeword candidate values
    brute_set = set(brute_positions)
    validation_map: dict[int, frozenset[int]] = {}
    for pos in partial_positions:
        if pos not in brute_set:
            validation_map[pos] = frozenset(partial_codewords[pos].candidates)

    chunk_args = (list(block), brute_positions, candidate_lists, combo_sizes)

    if parallel and total_combos > 1:
        import math
        import os
        from concurrent.futures import ProcessPoolExecutor, as_completed

        n_workers = os.cpu_count() or 1
        chunk_size = max(1, math.ceil(total_combos / n_workers))

        all_valid: list[list[int]] = []
        with ProcessPoolExecutor(max_workers=n_workers) as executor:
            futures = [
                executor.submit(
                    _check_combos_chunk,
                    (
                        *chunk_args,
                        start,
                        min(start + chunk_size, total_combos),
                        n_data,
                        n_ec,
                        validation_map,
                    ),
                )
                for start in range(0, total_combos, chunk_size)
            ]
            for future in as_completed(futures):
                all_valid.extend(future.result())
                if len(all_valid) > 1:
                    for f in futures:
                        f.cancel()
                    return list(block)
    else:
        all_valid = _check_combos_chunk(
            (*chunk_args, 0, total_combos, n_data, n_ec, validation_map)
        )

    return _collect_valid_results(block, all_valid)
