from __future__ import annotations
from pydantic import BaseModel  # type: ignore
from typing import (
    Any,
    Set
)
import importlib
import json

from . import defines
from .message import Message, Tunables, MessageMetaData
from .user import User
from .context import Context
from .conversation import Conversation
from .rag import ChromaDBFileWatcher, ChromaDBGetResponse, start_file_watcher, RagEntry
from .setup_logging import setup_logging
from .agents import class_registry, AnyAgent, Agent, __all__ as agents_all
from .metrics import Metrics
from .check_serializable import check_serializable
from .profile_image import generate_image, ImageRequest
from .redis_client import redis_manager

__all__ = [
    "Agent",
    "Message",
    "Tunables",
    "MessageMetaData",
    "Context",
    "Conversation",
    "Metrics",
    "RagEntry",
    "redis_manager",
    "ChromaDBFileWatcher",
    'ChromaDBGetResponse', 
    "start_file_watcher",
    "check_serializable",
    "logger",
    "User",
    "generate_image", "ImageRequest"
]

__all__.extend(agents_all)  # type: ignore

logger = setup_logging(level=defines.logging_level)


def rebuild_models():
    for class_name, (module_name, _) in class_registry.items():
        try:
            module = importlib.import_module(module_name)
            cls = getattr(module, class_name, None)

            logger.debug(f"Checking: {class_name} in module {module_name}")
            logger.debug(f"  cls: {True if cls else False}")
            logger.debug(f"  isinstance(cls, type): {isinstance(cls, type)}")
            logger.debug(
                f"  issubclass(cls, BaseModel): {issubclass(cls, BaseModel) if cls else False}"
            )
            logger.debug(
                f"  issubclass(cls, AnyAgent): {issubclass(cls, AnyAgent) if cls else False}"
            )
            logger.debug(
                f"  cls is not AnyAgent: {cls is not AnyAgent if cls else True}"
            )

            if (
                cls
                and isinstance(cls, type)
                and issubclass(cls, BaseModel)
                and issubclass(cls, AnyAgent)
                and cls is not AnyAgent
            ):
                logger.debug(f"Rebuilding {class_name} from {module_name}")
                from .agents import Agent
                from .context import Context

                cls.model_rebuild()
        except ImportError as e:
            logger.error(f"Failed to import module {module_name}: {e}")
        except Exception as e:
            logger.error(f"Error processing {class_name} in {module_name}: {e}")


# Call this after all modules are imported
rebuild_models()
