from functools import cache
from time import perf_counter
from typing import Type, TypeVar, Iterable
from weakref import WeakValueDictionary

from saengra.adapter import Adapter
from saengra.c_extension import DirectAdapter
from saengra.socket_adapter import SocketAdapter
from saengra.api import Observation
from saengra.client import SaengraClient
from saengra.entity import Entity, EnvProtocol
from saengra.graph import Primitive, Update, AddVertex, RemoveVertex, Edge, Refs
from saengra.observer import Observer, RefsHandler
from saengra.utilities.colors import light_green
from saengra.utilities.itertools import group_by
from saengra.utilities.loggers import logger, LazyRefsFormatter

T = TypeVar("T")


class Environment(EnvProtocol):
    def __init__(
        self,
        *,
        client: SaengraClient | None = None,
        adapter: Adapter | None = None,
        entity_types: Iterable[Type[Entity]] | None = None,
    ) -> None:
        self._entity_types = [] if entity_types is None else [*entity_types]
        self._primitive_to_entity = {
            entity_type.primitive_type: entity_type
            for entity_type in self._entity_types
        }
        self._currently_committing = False

        if adapter is not None:
            self._adapter = adapter
        elif client is not None:
            self._adapter = SocketAdapter(client=client)
        else:
            self._adapter = DirectAdapter()
        # self._profiler = Profiler()
        self._observers: dict[str, Observer] = {}

        self._entities: dict[Primitive, Entity] = {}
        self._temporary_entities: WeakValueDictionary[Primitive, Entity] = (
            WeakValueDictionary()
        )
        """ 
        Temporary entities are entities that refer to primitives that are not in the graph.
        E.g. in the following code

              e = MyEntity(graph=g, ...)
              if not e.alive:
                  env.add(e)
                  ...

        a temporary entity will be created in the line 1 which will be converted to a normal 
        entity in the line 3.
        """

    def add(self, *entities: Entity) -> None:
        updates: list[Update] = []
        for entity in entities:
            if entity.alive:
                continue
            updates.append(AddVertex(entity.primitive))
            entity._after_add()
            self._entities[entity.primitive] = entity
            del self._temporary_entities[entity.primitive]
        self._adapter.update(updates)

    def remove(self, *entities: Entity) -> None:
        updates: list[Update] = []
        for entity in entities:
            if not entity.alive:
                continue
            updates.append(RemoveVertex(entity.primitive))
            entity._after_remove()
            del self._entities[entity.primitive]
            self._temporary_entities[entity.primitive] = entity
        self._adapter.update(updates)

    def update(self, *updates: Update) -> None:
        """
        Apply low-level updates to the graph manually.

        Not compatible with Entity manipulation, so use at own risk.
        """
        self._adapter.update(updates)

    def commit(self) -> None:
        if self._currently_committing:
            raise RuntimeError("should not invoke .commit() in pre-commit hooks")
        try:
            self._currently_committing = True
            self._do_commit()
        except Exception:
            self.rollback()
            raise
        finally:
            self._currently_committing = False

    def rollback(self) -> None:
        self._adapter.rollback()
        self.sync()

    def sync(self) -> None:
        """Synchronize entities with current graph state, probably after rollback."""
        self._entities.clear()
        self._temporary_entities.clear()

        vertices, edges = self._adapter.find_all()

        # Create entity objects
        for vertex in vertices:
            entity = self._wrap_into_new_entity(vertex)
            if isinstance(entity, Entity):
                entity._after_add()
                self._entities[entity.primitive] = entity
                del self._temporary_entities[entity.primitive]

        # Fill entity relations
        vertex_to_edges = group_by(edges, key=lambda e: e.from_)
        for entity in self._entities.values():
            edges = vertex_to_edges.get(entity.primitive, [])
            labels_and_tos = [
                (edge.label, self._wrap_into_existing_entity(edge.to)) for edge in edges
            ]
            entity._sync(edges=labels_and_tos)

    def find_all(self) -> tuple[list[Primitive], list[Edge]]:
        return self._adapter.find_all()

    def find_edges(
        self,
        *,
        with_label: str | None = None,
        to: Primitive | Entity | list[Primitive | Entity] | None = None,
    ) -> list[Edge]:
        container = [] if to is None else to if isinstance(to, (list, tuple)) else [to]
        edges = self._adapter.find_edges(
            with_label=with_label, to_vertices=[_unwrap_primitive(v) for v in container]
        )
        return edges

    def find_all_of_type(self, type_: Type[T]) -> set[T]:
        primitive_type = _get_primitive_type(type_)
        vertices = self._adapter.find_vertices(with_type_name=primitive_type.__name__)
        return {
            self._wrap_into_existing_entity(v)
            for v in vertices
            if isinstance(v, primitive_type)
        }

    def match(
        self, expression: str, *placeholder_values: Primitive
    ) -> list[dict[str, Primitive | Entity]]:
        subgraphs = self._adapter.match(expression, *placeholder_values)
        return [
            {k: self._wrap_into_existing_entity(v) for k, v in sg.refs.items()}
            for sg in subgraphs
        ]

    def register_observers(self, observers: list[Observer]) -> None:
        new_observers: dict[str, Observer] = {}
        for w in observers:
            if w.id in self._observers:
                raise RuntimeError(
                    f"duplicate observer ID: {w.id}\n\n"
                    "To avoid this error, register observers all at once, or make an extra ID deduplication step"
                )
            new_observers[w.id] = w
        self._observers.update(new_observers)
        self._adapter.observe(w.to_api() for w in observers)

    def _get_living_entity(self, primitive: Primitive) -> Entity | None:
        return self._entities.get(primitive, None)

    def _get_temporary_entity(self, primitive: Primitive) -> Entity | None:
        return self._temporary_entities.get(primitive, None)

    def _register_temporary_entity(
        self, primitive: Primitive, entity: "Entity"
    ) -> None:
        self._temporary_entities[primitive] = entity

    def _do_commit(self) -> None:
        before = perf_counter()
        should_commit_again = True
        num_iterations = 0
        while should_commit_again:
            should_commit_again, observations = self._adapter.commit()
            for o in observations:
                self._invoke_handlers(o)

            num_iterations += 1
            if should_commit_again and num_iterations > 100:
                raise RuntimeError("stuck in pre-commit loop")
        after = perf_counter()
        logger.debug(
            f"Commit took {num_iterations} iterations, {after - before:.3f} seconds"
        )

    def _invoke_handlers(self, o: Observation) -> None:
        observer = self._observers[o.observer_id]
        handlers = observer.pick_handlers(o.type)
        for handler in handlers:
            self._invoke_handler(handler, o.variables)

    def _invoke_handler(self, handler: RefsHandler, refs: dict[str, Primitive]) -> None:
        logger.debug(
            f"Invoking {light_green(handler.__name__)}(%s)", LazyRefsFormatter(refs)
        )
        enriched_refs: dict[str, Primitive | Entity | Environment] = {}
        if env_arg := _get_env_arg(handler):
            enriched_refs[env_arg] = self
        for ref, primitive in refs.items():
            if entity_type := self._primitive_to_entity.get(type(primitive)):
                enriched_refs[ref] = entity_type(self, primitive)
            else:
                enriched_refs[ref] = primitive
        handler(**enriched_refs)

    def _wrap_into_new_entity(self, v: Primitive) -> Primitive | Entity:
        if type_ := self._primitive_to_entity.get(type(v)):
            return type_(env=self, primitive=v)
        return v

    def _wrap_into_existing_entity(self, v: Primitive) -> Primitive | Entity:
        if type_ := self._primitive_to_entity.get(type(v)):
            # Corresponding entity is supposed to exist.
            return self._entities[v]
        return v


def _unwrap_primitive(v: Primitive | Entity) -> Primitive:
    if isinstance(v, Entity):
        return v.primitive
    return v


def _get_primitive_type(t: Type) -> Type:
    if issubclass(t, Entity):
        return t.primitive_type
    return t


@cache
def _get_env_arg(handler: RefsHandler) -> str | None:
    for key, value in handler.__annotations__.items():
        if value is Environment:
            return key
    return None
