from typing import Any, List, Dict, Union, Iterable, Callable, Generator, Optional
from itertools import islice, chain
import logging
import jmespath  # type: ignore


# We are utilizing here the fact that Model is a singletone and set up
# in the thebeast.conf.mapping
from followthemoney import model as ftm  # type: ignore
from followthemoney.types import registry  # type: ignore
from followthemoney.util import make_entity_id  # type: ignore
from followthemoney.schema import Schema  # type: ignore
from followthemoney.exc import InvalidData

from thebeast.contrib.ftm_ext.rigged_entity_proxy import RiggedEntityProxy
from thebeast.conf.utils import import_string
from thebeast.types import RedGreenEntity


ENTITY_TYPE = registry.get("entity")
CALLABLE_CACHE: Dict[str, Callable] = {}
logger = logging.getLogger(__name__)


def ensure_list(values: Any) -> List[Any]:
    """
    Wrap the value into the list if needed
    """
    if not isinstance(values, list):
        return [values]
    return values


def generate_pseudo_id(temporary_entity_id: str):
    """
    Generates pseudo id for the entity to fool FTM, which is going to be resolved later
    """
    return make_entity_id(
        temporary_entity_id, key_prefix="thebeast_temporary_entity_id"
    )


def jmespath_results_as_array(path: str, record: Union[List, Dict]) -> List[Any]:
    """
    Extracts the values from the record according to the path, wraps everything into array
    """
    return ensure_list(jmespath.search(path, record) or [])


def resolve_entity_refs(
    entities: Iterable[Schema], context_entities: Dict[str, str]
) -> Generator[Schema, None, None]:
    """Replace temporary pseudo-IDs in entity-type properties with real entity IDs.

    During entity construction (make_entities), cross-entity references like
    ``employee: person`` are stored as pseudo-IDs (SHA1 hashes of the entity
    name, e.g. ``generate_pseudo_id("person")``).  After all entities in a
    collection receive their real IDs via ``make_id()``, this function walks
    every entity-type property and swaps pseudo-IDs for actual IDs using the
    ``context_entities`` lookup table.

    Metadata carried by StrProxy values is preserved through the swap via
    ``inject_meta_to_str``.  If a property value is a plain ``str`` (not a
    StrProxy) — which shouldn't happen in the normal pipeline but can occur
    if values were round-tripped through raw FTM internals — a warning is
    logged and the resolved ID is used without metadata.

    Args:
        entities: Iterable of entity proxies whose entity-type properties
            still contain pseudo-IDs.
        context_entities: Mapping from pseudo-ID → real entity ID, built up
            by ``make_entities`` as each entity gets its real ID.

    Yields:
        The same entity proxies, mutated in-place with resolved IDs.
    """
    for entity in entities:
        for prop in entity.iterprops():
            if prop.type == ENTITY_TYPE:
                resolved_properties = []

                # TODO: errors (probably red/green sorting) for the properties that cannot be resolved
                for prop_val in entity.get(prop):
                    resolved = context_entities.get(prop_val)
                    if hasattr(prop_val, 'inject_meta_to_str'):
                        resolved_properties.append(
                            prop_val.inject_meta_to_str(resolved)
                        )
                    else:
                        logger.warning(
                            "Entity %s property %s has a plain str value %r "
                            "instead of StrProxy — metadata will be lost",
                            entity.id or entity.key_prefix,
                            prop.name,
                            prop_val,
                        )
                        if resolved is not None:
                            resolved_properties.append(resolved)

                entity.set(prop, resolved_properties)

        yield entity


def make_entity(
    schema: Union[str, Schema], key_prefix: Optional[str] = None
) -> RiggedEntityProxy:
    """Instantiate an empty entity proxy of the given schema type."""

    schema_obj = ftm.get(schema) if isinstance(schema, str) else schema
    return RiggedEntityProxy(schema_obj, {}, key_prefix=key_prefix)


def resolve_callable(fqfn: str) -> Callable:
    """
    Caching callables, resolved by fully qualified function name (fqfn)
    """
    if fqfn in CALLABLE_CACHE:
        return CALLABLE_CACHE[fqfn]

    func = import_string(fqfn)
    CALLABLE_CACHE[fqfn] = func

    return func


def chunks(iterable: Iterable, size: int) -> Iterable:
    """Generate adjacent chunks of data"""
    it = iter(iterable)
    return iter(lambda: tuple(islice(it, size)), ())


def flatten(list_of_lists: Iterable[Iterable]) -> Iterable:
    "Flatten one level of nesting"
    return chain.from_iterable(list_of_lists)


def inflate_entity(entity_dict: Dict) -> RiggedEntityProxy:
    return RiggedEntityProxy.from_dict(entity_dict)


def deflate_entity(entity: RiggedEntityProxy) -> RedGreenEntity:
    asdict: Dict = entity.to_dict()
    valid: bool = True

    if entity.id is None:
        valid = False
    else:
        try:
            entity.schema.validate(asdict)
        except InvalidData:
            valid = False

    return RedGreenEntity(payload=asdict, valid=valid)
