# Copyright 2017-present Kensho Technologies, LLC.
"""Functions that ensure the IR generated by the front-end satisfies all invariants."""
from typing import Dict, List

from funcy import pairwise

from .blocks import (
    Backtrack,
    CoerceType,
    ConstructResult,
    Filter,
    Fold,
    MarkLocation,
    OutputSource,
    QueryRoot,
    Recurse,
    Traverse,
    Unfold,
)
from .compiler_entities import BasicBlock
from .helpers import FoldScopeLocation
from .ir_lowering_common.common import extract_folds_from_ir_blocks
from .metadata import QueryMetadataTable


def self_consistency_check_ir_blocks_from_frontend(
    ir_blocks: List[BasicBlock], query_metadata_table: QueryMetadataTable
) -> None:
    """Assert that IR blocks originating from the frontend do not have nonsensical structure.

    Args:
        ir_blocks: list of BasicBlocks representing the IR to self-consistency check.
        query_metadata_table: QueryMetadataTable object that captures information about the query.

    Raises:
        AssertionError, if the IR has unexpected structure. If the IR produced by the front-end
        cannot be successfully and correctly used to generate MATCH or Gremlin due to a bug,
        this is the method that should catch the problem.
    """
    if not ir_blocks:
        raise AssertionError("Received no ir_blocks: {}".format(ir_blocks))

    _assert_fold_scope_locations_are_unique(ir_blocks)
    _assert_no_nested_folds(ir_blocks)
    _assert_query_root_block(ir_blocks)
    _assert_output_source_follower_blocks(ir_blocks)
    _assert_block_pairwise_constraints(ir_blocks)
    _assert_mark_location_preceding_optional_traverse(ir_blocks)
    _assert_every_location_is_marked(ir_blocks)
    _assert_coerce_type_outside_of_fold(ir_blocks)
    _assert_all_marked_locations_are_registered(ir_blocks, query_metadata_table)
    _assert_registered_locations_parent_locations(query_metadata_table)


def _assert_registered_locations_parent_locations(
    query_metadata_table: QueryMetadataTable,
) -> None:
    """Assert that all registered locations' parent locations are also registered."""
    for location, location_info in query_metadata_table.registered_locations:
        if (
            location != query_metadata_table.root_location
            and not query_metadata_table.root_location.is_revisited_at(location)
        ):
            # If the location is not the root location and is not a revisit of the root,
            # then it must have a parent location.
            if location_info.parent_location is None:
                raise AssertionError(
                    "Found a location that is not the root location of the query "
                    "or a revisit of the root, but does not have a parent: "
                    "{} {}".format(location, location_info)
                )

        if location_info.parent_location is not None:
            # Make sure the parent_location is also registered.
            # If the location is not registered, the following line will raise an error.
            query_metadata_table.get_location_info(location_info.parent_location)


def _assert_all_marked_locations_are_registered(
    ir_blocks: List[BasicBlock], query_metadata_table: QueryMetadataTable
) -> None:
    """Assert that all locations in MarkLocation blocks have registered and valid metadata."""
    # Grab all the registered locations, then make sure that:
    # - Any location that appears in a MarkLocation block is also registered.
    # - There are no registered locations that do not appear in a MarkLocation block.
    registered_locations = {location for location, _ in query_metadata_table.registered_locations}

    ir_encountered_locations = {
        block.location for block in ir_blocks if isinstance(block, MarkLocation)
    }

    unregistered_locations = ir_encountered_locations - registered_locations
    unencountered_locations = registered_locations - ir_encountered_locations
    if unregistered_locations:
        raise AssertionError(
            "IR blocks unexpectedly contain locations not registered in the "
            "QueryMetadataTable: {}".format(unregistered_locations)
        )
    if unencountered_locations:
        raise AssertionError(
            "QueryMetadataTable unexpectedly contains registered locations that "
            "never appear in the IR blocks: {}".format(unencountered_locations)
        )


def _assert_fold_scope_locations_are_unique(ir_blocks: List[BasicBlock]) -> None:
    """Assert that every FoldScopeLocation that exists on a Fold block is unique."""
    observed_locations: Dict[FoldScopeLocation, Fold] = dict()
    for block in ir_blocks:
        if isinstance(block, Fold):
            alternate = observed_locations.get(block.fold_scope_location, None)
            if alternate is not None:
                raise AssertionError(
                    "Found two Fold blocks with identical FoldScopeLocations: "
                    "{} {} {}".format(alternate, block, ir_blocks)
                )
            observed_locations[block.fold_scope_location] = block


def _assert_no_nested_folds(ir_blocks: List[BasicBlock]) -> None:
    """Assert that there are no nested Fold contexts, and that every Fold has a matching Unfold."""
    fold_seen = False
    for block in ir_blocks:
        if isinstance(block, Fold):
            if fold_seen:
                raise AssertionError("Found a nested Fold contexts: {}".format(ir_blocks))
            else:
                fold_seen = True
        elif isinstance(block, Unfold):
            if not fold_seen:
                raise AssertionError(
                    "Found an Unfold block without a matching Fold: {}".format(ir_blocks)
                )
            else:
                fold_seen = False


def _assert_query_root_block(ir_blocks: List[BasicBlock]) -> None:
    """Assert that QueryRoot is always the first block, and only the first block."""
    if not isinstance(ir_blocks[0], QueryRoot):
        raise AssertionError("The first block was not QueryRoot: {}".format(ir_blocks))
    for block in ir_blocks[1:]:
        if isinstance(block, QueryRoot):
            raise AssertionError("Found QueryRoot after the first block: {}".format(ir_blocks))


def _self_consistency_check_construct_result_block(ir_blocks: List[BasicBlock]) -> None:
    """Assert that ConstructResult is always the last block, and only the last block."""
    if not isinstance(ir_blocks[-1], ConstructResult):
        raise AssertionError("The last block was not ConstructResult: {}".format(ir_blocks))
    for block in ir_blocks[:-1]: