// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

mod arrays;
mod assigns;
mod bindings;
mod branching;
mod calls;
mod classical_args;
mod debug_metadata;
mod dynamic_vars;
mod intrinsics;
mod loops;
mod misc;
mod operators;
mod output_recording;
mod qubits;
mod results;
mod returns;

use crate::{Error, PartialEvalConfig, ProgramEntry, partially_evaluate};
use expect_test::Expect;
use qsc::{PackageType, incremental::Compiler};
use qsc_data_structures::{
    language_features::LanguageFeatures, source::SourceMap, target::TargetCapabilityFlags,
};
use qsc_fir::fir::PackageStore;
use qsc_frontend::compile::PackageStore as HirPackageStore;
use qsc_lowerer::{Lowerer, map_hir_package_to_fir};
use qsc_rca::{Analyzer, PackageStoreComputeProperties};
use qsc_rir::{
    passes::check_and_transform,
    rir::{BlockId, CallableId, Program},
};

pub fn assert_block_instructions(program: &Program, block_id: BlockId, expected_insts: &Expect) {
    let block = program.get_block(block_id);
    expected_insts.assert_eq(&block.to_string());
}

pub fn assert_blocks(program: &Program, expected_blocks: &Expect) {
    let mut str = program
        .blocks
        .iter()
        .fold("Blocks:".to_string(), |acc, (id, block)| {
            acc + &format!("\nBlock {}:", id.0) + &block.to_string()
        });

    let dbg_info = program.dbg_info.to_string();
    if !dbg_info.is_empty() {
        str += "\n";
        str += &dbg_info;
    }
    expected_blocks.assert_eq(&str);
}

pub fn assert_callable(program: &Program, callable_id: CallableId, expected_callable: &Expect) {
    let actual_callable = program.get_callable(callable_id);
    expected_callable.assert_eq(&actual_callable.to_string());
}

pub fn assert_error(error: &Error, expected_error: &Expect) {
    expected_error.assert_eq(format!("{error:?}").as_str());
}

#[must_use]
pub fn get_partial_evaluation_error(source: &str) -> Error {
    let maybe_program = compile_and_partially_evaluate(
        source,
        TargetCapabilityFlags::all(),
        PartialEvalConfig {
            generate_debug_metadata: false,
        },
    );
    match maybe_program {
        Ok(_) => panic!("partial evaluation succeeded"),
        Err(error) => error,
    }
}

#[must_use]
pub fn get_partial_evaluation_error_with_capabilities(
    source: &str,
    capabilities: TargetCapabilityFlags,
) -> Error {
    let maybe_program = compile_and_partially_evaluate(
        source,
        capabilities,
        PartialEvalConfig {
            generate_debug_metadata: false,
        },
    );
    match maybe_program {
        Ok(_) => panic!("partial evaluation succeeded"),
        Err(error) => error,
    }
}

#[must_use]
pub fn get_rir_program(source: &str) -> Program {
    let maybe_program = compile_and_partially_evaluate(
        source,
        TargetCapabilityFlags::all(),
        PartialEvalConfig {
            generate_debug_metadata: false,
        },
    );
    match maybe_program {
        Ok(program) => {
            // Verify the program can go through transformations.
            check_and_transform(&mut program.clone());
            program
        }
        Err(error) => panic!("partial evaluation failed: {error:?}"),
    }
}

#[must_use]
pub fn get_rir_program_with_dbg_metadata(source: &str) -> Program {
    let maybe_program = compile_and_partially_evaluate(
        source,
        TargetCapabilityFlags::all(),
        PartialEvalConfig {
            generate_debug_metadata: true,
        },
    );
    match maybe_program {
        Ok(program) => {
            // Verify the program can go through transformations.
            check_and_transform(&mut program.clone());
            validate(&program);
            program
        }
        Err(error) => panic!("partial evaluation failed: {error:?}"),
    }
}

fn validate(program: &Program) {
    let mut dbg_scopes = program.dbg_info.dbg_scopes.clone();
    let mut dbg_locations = program.dbg_info.dbg_locations.clone();

    // All scope and inlined_at references should be to existing dbg scopes and locations.
    for (dbg_location, _) in dbg_locations.values() {
        assert!(dbg_scopes.contains_key(dbg_location.scope));
        if let Some(inlined_at) = dbg_location.inlined_at {
            assert!(dbg_locations.contains_key(inlined_at));
        }
    }

    // All dbg location references in instructions should be to existing dbg locations.
    for instruction in program.blocks.iter().flat_map(|(_, block)| &block.0) {
        if let Some(dbg_location) = instruction.metadata().map(|metadata| metadata.dbg_location) {
            assert!(dbg_locations.contains_key(dbg_location));
        }
    }

    // Ensure all entries are referenced by removing referenced scopes/locations from the lists and then checking if any remain at the end.
    for instruction in program.blocks.iter().flat_map(|(_, block)| &block.0) {
        if let Some(dbg_location) = instruction.metadata().map(|metadata| metadata.dbg_location) {
            let mut to_remove = vec![dbg_location];
            let mut next = dbg_locations.get(dbg_location);

            while let Some(entry) = next {
                // remove referenced scope
                dbg_scopes.remove(entry.0.scope);

                if let Some(inlined_at) = entry.0.inlined_at {
                    // collect referenced dbg locations
                    next = dbg_locations.get(inlined_at);
                    to_remove.push(inlined_at);
                } else {
                    break;
                }
            }
            for id in to_remove {
                dbg_locations.remove(id);
            }
        }
    }

    assert!(
        dbg_locations.is_empty(),
        "unreferenced entry in dbg locations"
    );

    assert!(dbg_scopes.is_empty(), "unreferenced entry in dbg scopes");
}

#[must_use]
pub fn get_rir_program_with_capabilities(
    source: &str,
    capabilities: TargetCapabilityFlags,
) -> Program {
    let maybe_program = compile_and_partially_evaluate(
        source,
        capabilities,
        PartialEvalConfig {
            generate_debug_metadata: false,
        },
    );
    match maybe_program {
        Ok(program) => program,
        Err(error) => panic!("partial evaluation failed: {error:?}"),
    }
}

fn compile_and_partially_evaluate(
    source: &str,
    capabilities: TargetCapabilityFlags,
    config: PartialEvalConfig,
) -> Result<Program, Error> {
    let compilation_context = CompilationContext::new(source, capabilities);
    partially_evaluate(
        &compilation_context.fir_store,
        &compilation_context.compute_properties,
        &compilation_context.entry,
        capabilities,
        config,
    )
}

struct CompilationContext {
    fir_store: PackageStore,
    compute_properties: PackageStoreComputeProperties,
    entry: ProgramEntry,
}

impl CompilationContext {
    fn new(source: &str, capabilities: TargetCapabilityFlags) -> Self {
        let source_map = SourceMap::new([("test".into(), source.into())], Some("".into()));
        let (std_id, store) = qsc::compile::package_store_with_stdlib(capabilities);
        let compiler = Compiler::new(
            source_map,
            PackageType::Exe,
            capabilities,
            LanguageFeatures::default(),
            store,
            &[(std_id, None)],
        )
        .expect("should be able to create a new compiler");
        let package_id = map_hir_package_to_fir(compiler.source_package_id());
        let fir_store = lower_hir_package_store(compiler.package_store());
        let analyzer = Analyzer::init(&fir_store);
        let compute_properties = analyzer.analyze_all();
        let package = fir_store.get(package_id);
        let entry = ProgramEntry {
            exec_graph: package.entry_exec_graph.clone(),
            expr: (
                package_id,
                package
                    .entry
                    .expect("package must have an entry expression"),
            )
                .into(),
        };

        Self {
            fir_store,
            compute_properties,
            entry,
        }
    }
}

fn lower_hir_package_store(hir_package_store: &HirPackageStore) -> PackageStore {
    let mut fir_store = PackageStore::new();
    for (id, unit) in hir_package_store {
        let mut lowerer = Lowerer::new();
        let lowered_package = lowerer.lower_package(&unit.package, &fir_store);
        fir_store.insert(map_hir_package_to_fir(id), lowered_package);
    }
    fir_store
}
