import os
import pandas as pd
import xgboost as xgb
from typing import Dict, List, Any

class EndpointHandler:
    def __init__(self, path=""):
        self.model = xgb.XGBClassifier()
        self.model.load_model(os.path.join(path, "xgboost_model.json"))
        self.features = [
            'Speed (RPM)', 'Torque (Nm)', 'Flux (Wb)', 'Voltage (V)', 'Current (A)',
            'Power (W)', 'Kp', 'Ki', 'Kd', 'PID Gain Adjustment', 
            'Speed Overshoot (RPM)', 'Torque Ripple (Nm)', 'Flux Ripple (Wb)', 
            'Control Error (RPM)', 'ISE', 'Stabilization Time (s)', 
            'Load Disturbance (Nm)', 'Internal Disturbance (Friction, Temp)', 
            'Simulation Step (s)', 'Optimization Iteration'
        ]

    def __call__(self, data: Any) -> List[Dict[str, Any]]:
        inputs = data.get("inputs", data)
        if isinstance(inputs, dict): inputs = [inputs]
        df = pd.DataFrame(inputs)
        for col in self.features: 
            if col not in df.columns: df[col] = 0
        df = df[self.features]
        preds = self.model.predict(df)
        probs = self.model.predict_proba(df)[:, 1]
        return [{"label": "Success" if int(p) == 1 else "Failure", "score": float(pr)} for p, pr in zip(preds, probs)]