"""Entity registry for the Agentbyte WebUI."""

from __future__ import annotations

from typing import Any

from agentbyte.agents.base import BaseAgent
from agentbyte.orchestration.base import BaseOrchestrator
from agentbyte.workflow.core.workflow import BaseWorkflow

from .discovery import AgentbyteScanner
from .models import AgentInfo, Entity, OrchestratorInfo, WorkflowInfo


class EntityRegistry:
    """Combined registry for directory-discovered and in-memory entities."""

    def __init__(self, entities_dir: str | None = None) -> None:
        self.entities_dir = entities_dir
        self.scanner = AgentbyteScanner(entities_dir) if entities_dir else None
        self._registered: dict[str, tuple[Any, Entity]] = {}

    def register_entity(self, entity_id: str, entity_obj: Any) -> Entity:
        """Register an in-memory entity."""
        info = self._build_info(entity_id, entity_obj, source="memory")
        self._registered[entity_id] = (entity_obj, info)
        return info

    def unregister_entity(self, entity_id: str) -> bool:
        """Remove an in-memory entity."""
        return self._registered.pop(entity_id, None) is not None

    def list_entities(self) -> list[Entity]:
        """Return all visible entities."""
        entities = list(self._registered.values())
        discovered = self.scanner.discover_entities() if self.scanner else []
        infos = [info for _, info in entities] + discovered
        infos.sort(key=lambda item: (item.type, item.name.lower()))
        return infos

    def get_entity_info(self, entity_id: str) -> Entity | None:
        """Return metadata for an entity id."""
        if entity_id in self._registered:
            return self._registered[entity_id][1]
        for info in self.list_entities():
            if info.id == entity_id:
                return info
        return None

    def get_entity_object(self, entity_id: str) -> Any | None:
        """Return the underlying runtime object."""
        if entity_id in self._registered:
            return self._registered[entity_id][0]
        if self.scanner is None:
            return None
        return self.scanner.get_entity_object(entity_id)

    def clear_cache(self) -> None:
        """Clear discovery caches."""
        if self.scanner is not None:
            self.scanner.clear_cache()

    def _build_info(self, entity_id: str, entity_obj: Any, source: str) -> Entity:
        if isinstance(entity_obj, BaseAgent):
            model = getattr(entity_obj.model_client, "model", None)
            return AgentInfo(
                id=entity_id,
                name=entity_obj.name,
                description=entity_obj.description,
                source=source,
                module_path=None,
                example_tasks=list(entity_obj.example_tasks),
                has_env=False,
                model=str(model) if model else None,
                tools=[tool.name for tool in entity_obj.tools],
                memory_type=type(entity_obj.memory).__name__ if entity_obj.memory else None,
                middleware_count=len(entity_obj.middleware_chain.middlewares),
            )
        if isinstance(entity_obj, BaseOrchestrator):
            return OrchestratorInfo(
                id=entity_id,
                name=getattr(entity_obj, "name", entity_id),
                description=getattr(entity_obj, "description", None),
                source=source,
                module_path=None,
                example_tasks=list(getattr(entity_obj, "example_tasks", []) or []),
                has_env=False,
                orchestrator_type=entity_obj.__class__.__name__,
                agents=[agent.name for agent in entity_obj.agents],
                termination_conditions=self._extract_termination_conditions(
                    entity_obj.termination
                ),
            )
        if isinstance(entity_obj, BaseWorkflow):
            start_step_id = getattr(entity_obj, "start_step_id", None)
            input_schema = None
            if start_step_id and start_step_id in entity_obj.steps:
                input_type = getattr(entity_obj.steps[start_step_id], "input_type", None)
                if input_type and hasattr(input_type, "model_json_schema"):
                    input_schema = input_type.model_json_schema()
            return WorkflowInfo(
                id=entity_id,
                name=entity_obj.config.name,
                description=entity_obj.config.description,
                source=source,
                module_path=None,
                example_tasks=[],
                has_env=False,
                steps=sorted(entity_obj.steps.keys()),
                start_step=start_step_id,
                start_step_id=start_step_id,
                end_step_ids=list(entity_obj.end_step_ids),
                input_schema=input_schema,
            )
        raise TypeError(
            "Entity must be an Agentbyte Agent, Orchestrator, or Workflow instance"
        )

    def _extract_termination_conditions(self, termination: Any) -> list[str]:
        metadata = getattr(termination, "get_metadata", None)
        if callable(metadata):
            details = metadata()
            if isinstance(details, dict):
                conditions = details.get("conditions")
                if isinstance(conditions, list):
                    return [str(item) for item in conditions]
        return [termination.__class__.__name__]


__all__ = ["EntityRegistry"]
