"""
===========================================================
 Machine Learning Model Loader
 - Supports Auto '.pkl' Extension
 - Auto loads vectorizer (if given)
===========================================================
"""

# ---------------- IMPORTS ----------------
import os
import joblib
import numpy as np
from .base_model import BaseModel

# ---------------- CLASS ----------------
class DarkNeuronML(BaseModel):
    """Handles loading and inference for ML models."""

    BASE_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "models")

    def __init__(self, model_path: str, vectorizer_path: str = None):
        # Auto add ".pkl" if missing
        if not model_path.endswith(".pkl"):
            model_path += ".pkl"
        if vectorizer_path and not vectorizer_path.endswith(".pkl"):
            vectorizer_path += ".pkl"

        # Build full paths inside models/
        self.model_path = os.path.join(self.BASE_MODEL_DIR, model_path)
        self.vectorizer_path = (
            os.path.join(self.BASE_MODEL_DIR, vectorizer_path)
            if vectorizer_path
            else None
        )

        self.model = None
        self.vectorizer = None

    # ---------------- LOAD MODEL ----------------
    def load_model(self):
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(f"Model file not found: {self.model_path}")

        self.model = joblib.load(self.model_path)

        # Load vectorizer if available
        if self.vectorizer_path and os.path.exists(self.vectorizer_path):
            self.vectorizer = joblib.load(self.vectorizer_path)
        return self

    # ---------------- PREDICT ----------------
    def predict(self, input_data):
        if self.model is None:
            raise ValueError("Model not loaded. Call load_model() first.")

        # Handle text models
        if self.vectorizer:
            if isinstance(input_data, str):
                X = self.vectorizer.transform([input_data])
            else:
                X = self.vectorizer.transform(input_data)
        else:
            X = np.array(input_data).reshape(1, -1)

        return self.model.predict(X)[0]