import logging
from enum import Enum
from typing import List, Optional

from fine_chunker.engine.base import BaseEngine
from fine_chunker.model.base import BaseChunkingModel
from fine_chunker.model.loader import ModelLoader
from fine_chunker.types import Chunk

logger = logging.getLogger("fine_chunker.chunker")


class ChunkerType(Enum):
    RECURSIVE = "recursive"


def _create_engine(type: ChunkerType, model: BaseChunkingModel, **kwargs) -> BaseEngine:
    if type == ChunkerType.RECURSIVE:
        from fine_chunker.engine.recursive import RecursiveEngine

        return RecursiveEngine(model=model, **kwargs)


class Chunker:
    """
    Main entry point for the library. Provides a simple interface to load models
    and chunk text.
    """

    def __init__(self, engine: BaseEngine):
        self.engine = engine

    @classmethod
    def from_pretrained(
        cls,
        type: ChunkerType = ChunkerType.RECURSIVE,
        model_path: str = "jboksa/modbert-chunker-base",
        device: Optional[str] = None,
        use_onnx: bool = True,
        **kwargs,
    ) -> "Chunker":
        """
        Factory method to create a chunker from a model checkpoint.

        Args:
            model_path: Path to the .pt model file or HF repo ID. Defaults to jboksa/modbert-chunker-base.
            device: 'cuda' or 'cpu'. Auto-detected if None.
            **kwargs: Arguments for RecursiveEngine (threshold_start, max_chunk_size, etc.)
        """

        logger.info(
            "Loading model from '%s' with device '%s' and ONNX=%s",
            model_path,
            device or "auto",
            use_onnx,
        )
        loader = ModelLoader(model_path, device=device, use_onnx=use_onnx)
        engine = _create_engine(type, loader.model, **kwargs)
        return cls(engine)

    def chunk(self, text: str) -> List[Chunk]:
        """Chunk a plain text string."""
        return self.engine.chunk(text)
