import os
import pandas as pd
import joblib
from huggingface_hub import hf_hub_download
from utils.config import AppConfig
from utils.tracing import Tracer

class PredictTool:
    def __init__(self, cfg: AppConfig, tracer: Tracer):
        self.cfg = cfg
        self.tracer = tracer
        self._model = None
        self._feature_meta = None

    def _ensure_loaded(self):
        if self._model is None:
            path = hf_hub_download(repo_id=self.cfg.hf_model_repo, filename="model.pkl", token=os.getenv("HF_TOKEN"))
            self._model = joblib.load(path)
            meta = hf_hub_download(repo_id=self.cfg.hf_model_repo, filename="feature_metadata.json", token=os.getenv("HF_TOKEN"))
            import json
            with open(meta, "r") as f:
                self._feature_meta = json.load(f)

    def run(self, df: pd.DataFrame) -> pd.DataFrame:
        self._ensure_loaded()
        use_cols = self._feature_meta.get("feature_order", list(df.columns))
        X = df[use_cols].copy()
        preds = self._model.predict_proba(X)[:, 1] if hasattr(self._model, "predict_proba") else self._model.predict(X)
        out = df.copy()
        out[self._feature_meta.get("prediction_column", "prediction")] = preds
        self.tracer.trace_event("predict", {"rows": len(out)})
        return out