from starkware.cairo.common.bool import FALSE
from starkware.cairo.common.builtin_poseidon.poseidon import poseidon_hash, poseidon_hash_many
from starkware.cairo.common.cairo_builtins import PoseidonBuiltin
from starkware.cairo.common.math import assert_le, assert_lt, assert_nn_le
from starkware.cairo.common.uint256 import Uint256
from starkware.starknet.core.os.data_availability.bls_field import (
    BigInt3,
    bigint3_to_uint256,
    felt_to_bigint3,
    horner_eval,
    reduced_mul,
)

const BLOB_LENGTH = 4096;

// Represents an integer in the range [0, 2^384).
struct Uint384 {
    // The low 192 bits of the value.
    low: felt,
    // The high 192 bits of the value.
    high: felt,
}

// Data required for KZG commitment on data-availability.
struct OsKzgCommitmentInfo {
    // Generated by hashing the Stark-friendly commitment with the KZG commitments.
    z: felt,
    // Number of blobs the data is divided to.
    n_blobs: felt,
    // KZG commitment per blob.
    kzg_commitments: Uint384*,
    // The value of the data polynomial at `z` per blob.
    evals: Uint256*,
}

// Computes information needed for publishing data on L1 using KZG commitment.
//
// Specifically, divides the data into blobs of size <= BLOB_LENGTH.
// For each blob, let `p` be the data polynomial. The OS samples `p` at a point `z`,
// where `z` is generated using Fiat-Shamir heuristic.
//
// The OS uses a Stark-friendly commitment scheme, while the L1 side uses KZG.
// This requires combining them both in the generation of `z`; i.e.,
// the OS guesses (and outputs):
//   * kzg_commitments = [KZG(blob) for blob in blobs],
// computes:
//   * z = Poseidon(Poseidon(data), Poseidon(kzg_commitments)),
// and samples the polynomials:
//   * evals = [p(z) for p in blobs_p].
//
// The L1 side will need to verify that `kzg_commitment` is correct
// and that `p(z) = y` (using a KZG proof) for each blob.
func compute_os_kzg_commitment_info{range_check_ptr, poseidon_ptr: PoseidonBuiltin*}(
    state_updates_start: felt*, state_updates_end: felt*
) -> OsKzgCommitmentInfo* {
    alloc_locals;
    local da_size = state_updates_end - state_updates_start;
    local n_blobs;
    local kzg_commitments: Uint384*;
    local evals: Uint256*;
    %{
        import itertools

        from starkware.python.utils import blockify

        kzg_manager.store_da_segment(
            da_segment=memory.get_range_as_ints(addr=ids.state_updates_start, size=ids.da_size)
        )
        kzg_commitments = [
            kzg_manager.polynomial_coefficients_to_kzg_commitment_callback(chunk)
            for chunk in blockify(kzg_manager.da_segment, chunk_size=ids.BLOB_LENGTH)
        ]

        ids.n_blobs = len(kzg_commitments)
        ids.kzg_commitments = segments.add_temp_segment()
        ids.evals = segments.add_temp_segment()

        segments.write_arg(ids.kzg_commitments.address_, list(itertools.chain(*kzg_commitments)))
    %}

    // Note: the number of blobs per L1 transaction is limited, and should be checked ouside
    // of this program.
    with_attr error_message("Invalid number of blobs.") {
        assert_le(1, n_blobs);
    }

    // Calculate `z`.
    let (da_hash) = poseidon_hash_many(n=da_size, elements=state_updates_start);
    let (kzg_commitment_hash) = poseidon_hash_many(
        n=n_blobs * Uint384.SIZE, elements=cast(kzg_commitments, felt*)
    );
    let (z) = poseidon_hash(da_hash, kzg_commitment_hash);
    let z_bigint = felt_to_bigint3(z);

    // Evaluate each blob polynomial at `z`.
    eval_in_chunks(
        z=z_bigint, n_blobs=n_blobs, data_size=da_size, data=state_updates_start, evals=evals
    );

    tempvar os_kzg_commitment_info = new OsKzgCommitmentInfo(
        z=z, n_blobs=n_blobs, kzg_commitments=kzg_commitments, evals=evals
    );
    return os_kzg_commitment_info;
}

// Splits the data into chunks of size BLOB_LENGTH, and evaluates the induced
// polynomials at the given point `z`.
// Writes the results to `evals`.
func eval_in_chunks{range_check_ptr}(
    z: BigInt3, n_blobs: felt, data_size: felt, data: felt*, evals: Uint256*
) {
    if (n_blobs == 1) {
        assert_nn_le(data_size, BLOB_LENGTH);
        let res = eval_uint256(z=z, n_coefficients=data_size, coefficients=data);
        assert evals[0] = res;
        return ();
    }

    assert_lt(BLOB_LENGTH, data_size);
    let res = eval_uint256(z=z, n_coefficients=BLOB_LENGTH, coefficients=data);
    assert evals[0] = res;
    return eval_in_chunks(
        z=z,
        n_blobs=n_blobs - 1,
        data_size=data_size - BLOB_LENGTH,
        data=&data[BLOB_LENGTH],
        evals=&evals[1],
    );
}

// A Wrapper for `horner_eval` that returns Uint256.
func eval_uint256{range_check_ptr}(
    z: BigInt3, n_coefficients: felt, coefficients: felt*
) -> Uint256 {
    let unreduced_y = horner_eval(
        n_coefficients=n_coefficients, coefficients=coefficients, point=z
    );
    // Call reduced_mul with 1 to get a reduced form of the value that satisfies the
    // assumptions of `bigint3_to_uint256`.
    let y_bigint = reduced_mul(unreduced_y, BigInt3(1, 0, 0));
    return bigint3_to_uint256(y_bigint);
}
