//
//  Copyright 2025-2026 Shuntaro Kasatani
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

use crate::{
    ast::{FunctionID, compilation_data::ProgramContext, scope_manager::IOBlueprint},
    backend::{
        Backend,
        func_translator::{FuncTranslator, TranslatorParams},
    },
    builtin::BuiltinRegistry,
};
use cranelift::prelude::{AbiParam, FunctionBuilder, InstBuilder, types};
use cranelift_codegen::{settings, verify_function};
use cranelift_module::{Linkage, Module};

impl Backend {
    /// Compiles the program that processes a buffer.
    ///
    /// Signature:
    /// `fn(input, output, state, should_init, buffer_size)`
    pub fn compile_buffer(
        &mut self,
        prog_ctx: &ProgramContext,
        builtin_registry: &BuiltinRegistry,
        blueprint: &IOBlueprint,
        entry_point: &FunctionID,
    ) -> Result<*const u8, String> {
        self.translate_buffer(prog_ctx, builtin_registry, blueprint, entry_point);

        // Verify the function
        let verifier_flags = settings::Flags::new(settings::builder());
        verify_function(&self.ctx.func, &verifier_flags).map_err(|e| e.to_string())?;

        let module = self.module.as_mut().unwrap();

        let id = module
            .declare_function("main", Linkage::Export, &self.ctx.func.signature)
            .map_err(|e| e.to_string())?;
        module
            .define_function(id, &mut self.ctx)
            .map_err(|e| e.to_string())?;

        module.clear_context(&mut self.ctx);
        module.finalize_definitions().unwrap();

        let code = module.get_finalized_function(id);
        Ok(code)
    }

    pub fn translate_buffer(
        &mut self,
        prog_ctx: &ProgramContext,
        builtin_registry: &BuiltinRegistry,
        blueprint: &IOBlueprint,
        entry_point: &FunctionID,
    ) {
        let module = self.module.as_mut().unwrap();

        // Add parameter for the input and output pointers
        let pointer_type = module.target_config().pointer_type();
        self.ctx.func.signature.params.extend(&[
            AbiParam::new(pointer_type), // Input pointers
            AbiParam::new(pointer_type), // Output pointers
            AbiParam::new(pointer_type), // State pointers
            AbiParam::new(types::I8),    // Should init
            AbiParam::new(types::I32),   // Size of the buffer
        ]);

        // Create a function builder
        let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_ctx);

        // Create needed blocks
        let entry_block = builder.create_block();
        // Append block params for function params
        builder.append_block_params_for_function_params(entry_block);

        // Switch to the entry block
        builder.switch_to_block(entry_block);
        builder.seal_block(entry_block);

        // Get the block parameters
        let block_params = builder.block_params(entry_block).to_vec();
        let buffer_size = block_params[4];
        let translator_params = TranslatorParams {
            input_ptr_ptr: block_params[0],
            output_ptr_ptr: block_params[1],
            state_ptr_ptr: block_params[2],
            should_init: block_params[3],
        };

        // Create a FuncTranslator and translate the loop
        let mut translator = FuncTranslator::new(builder, module, prog_ctx, builtin_registry);

        translator.create_loop(buffer_size, |translator, i, increment_block| {
            translator.translate(
                translator_params,
                Some(i),
                entry_point,
                blueprint,
                increment_block,
            );
        });

        // Add return statement after finishing the loop
        translator.builder.ins().return_(&[]);
        // Finalize the function
        translator.builder.finalize();
    }
}
