/*
 * AquaVM Workflow Engine
 *
 * Copyright (C) 2024 Fluence DAO
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation version 3 of the
 * License.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use super::ExecutionCtx;
use super::ExecutionResult;
use super::TraceHandler;
use crate::log_instruction;

use air_parser::ast::New;
use air_parser::ast::NewArgument;

impl<'i> super::ExecutableInstruction<'i> for New<'i> {
    fn execute(&self, exec_ctx: &mut ExecutionCtx<'i>, trace_ctx: &mut TraceHandler) -> ExecutionResult<()> {
        log_instruction!(new, exec_ctx, trace_ctx);

        prolog(self, exec_ctx);
        // it should be a lazy error evaluating after execution of epilog block, since it's
        // necessary to return a restricted variable to it's previous state in case of
        // any error. It's highly important to distinguish between global and restricted streams
        // at the end of execution to make a correct data.
        let instruction_result = self.instruction.execute(exec_ctx, trace_ctx);
        let epilog_result = epilog(self, exec_ctx, trace_ctx);

        match (instruction_result, epilog_result) {
            (Ok(()), Ok(())) => Ok(()),
            // instruction error has higher priority over epilog result error,
            // because epilog returns "utility" errors that normally (meaning that AquaVM
            // scalar handling code doesn't contain errors) shouldn't happen,
            // additionally see test new_with_randomly_set_scalars_in_fold_2
            (err @ Err(_), _) => err,
            (_, err @ Err(_)) => err,
        }
    }
}

fn prolog<'i>(new: &New<'i>, exec_ctx: &mut ExecutionCtx<'i>) {
    let position = new.span.left;
    match &new.argument {
        NewArgument::Stream(stream) => exec_ctx.streams.meet_scope_start(stream.name, new.span),
        NewArgument::StreamMap(stream_map) => exec_ctx.stream_maps.meet_scope_start(stream_map.name, new.span),
        NewArgument::Scalar(scalar) => exec_ctx.scalars.meet_new_start_scalar(scalar.name.to_string()),
        NewArgument::CanonStream(canon_stream) => exec_ctx
            .scalars
            .meet_new_start_canon_stream(canon_stream.name.to_string()),
        NewArgument::CanonStreamMap(canon_stream_map) => exec_ctx
            .scalars
            .meet_new_start_canon_stream_map(canon_stream_map.name.to_string()),
    }

    exec_ctx.tracker.meet_new(position);
}

fn epilog<'i>(new: &New<'i>, exec_ctx: &mut ExecutionCtx<'i>, trace_ctx: &mut TraceHandler) -> ExecutionResult<()> {
    match &new.argument {
        NewArgument::Stream(stream) => exec_ctx.streams.meet_scope_end(stream.name.to_string(), trace_ctx),
        NewArgument::StreamMap(stream_map) => exec_ctx
            .stream_maps
            .meet_scope_end(stream_map.name.to_string(), trace_ctx),
        NewArgument::Scalar(scalar) => exec_ctx.scalars.meet_new_end_scalar(scalar.name),
        NewArgument::CanonStream(canon_stream) => exec_ctx.scalars.meet_new_end_canon_stream(canon_stream.name),
        NewArgument::CanonStreamMap(canon_stream_map) => {
            exec_ctx.scalars.meet_new_end_canon_stream_map(canon_stream_map.name)
        }
    }
}
