import os
import json
import joblib
import pandas as pd
from pathlib import Path
from flask import Flask, request, jsonify
from flask_cors import CORS

app = Flask(__name__)
CORS(app)

# Model directory in HF Space (root/models/)
MODEL_PATH = Path("models") / "best_model.joblib"

if MODEL_PATH.exists():
    PIPELINE = joblib.load(MODEL_PATH)
else:
    PIPELINE = None


@app.get("/")
def root():
    return jsonify({
        "status": "ok",
        "message": "ExtraaLearn Lead Conversion API",
        "model_loaded": PIPELINE is not None
    })


@app.post("/predict")
def predict_single():
    """Predict for a single JSON input"""
    if PIPELINE is None:
        return jsonify({"error": "Model not loaded"}), 503

    payload = request.get_json(force=True)
    X = pd.DataFrame([payload])

    proba = float(PIPELINE.predict_proba(X)[:, 1][0])
    pred = int(proba >= 0.5)

    return jsonify({
        "probability": proba,
        "prediction": pred
    })


@app.post("/predict-batch")
def predict_batch():
    """Predict for multiple JSON rows"""
    if PIPELINE is None:
        return jsonify({"error": "Model not loaded"}), 503

    payload = request.get_json(force=True)

    if isinstance(payload, dict) and "records" in payload:
        records = payload["records"]
    elif isinstance(payload, list):
        records = payload
    else:
        return jsonify({"error": "Invalid payload format"}), 400

    df = pd.DataFrame(records)
    probas = PIPELINE.predict_proba(df)[:, 1]
    preds = (probas >= 0.5).astype(int)

    df["conversion_proba"] = probas
    df["prediction"] = preds

    return df.to_json(orient="records")


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port)