// Copyright (C) 2019-2026 Provable Inc.
// This file is part of the Leo library.

// The Leo library is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// The Leo library 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 General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.

use crate::TypeTable;

use crate::errors::static_analyzer;
use leo_ast::{AstVisitor, DynamicOpKind, Expression, Function, Intrinsic, Node, Type};
use leo_errors::Handler;

/// Error if futures are used improperly.
///
/// This prevents, for instance, a bare call which creates an unused future.
pub fn future_check_function(function: &Function, type_table: &TypeTable, handler: &Handler) {
    let mut future_checker = FutureChecker { type_table, handler };
    future_checker.visit_block(&function.block);
}

#[derive(Clone, Copy, Debug, Default)]
enum Position {
    #[default]
    Misc,
    Await,
    TupleAccess,
    Return,
    FunctionArgument,
    LastTupleLiteral,
    Definition,
}

struct FutureChecker<'a> {
    type_table: &'a TypeTable,
    handler: &'a Handler,
}

impl FutureChecker<'_> {
    fn emit_err(&self, err: leo_errors::Formatted) {
        self.handler.emit_err(err);
    }
}

impl AstVisitor for FutureChecker<'_> {
    /* Expressions */
    type AdditionalInput = Position;
    type Output = ();

    fn visit_expression(&mut self, input: &Expression, additional: &Self::AdditionalInput) -> Self::Output {
        use Position::*;
        let is_call = matches!(input, Expression::Call(..))
            || matches!(input, Expression::DynamicOp(op) if matches!(op.kind, DynamicOpKind::Call { .. }))
            || matches!(input, Expression::Intrinsic(e) if e.name == leo_span::sym::_dynamic_call);
        let is_async_block = matches!(input, Expression::Async(..));
        // The nested `if`s below must not be collapsed into match guards: doing so changes
        // semantics because later arms would pick up the fall-through when the inner `if` is
        // false, causing spurious errors.
        #[allow(clippy::collapsible_match, clippy::collapsible_if)]
        match self.type_table.get(&input.id()) {
            Some(Type::Future(..)) if is_call | is_async_block => {
                // A call producing a Future may appear in any of these positions.
                if !matches!(additional, Await | Return | FunctionArgument | LastTupleLiteral | Definition) {
                    self.emit_err(static_analyzer::misplaced_final(input.span()));
                }
            }
            Some(Type::Future(..)) => {
                // A Future expression that's not a call may appear in any of these positions.
                if !matches!(additional, Await | Return | FunctionArgument | LastTupleLiteral | TupleAccess) {
                    self.emit_err(static_analyzer::misplaced_final(input.span()));
                }
            }
            Some(Type::Tuple(tuple)) if !matches!(tuple.elements().last(), Some(Type::Future(_))) => {}
            Some(Type::Tuple(..)) if is_call => {
                // A call producing a Tuple ending in a Future may appear in any of these positions.
                if !matches!(additional, Return | Definition) {
                    self.emit_err(static_analyzer::misplaced_final(input.span()));
                }
            }
            Some(Type::Tuple(..)) => {
                // A Tuple ending in a Future that's not a call may appear in any of these positions.
                if !matches!(additional, Return | TupleAccess) {
                    self.emit_err(static_analyzer::misplaced_final(input.span()));
                }
            }
            _ => {}
        }

        match input {
            Expression::Array(array) => self.visit_array(array, &Position::Misc),
            Expression::ArrayAccess(access) => self.visit_array_access(access, &Position::Misc),
            Expression::Intrinsic(intr) => self.visit_intrinsic(intr, &Position::Misc),
            Expression::Async(async_) => self.visit_async(async_, &Position::Misc),
            Expression::Binary(binary) => self.visit_binary(binary, &Position::Misc),
            Expression::Call(call) => self.visit_call(call, &Position::Misc),
            Expression::DynamicOp(op) => self.visit_dynamic_op(op, &Position::Misc),
            Expression::Cast(cast) => self.visit_cast(cast, &Position::Misc),
            Expression::Composite(composite) => self.visit_composite_init(composite, &Position::Misc),
            Expression::Err(err) => self.visit_err(err, &Position::Misc),
            Expression::Path(path) => self.visit_path(path, &Position::Misc),
            Expression::Literal(literal) => self.visit_literal(literal, &Position::Misc),
            Expression::MemberAccess(access) => self.visit_member_access(access, &Position::Misc),
            Expression::Repeat(repeat) => self.visit_repeat(repeat, &Position::Misc),
            Expression::Ternary(ternary) => self.visit_ternary(ternary, &Position::Misc),
            Expression::Tuple(tuple) => self.visit_tuple(tuple, additional),
            Expression::TupleAccess(access) => self.visit_tuple_access(access, &Position::Misc),
            Expression::Unary(unary) => self.visit_unary(unary, &Position::Misc),
            Expression::Unit(unit) => self.visit_unit(unit, &Position::Misc),
        }
    }

    fn visit_array_access(
        &mut self,
        input: &leo_ast::ArrayAccess,
        _additional: &Self::AdditionalInput,
    ) -> Self::Output {
        self.visit_expression(&input.array, &Position::Misc);
        self.visit_expression(&input.index, &Position::Misc);
    }

    fn visit_member_access(
        &mut self,
        input: &leo_ast::MemberAccess,
        _additional: &Self::AdditionalInput,
    ) -> Self::Output {
        self.visit_expression(&input.inner, &Position::Misc);
    }

    fn visit_tuple_access(
        &mut self,
        input: &leo_ast::TupleAccess,
        _additional: &Self::AdditionalInput,
    ) -> Self::Output {
        self.visit_expression(&input.tuple, &Position::TupleAccess);
    }

    fn visit_intrinsic(
        &mut self,
        input: &leo_ast::IntrinsicExpression,
        _additional: &Self::AdditionalInput,
    ) -> Self::Output {
        let position = if let Some(Intrinsic::FinalRun) = Intrinsic::from_symbol(input.name, &input.type_parameters) {
            Position::Await
        } else {
            Position::Misc
        };

        input.arguments.iter().for_each(|arg| {
            self.visit_expression(arg, &position);
        });
    }

    fn visit_call(&mut self, input: &leo_ast::CallExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        input.arguments.iter().for_each(|expr| {
            self.visit_expression(expr, &Position::FunctionArgument);
        });
        Default::default()
    }

    fn visit_tuple(&mut self, input: &leo_ast::TupleExpression, additional: &Self::AdditionalInput) -> Self::Output {
        let next_position = match additional {
            Position::Definition | Position::Return => Position::LastTupleLiteral,
            _ => Position::Misc,
        };
        let mut iter = input.elements.iter().peekable();
        while let Some(expr) = iter.next() {
            let position = if iter.peek().is_some() { &Position::Misc } else { &next_position };
            self.visit_expression(expr, position);
        }
        Default::default()
    }

    /* Statments */
    fn visit_definition(&mut self, input: &leo_ast::DefinitionStatement) {
        if let Some(ty) = input.type_.as_ref() {
            self.visit_type(ty)
        }
        self.visit_expression(&input.value, &Position::Definition);
    }

    fn visit_return(&mut self, input: &leo_ast::ReturnStatement) {
        self.visit_expression(&input.expression, &Position::Return);
    }
}
