from __future__ import annotations

import uuid
from typing import Any, Optional, Callable

from flownet.api.actor_ref import ActorRef
from flownet.core.actor import Actor


class ActorRegistry:
    def __init__(self, local_address_fn: Callable[[], str]):
        self._local_address_fn = local_address_fn
        self._id_to_actor: dict[str, Actor] = {}
        self._object_to_actor: dict[Any, Actor] = {}

    async def create_actor(
        self,
        actor_class: type,
        *args,
        actor_config: Any | None = None,
        **kwargs,
    ) -> ActorRef:
        actor_id = uuid.uuid4().hex
        actor_instance = actor_class(*args, **kwargs)
        actor = Actor(actor_id, actor_instance, config=actor_config)
        self._id_to_actor[actor_id] = actor
        self._object_to_actor[actor_instance] = actor
        return ActorRef(address=self._local_address_fn(), actor_id=actor_id)

    def make_object_actor(
        self, actor_object: Any, actor_config: Any | None = None
    ) -> ActorRef:
        if actor_object in self._object_to_actor:
            actor = self._object_to_actor[actor_object]
            return ActorRef(address=self._local_address_fn(), actor_id=actor.actor_id)

        actor_id = uuid.uuid4().hex
        actor = Actor(actor_id, actor_object, config=actor_config)
        self._id_to_actor[actor_id] = actor
        self._object_to_actor[actor_object] = actor
        return ActorRef(address=self._local_address_fn(), actor_id=actor_id)

    def register_actor(
        self, actor_id: str, actor_object: Any, actor_config: Any | None = None
    ) -> ActorRef:
        if actor_id in self._id_to_actor:
            raise ValueError(f"Actor with id {actor_id} already exists.")
        actor = Actor(actor_id, actor_object, config=actor_config)
        self._id_to_actor[actor_id] = actor
        self._object_to_actor[actor_object] = actor
        return ActorRef(address=self._local_address_fn(), actor_id=actor_id)

    def resolve(self, actor_id: str) -> Actor:
        actor = self._id_to_actor.get(actor_id)
        if actor:
            return actor
        raise ValueError(f"Actor with id {actor_id} not found.")

    def lookup_actor(self, actor_id: str) -> Optional[Actor]:
        return self._id_to_actor.get(actor_id)

    def list_actors(self) -> list[dict[str, str]]:
        records: list[dict[str, str]] = []
        for actor_id, actor in list(self._id_to_actor.items()):
            actor_object = getattr(actor, "object", None)
            actor_type = type(actor_object)
            records.append(
                {
                    "actor_id": actor_id,
                    "class_name": actor_type.__name__,
                    "module": actor_type.__module__,
                }
            )
        records.sort(key=lambda item: item["actor_id"])
        return records
