use crate::gen::conditions::{
    parse_spends, process_single_spend, validate_conditions, ParseState, SpendBundleConditions,
};
use crate::gen::flags::ALLOW_BACKREFS;
use crate::gen::spend_visitor::SpendVisitor;
use crate::gen::validation_error::{first, ErrorCode, ValidationErr};
use crate::generator_rom::{COST_PER_BYTE, GENERATOR_ROM, KLVM_DESERIALIZER};
use klvm_utils::tree_hash;
use klvmr::allocator::{Allocator, NodePtr};
use klvmr::chik_dialect::ChikDialect;
use klvmr::cost::Cost;
use klvmr::reduction::Reduction;
use klvmr::run_program::run_program;
use klvmr::serde::{node_from_bytes, node_from_bytes_backrefs};

fn subtract_cost(a: &Allocator, cost_left: &mut Cost, subtract: Cost) -> Result<(), ValidationErr> {
    if subtract > *cost_left {
        Err(ValidationErr(a.nil(), ErrorCode::CostExceeded))
    } else {
        *cost_left -= subtract;
        Ok(())
    }
}

// Runs the generator ROM and passes in the program (transactions generator).
// The program is expected to return a list of spends. Each item being:

// (parent-coin-id puzzle-reveal amount solution)

// The puzzle-reveals are then executed with the corresponding solution being
// passed as the argument. The output from those puzzles are lists of
// conditions. The conditions are parsed and returned in the
// SpendBundleConditions. Some conditions are validated, and if invalid may
// cause the function to return an error.

// the only reason we need to pass in the allocator is because the returned
// SpendBundleConditions contains NodePtr fields. If that's changed, we could
// create the allocator inside this functions as well.
pub fn run_block_generator<GenBuf: AsRef<[u8]>, V: SpendVisitor>(
    a: &mut Allocator,
    program: &[u8],
    block_refs: &[GenBuf],
    max_cost: u64,
    flags: u32,
) -> Result<SpendBundleConditions, ValidationErr> {
    let mut cost_left = max_cost;
    let byte_cost = program.len() as u64 * COST_PER_BYTE;

    subtract_cost(a, &mut cost_left, byte_cost)?;

    let generator_rom = node_from_bytes(a, &GENERATOR_ROM)?;
    let program = if (flags & ALLOW_BACKREFS) != 0 {
        node_from_bytes_backrefs(a, program)?
    } else {
        node_from_bytes(a, program)?
    };

    // iterate in reverse order since we're building a linked list from
    // the tail
    let mut args = a.nil();
    for g in block_refs.iter().rev() {
        let ref_gen = a.new_atom(g.as_ref())?;
        args = a.new_pair(ref_gen, args)?;
    }

    args = a.new_pair(args, a.nil())?;
    let args = a.new_pair(args, a.nil())?;
    let args = a.new_pair(program, args)?;

    let dialect = ChikDialect::new(flags);
    let Reduction(klvm_cost, generator_output) =
        run_program(a, &dialect, generator_rom, args, cost_left)?;

    subtract_cost(a, &mut cost_left, klvm_cost)?;

    // we pass in what's left of max_cost here, to fail early in case the
    // cost of a condition brings us over the cost limit
    let mut result = parse_spends::<V>(a, generator_output, cost_left, flags)?;
    result.cost += max_cost - cost_left;
    Ok(result)
}

pub fn extract_n<const N: usize>(
    a: &Allocator,
    mut n: NodePtr,
    e: ErrorCode,
) -> Result<[NodePtr; N], ValidationErr> {
    let mut ret: [NodePtr; N] = [NodePtr::NIL; N];
    let mut counter = 0;
    assert!(N > 0);
    while let Some((item, rest)) = a.next(n) {
        if counter == N - 1 {
            break;
        }
        n = rest;
        ret[counter] = item;
        counter += 1;
    }
    if counter != N - 1 {
        return Err(ValidationErr(n, e));
    }
    ret[counter] = n;
    Ok(ret)
}

// This has the same behavior as run_block_generator() but implements the
// generator ROM in rust instead of using the KLVM implementation.
// it is not backwards compatible in the KLVM cost computation (in this version
// you only pay cost for the generator, the puzzles and the conditions).
// it also does not apply the stack depth or object allocation limits the same,
// as each puzzle run in its own environment.
pub fn run_block_generator2<GenBuf: AsRef<[u8]>, V: SpendVisitor>(
    a: &mut Allocator,
    program: &[u8],
    block_refs: &[GenBuf],
    max_cost: u64,
    flags: u32,
) -> Result<SpendBundleConditions, ValidationErr> {
    let byte_cost = program.len() as u64 * COST_PER_BYTE;

    let mut cost_left = max_cost;
    subtract_cost(a, &mut cost_left, byte_cost)?;

    let klvm_deserializer = node_from_bytes(a, &KLVM_DESERIALIZER)?;
    let program = if (flags & ALLOW_BACKREFS) != 0 {
        node_from_bytes_backrefs(a, program)?
    } else {
        node_from_bytes(a, program)?
    };

    // iterate in reverse order since we're building a linked list from
    // the tail
    let mut blocks = a.nil();
    for g in block_refs.iter().rev() {
        let ref_gen = a.new_atom(g.as_ref())?;
        blocks = a.new_pair(ref_gen, blocks)?;
    }

    // the first argument to the generator is the serializer, followed by a list
    // of the blocks it requested.
    let mut args = a.new_pair(blocks, a.nil())?;
    args = a.new_pair(klvm_deserializer, args)?;

    let dialect = ChikDialect::new(flags);

    let Reduction(klvm_cost, mut all_spends) = run_program(a, &dialect, program, args, cost_left)?;

    subtract_cost(a, &mut cost_left, klvm_cost)?;
    all_spends = first(a, all_spends)?;

    // at this point all_spends is a list of:
    // (parent-coin-id puzzle-reveal amount solution . extra)
    // where extra may be nil, or additional extension data

    let mut ret = SpendBundleConditions::default();
    let mut state = ParseState::default();

    while let Some((spend, rest)) = a.next(all_spends) {
        all_spends = rest;
        // process the spend
        let [parent_id, puzzle, amount, solution, _spend_level_extra] =
            extract_n::<5>(a, spend, ErrorCode::InvalidCondition)?;

        let Reduction(klvm_cost, conditions) =
            run_program(a, &dialect, puzzle, solution, cost_left)?;

        subtract_cost(a, &mut cost_left, klvm_cost)?;

        let buf = tree_hash(a, puzzle);
        let puzzle_hash = a.new_atom(&buf)?;

        process_single_spend::<V>(
            a,
            &mut ret,
            &mut state,
            parent_id,
            puzzle_hash,
            amount,
            conditions,
            flags,
            &mut cost_left,
        )?;
    }
    if a.atom_len(all_spends) != 0 {
        return Err(ValidationErr(all_spends, ErrorCode::GeneratorRuntimeError));
    }

    validate_conditions(a, &ret, state, a.nil(), flags)?;

    ret.cost = max_cost - cost_left;
    Ok(ret)
}
