from abc import ABC, abstractmethod
from plllm_mlx.helpers import *
from typing import Any, Optional, List, TYPE_CHECKING
from plllm_mlx.logging_config import get_logger

logger = get_logger(__name__)
import os

if TYPE_CHECKING:
    from .special_tokens import SpecialTokens


class PlStepProcessor(ABC):
    def __init__(self, special_tokens: Optional["SpecialTokens"] = None):
        self._special_tokens = special_tokens
        self._can_step_continue = True
        self._total_tokens = 0
        self._unprocessed_text = ""

    __STEPP_MAP__ = {}

    @staticmethod
    def registerStepProcessor(step_p_name: str, step_p_clz):
        PlStepProcessor.__STEPP_MAP__[step_p_name] = step_p_clz

    @staticmethod
    def findStepProcessor(step_p_name: str):
        return PlStepProcessor.__STEPP_MAP__.get(step_p_name, None)

    @staticmethod
    def listStepProcessors():
        return [s for s in PlStepProcessor.__STEPP_MAP__.keys()]

    @staticmethod
    @abstractmethod
    def step_clz_name() -> str:
        pass

    @property
    def special_tokens(self) -> "SpecialTokens":
        if self._special_tokens is None:
            from .special_tokens import SpecialTokens

            self._special_tokens = SpecialTokens()
        return self._special_tokens

    @property
    def is_running(self):
        return self._can_step_continue

    def stop(self):
        self._can_step_continue = False

    @property
    def total_tokens(self):
        return self._total_tokens

    @total_tokens.setter
    def total_tokens(self, tt):
        self._total_tokens = tt

    @property
    def unprocessed_text(self):
        return self._unprocessed_text

    @unprocessed_text.setter
    def unprocessed_text(self, ut):
        self._unprocessed_text = ut

    @abstractmethod
    def step(self, generate_response: Optional[Any] = None) -> Optional[PlChunk]:
        pass

    @abstractmethod
    def tool_calls(self) -> List[PlChunk]:
        pass

    @abstractmethod
    def finish(self) -> PlChunk:
        pass


# Load all step processors when current script is imported
_all_python_codes = PlUnpackPath(os.path.join(PlRootPath(), "models"), recursive=False)
for script in _all_python_codes:
    if script.endswith("_step_processor.py"):
        step_processor_clz = PlFindSpecifialSubclass(script, PlStepProcessor)
        for clz in step_processor_clz:
            logger.info(f"add step processor: {clz.step_clz_name()}")
            PlStepProcessor.registerStepProcessor(clz.step_clz_name(), clz)
