// 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/>.

mod ast;

mod program;

mod scope_state;

mod visitor;
use visitor::*;

use self::scope_state::ScopeState;
use crate::{CompilerState, Pass};

use leo_ast::{Ast, CallGraph, CompositeGraph, NetworkName, UnitVisitor};
use leo_errors::Result;

use snarkvm::prelude::{CanaryV0, MainnetV0, Network, TestnetV0};

use indexmap::{IndexMap, IndexSet};

/// Specify network limits for type checking.
#[derive(Clone)]
pub struct TypeCheckingInput {
    pub max_array_elements: usize,
    pub max_mappings: usize,
    pub max_functions: usize,
    pub max_inputs: usize,
    pub max_outputs: usize,
}

impl TypeCheckingInput {
    /// Create a new `TypeCheckingInput` from the given network.
    pub fn new(network: NetworkName) -> Self {
        let (max_array_elements, max_mappings, max_functions, max_inputs, max_outputs) = match network {
            NetworkName::MainnetV0 => (
                MainnetV0::LATEST_MAX_ARRAY_ELEMENTS(),
                MainnetV0::MAX_MAPPINGS,
                MainnetV0::MAX_FUNCTIONS,
                MainnetV0::MAX_INPUTS,
                MainnetV0::MAX_OUTPUTS,
            ),
            NetworkName::TestnetV0 => (
                TestnetV0::LATEST_MAX_ARRAY_ELEMENTS(),
                TestnetV0::MAX_MAPPINGS,
                TestnetV0::MAX_FUNCTIONS,
                TestnetV0::MAX_INPUTS,
                TestnetV0::MAX_OUTPUTS,
            ),
            NetworkName::CanaryV0 => (
                CanaryV0::LATEST_MAX_ARRAY_ELEMENTS(),
                CanaryV0::MAX_MAPPINGS,
                CanaryV0::MAX_FUNCTIONS,
                CanaryV0::MAX_INPUTS,
                CanaryV0::MAX_OUTPUTS,
            ),
        };
        Self { max_array_elements, max_mappings, max_functions, max_inputs, max_outputs }
    }
}

/// A pass to check types.
///
/// Also constructs the composite graph, call graph, and local symbol table data.
pub struct TypeChecking;

impl Pass for TypeChecking {
    type Input = TypeCheckingInput;
    type Output = ();

    const NAME: &'static str = "TypeChecking";

    fn do_pass(input: Self::Input, state: &mut CompilerState) -> Result<Self::Output> {
        let composite_names = state
            .symbol_table
            .iter_records()
            .map(|(loc, _)| loc.clone())
            .chain(state.symbol_table.iter_structs().map(|(loc, _)| loc.clone()))
            .collect();
        let function_names: IndexSet<leo_ast::Location> =
            state.symbol_table.iter_functions().map(|(loc, _)| loc.clone()).collect();

        let ast = std::mem::take(&mut state.ast);

        // Initialize the composite graph with all the composites in the program.
        state.composite_graph = CompositeGraph::new(composite_names);
        // Reinitialize the call counts
        state.call_count = function_names.iter().cloned().map(|n| (n, 0)).collect();
        // Initialize the call graph with all the functions in the program.
        state.call_graph = CallGraph::new(function_names);

        let mut visitor = TypeCheckingVisitor {
            state,
            scope_state: ScopeState::new(),
            async_function_input_types: IndexMap::new(),
            async_function_callers: IndexMap::new(),
            used_composites: IndexSet::new(),
            conditional_scopes: Vec::new(),
            limits: input,
            async_block_id: None,
        };

        match &ast {
            Ast::Program(program) => visitor.visit_program(program),
            Ast::Library(library) => visitor.visit_library(library),
        }

        visitor.state.handler.last_err()?;

        // Records are always kept regardless of whether they appear in function bodies, because
        // they are part of the program's public definition and required for interface conformance.
        visitor.used_composites.extend(visitor.state.symbol_table.iter_records().map(|(loc, _)| loc.clone()));

        // Remove unused composites from the composite graph.
        // This prevents unused composite definitions from being included in the generated bytecode.
        visitor.state.composite_graph.retain_nodes(&visitor.used_composites);
        visitor.state.ast = ast;

        Ok(())
    }
}
