import logging
import os
from typing import Optional

from huggingface_hub import hf_hub_download

from fine_chunker.model.base import BaseChunkingModel

logger = logging.getLogger("model")
DEFAULT_MODEL = "jboksa/modbert-chunker-base"


class ModelLoader:
    def __init__(
        self,
        checkpoint_path: str = DEFAULT_MODEL,
        device: Optional[str] = None,
        use_onnx: bool = False,
    ):
        self.checkpoint_path = checkpoint_path
        self.use_onnx = use_onnx or checkpoint_path.endswith(".onnx")

        if device is None:
            if self.use_onnx:
                self.device = "cpu"
            else:
                import torch

                self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        self.model = self._load()

    def _load(self) -> BaseChunkingModel:
        if self.use_onnx:
            return self._load_onnx()
        return self._load_torch()

    def _load_torch(self) -> BaseChunkingModel:
        import torch

        from fine_chunker.model.torch_model import TorchChunkingModel

        actual_path = self._ensure_local_file(
            self.checkpoint_path, filenames=["final_model.pt", "pytorch_model.bin"]
        )

        tokenizer_source = self._get_tokenizer_source()

        model = TorchChunkingModel(model_name=tokenizer_source)
        model.to(self.device)

        state_dict = torch.load(actual_path, map_location=self.device)
        model.load_state_dict(state_dict, strict=False)
        model.eval()
        return model

    def _load_onnx(self) -> BaseChunkingModel:
        from fine_chunker.model.onnx_model import ONNXChunkingModel

        onnx_filename = "onnx/model.onnx"
        tokenizer_filename = "onnx/tokenizer.json"

        model_path = self._ensure_local_file(
            self.checkpoint_path, filenames=[onnx_filename, "model.onnx"]
        )

        tokenizer_path = self._ensure_local_file(
            self.checkpoint_path, filenames=[tokenizer_filename, "tokenizer.json"]
        )

        return ONNXChunkingModel(
            model_path=model_path, tokenizer_path=tokenizer_path, device=self.device
        )

    def _ensure_local_file(self, path: str, filenames: list) -> str:
        if os.path.exists(path) and os.path.isfile(path):
            return path

        if os.path.isdir(path):
            for f in filenames:
                full_path = os.path.join(path, f)
                if os.path.exists(full_path):
                    return full_path
        else:
            logger.info(
                f"Path '{path}' not found locally. Fetching from Hugging Face Hub..."
            )

        for fname in filenames:
            try:
                return hf_hub_download(repo_id=path, filename=fname)
            except Exception:
                continue

        raise FileNotFoundError(
            f"Could not find or download any of {filenames} from {path}"
        )

    def _get_tokenizer_source(self) -> str:
        if not os.path.exists(self.checkpoint_path):
            return self.checkpoint_path

        if os.path.isfile(self.checkpoint_path):
            return os.path.dirname(self.checkpoint_path) or DEFAULT_MODEL

        return self.checkpoint_path
