from typing import Dict, Optional

from mineru_flow.internal.processor.base import BasePhaseProcessor
from mineru_flow.internal.processor.chunker.default import DefaultDocumentChunker
from mineru_flow.internal.processor.knowledge_base_import.ragflow import (
    KnowledgeBaseImportProcessor,
)
from mineru_flow.internal.processor.parser.mineru import MinerUParserProcessor
from mineru_flow.internal.schema.state import Phase


class ProcessorRegistry:
    """处理器注册表"""

    def __init__(self):
        self._processors: Dict[str, BasePhaseProcessor] = {}
        self._register_default_processors()

    def _register_default_processors(self):
        """注册默认处理器"""
        self.register(Phase.PARSE, MinerUParserProcessor())
        self.register(Phase.CHUNK, DefaultDocumentChunker())
        self.register(Phase.IMPORT, KnowledgeBaseImportProcessor())

    def register(self, phase: str, processor: BasePhaseProcessor):
        """注册处理器"""
        self._processors[phase] = processor

    def get_processor(self, phase: str) -> Optional[BasePhaseProcessor]:
        """获取处理器"""
        return self._processors.get(phase)
