import pandas as pd
import numpy as np
import joblib
import os
from huggingface_hub import hf_hub_download, login

class EngineConditionModel:
    def __init__(self, model_name):
        self.model_name = model_name
        self.model = None
        self.scaler = None
        self.features = [
            'Engine rpm', 'Lub oil pressure', 'Fuel pressure',
            'Coolant pressure', 'lub oil temp', 'Coolant temp'
        ]
        self.load_model()

    def load_model(self):
        try:
            hf_token = os.getenv("HF_TOKEN")
            if hf_token:
                try:
                    login(token=hf_token)
                except Exception as e:
                    print(f"HF login warning: {e}")

            model_path = hf_hub_download(
                repo_id=self.model_name,
                filename="model/best_engine_model.joblib",  # adjust if different
                repo_type="model",
                token=hf_token
            )
            self.model = joblib.load(model_path)
            print("✅ Loaded model artifact.")

            # Optional scaler
            try:
                scaler_path = hf_hub_download(
                    repo_id=self.model_name,
                    filename="scaler.joblib",
                    repo_type="model",
                    token=hf_token
                )
                self.scaler = joblib.load(scaler_path)
                print("✅ Loaded scaler.")
            except Exception:
                self.scaler = None
        except Exception as e:
            print(f"❌ HF Hub load error: {e}
⚠️ Falling back to dummy model.")
            from sklearn.ensemble import RandomForestClassifier
            self.model = RandomForestClassifier(n_estimators=10, random_state=42)
            X_dummy = np.random.rand(10, len(self.features))
            y_dummy = np.random.randint(0, 2, 10)
            self.model.fit(X_dummy, y_dummy)
            self.scaler = None

    def preprocess(self, data):
        df = pd.DataFrame([data]) if isinstance(data, dict) else data.copy()
        X = df[self.features]
        if self.scaler is not None:
            return self.scaler.transform(X)
        return X.values

    def predict(self, data):
        try:
            X = self.preprocess(data)
            pred = int(self.model.predict(X)[0])
            if hasattr(self.model, "predict_proba"):
                proba = self.model.predict_proba(X)
                conf = float(proba[0, 1]) if proba.shape[1] > 1 else float(proba[0, 0])
            else:
                conf = 0.5
            condition = 'Maintenance Required' if pred == 1 else 'Normal'
            return {"prediction": pred, "probability": conf, "condition": condition}
        except Exception as e:
            print(f"❌ Prediction error: {e}")
            return {"prediction": -1, "probability": 0.0, "condition": "Error"}

_engine_model_instance = None

def load_engine_model(model_name="dhani10/engine-condition-model"):
    global _engine_model_instance
    if _engine_model_instance is None:
        _engine_model_instance = EngineConditionModel(model_name)
    return _engine_model_instance