from typing import Generic, Type

from openkosmos_ai.common.ai_template_loader import AITemplateLoader, AITemplateConfig, PROMPT_ROLE, \
    TEMPLATE_SOURCE
from openkosmos_ai.common.model import AI_CONFIG_T
from openkosmos_core.config import Config


class BaseStateNode(Generic[AI_CONFIG_T]):

    def __init__(self, config=None):
        self._config = config
        c_instance = Config.instance()
        if c_instance is None:
            self._default_template_loader = AITemplateLoader()
        else:
            self._default_template_loader = AITemplateLoader(c_instance.get_template_config())
        self._config_template_loader = None

    def config_template_loader(self, config: TEMPLATE_SOURCE | AITemplateConfig):
        self._config_template_loader = AITemplateLoader(config)
        return self

    def config(self) -> Type[AI_CONFIG_T]:
        return self._config

    def get_prompt(self, name: str, role: PROMPT_ROLE = "system", use_default=False):
        if use_default or self._config_template_loader is None:
            return self._default_template_loader.get_prompt(name=name, role=role)
        else:
            return self._config_template_loader.get_prompt(name=name, role=role)
