from __future__ import annotations
from typing import TypeAlias, Dict, Tuple, Optional
import importlib
import pathlib
import inspect

from . types import registry
from .. setup_logging import setup_logging
from .. import defines
from . base import Agent

logger = setup_logging(defines.logging_level)

__all__ = [ "AnyAgent", "Agent", "registry", "class_registry" ]

# Type alias for Agent or any subclass
AnyAgent: TypeAlias = Agent  # BaseModel covers Agent and subclasses

class_registry: Dict[str, Tuple[str, str]] = {}  # Maps class_name to (module_name, class_name)

package_dir = pathlib.Path(__file__).parent
package_name = __name__

for path in package_dir.glob("*.py"):
    if path.name in ("__init__.py", "base.py") or path.name.startswith("_"):
        continue

    module_name = path.stem
    full_module_name = f"{package_name}.{module_name}"

    try:
        module = importlib.import_module(full_module_name)

        # Find all Agent subclasses in the module
        for name, obj in inspect.getmembers(module, inspect.isclass):
            if (
                issubclass(obj, AnyAgent)
                and obj is not AnyAgent
                and obj is not Agent
                and name not in class_registry
            ):
                class_registry[name] = (full_module_name, name)
                globals()[name] = obj
                logger.info(f"Adding agent: {name}")
                __all__.append(name) # type: ignore
    except ImportError as e:
        logger.error(f"Error importing {full_module_name}: {e}")
        raise e
    except Exception as e:
        logger.error(f"Error processing {full_module_name}: {e}")
        raise e