# services/agent_client.py

import logging
from abc import ABC
from typing import Any, AsyncGenerator, Generic, List

from agentic_chat_ui.src.constants import TCController
from agentic_chat_ui.src.mixins import FromConfigMixin
from agentic_chat_ui.src.models import AgentClient
from agentic_chat_ui.src.utils import load_class

_logger = logging.getLogger(__name__)


class Controller(FromConfigMixin[TCController], ABC, Generic[TCController]):

    def __init__(self, config: TCController) -> None:
        _logger.info("Initializing Controller")
        _logger.debug(f"Controller config received: {config}")

        self.config = config

        module_path = self.config.agent_client.module_path
        _logger.info(f"Loading AgentClient class from module_path='{module_path}'")

        try:
            client_cls = load_class(module_path)
            _logger.debug(f"Loaded client class: {client_cls}")

            self.client: AgentClient[Any] = client_cls.from_config(self.config.agent_client)
            _logger.info(f"AgentClient successfully instantiated from '{module_path}'")
        except Exception as e:
            _logger.exception(f"Failed to initialize AgentClient from '{module_path}': {e}")
            raise

    async def astream_message(
        self,
        message: str,
        history: List[str],
    ) -> AsyncGenerator[str, None]:

        _logger.info("Starting async message stream")
        _logger.debug(f"Message length: {len(message)} | History size: {len(history)}")
        try:
            async for token in self.client.astream_chat(message):  # type: ignore[attr-defined]
                yield token

            _logger.info("Completed async message stream")
        except Exception as e:
            _logger.exception(f"Streaming failed: {e}")
            raise
