
# Unified model loader for SG

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Optional

import yaml

logger = logging.getLogger(__name__)




class ModelSource(str, Enum):
    HF_PUBLIC     = "hf_public"      # e.g. "Qwen/Qwen3-0.6B"
    HF_PRIVATE    = "hf_private"     # private org repo, needs HF token
    LOCAL         = "local"          # /path/to/finetuned-checkpoint
    OPENAI_COMPAT = "openai_compat"  # ollama, vLLM, lm-studio, etc.
    MLX           = "mlx"            # Apple Silicon — mlx-lm 4-bit inference


@dataclass
class ModelConfig:
    
    source: ModelSource
    model_id_or_path: str
    hf_token: Optional[str] = None
    base_url: Optional[str] = None
    api_key: Optional[str] = None
    device: Optional[str] = None
    max_new_tokens: int = 256
    trust_remote_code: bool = False
    prompt_template: Optional[str] = None
    label: str = ""
    param_count_millions: Optional[int] = None

    def __post_init__(self):
        if not self.label:
            self.label = Path(self.model_id_or_path).name or self.model_id_or_path
        # Pull token from env if not explicitly set
        if self.hf_token is None:
            self.hf_token = os.environ.get("HF_TOKEN")

    @classmethod
    def from_dict(cls, d: dict) -> "ModelConfig":
        d = dict(d)
        d["source"] = ModelSource(d["source"])
        return cls(**d)

    @classmethod
    def from_yaml(cls, path: str | Path) -> "ModelConfig":
        with open(path) as f:
            data = yaml.safe_load(f)
        return cls.from_dict(data)

    def to_dict(self) -> dict:
        return {
            "source": self.source.value,
            "model_id_or_path": self.model_id_or_path,
            "hf_token": self.hf_token,
            "base_url": self.base_url,
            "api_key": self.api_key,
            "device": self.device,
            "max_new_tokens": self.max_new_tokens,
            "trust_remote_code": self.trust_remote_code,
            "prompt_template": self.prompt_template,
            "label": self.label,
            "param_count_millions": self.param_count_millions,
        }


def _detect_device() -> str:
    try:
        import torch
        if torch.cuda.is_available():
            return "cuda"
        if torch.backends.mps.is_available():
            return "mps"
    except ImportError:
        pass
    return "cpu"


class BaseModelAdapter:

    def __init__(self, config: ModelConfig):
        self.config = config

    def query(self, prompt: str, **kwargs) -> str:
        raise NotImplementedError

    def is_reachable(self) -> bool:
        try:
            self.query("ping", max_new_tokens=5)
            return True
        except Exception:
            return False

    def _apply_template(self, prompt: str) -> str:
        if self.config.prompt_template:
            return self.config.prompt_template.format(prompt=prompt)
        return prompt


class HFModelAdapter(BaseModelAdapter):
    """
    Loads HuggingFace models locally via transformers pipeline.
    Works for both public models and local fine-tuned checkpoints.
    """

    def __init__(self, config: ModelConfig):
        super().__init__(config)
        self._pipeline = None

    def _load(self):
        if self._pipeline is not None:
            return
        try:
            from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
            import torch
        except ImportError:
            raise RuntimeError("transformers and torch required: pip install transformers torch")

        device = self.config.device or _detect_device()
        model_id = self.config.model_id_or_path
        token = self.config.hf_token

        logger.info(f"Loading model '{model_id}' on {device} ...")

        # Resolve device_map for MPS/CPU (pipeline doesn't accept device_map on MPS)
        if device in ("mps", "cpu"):
            device_kwargs = {"device": device}
        else:
            device_kwargs = {"device_map": "auto"}

        self._pipeline = pipeline(
            "text-generation",
            model=model_id,
            tokenizer=model_id,
            token=token,
            trust_remote_code=self.config.trust_remote_code,
            torch_dtype="auto",
            **device_kwargs,
        )
        logger.info(f"Model loaded: {self.config.label}")

    def query(self, prompt: str, **kwargs) -> str:
        self._load()
        formatted = self._apply_template(prompt)
        max_new = kwargs.get("max_new_tokens", self.config.max_new_tokens)

        outputs = self._pipeline(
            formatted,
            max_new_tokens=max_new,
            do_sample=False,
            return_full_text=False,
        )
        return outputs[0]["generated_text"].strip()


class OpenAICompatAdapter(BaseModelAdapter):
  

    def __init__(self, config: ModelConfig):
        super().__init__(config)
        self._client = None

    def _load(self):
        if self._client is not None:
            return
        try:
            from openai import OpenAI
        except ImportError:
            raise RuntimeError("openai required: pip install openai")

        self._client = OpenAI(
            base_url=self.config.base_url,
            api_key=self.config.api_key or "none",
        )

    def query(self, prompt: str, **kwargs) -> str:
        self._load()
        formatted = self._apply_template(prompt)
        max_new = kwargs.get("max_new_tokens", self.config.max_new_tokens)

        response = self._client.chat.completions.create(
            model=self.config.model_id_or_path,
            messages=[{"role": "user", "content": formatted}],
            max_tokens=max_new,
            temperature=0,
        )
        return response.choices[0].message.content.strip()

    def is_reachable(self) -> bool:
        try:
            self._load()
            self.query("ping", max_new_tokens=5)
            return True
        except Exception as e:
            logger.warning(f"Endpoint not reachable: {e}")
            return False



class MLXModelAdapterWrapper(BaseModelAdapter):
    """
    Wraps sichgate.models.mlx_adapter.MLXModelAdapter behind the
    BaseModelAdapter interface used by model_loader.load_model().
    """

    def __init__(self, config: ModelConfig):
        super().__init__(config)
        from sichgate.models.mlx_adapter import MLXModelAdapter
        self._mlx = MLXModelAdapter(
            model_id=config.model_id_or_path,
            max_tokens=config.max_new_tokens,
        )

    def query(self, prompt: str, **kwargs) -> str:
        return self._mlx.query(prompt)


def load_model(config: ModelConfig) -> BaseModelAdapter:
    if config.source == ModelSource.MLX:
        return MLXModelAdapterWrapper(config)
    elif config.source in (ModelSource.HF_PUBLIC, ModelSource.HF_PRIVATE, ModelSource.LOCAL):
        return HFModelAdapter(config)
    elif config.source == ModelSource.OPENAI_COMPAT:
        return OpenAICompatAdapter(config)
    else:
        raise ValueError(f"Unknown model source: {config.source}")