import os
import uvicorn
import pandas as pd
from pydantic import BaseModel
from fastapi import FastAPI
import mlflow
import os
from dotenv import load_dotenv

description = """

# Fraud Detection API

Fast API backend to predict if a payment is fraud or not from : https://huggingface.co/spaces/sdacelo/real-time-fraud-detection

## Machine-Learning 

Where you can:
* `/predict` : prediction for a single value

Check out documentation for more information on each endpoint. 
"""

tags_metadata = [
    {
        "name": "Predictions",
        "description": "Endpoints that uses our Machine Learning model",
    },
]

load_dotenv()

# Mlflow variables
MLFLOW_TRACKING_APP_URI = os.getenv("MLFLOW_TRACKING_APP_URI")
MODEL_NAME = os.getenv("MODEL_NAME", "fraud_detection_dtc")
STAGE = os.getenv("STAGE", "production")

# AWS variables
os.environ["AWS_ACCESS_KEY_ID"] = os.getenv("AWS_ACCESS_KEY_ID")
os.environ["AWS_SECRET_ACCESS_KEY"] = os.getenv("AWS_SECRET_ACCESS_KEY")

# Load model
mlflow.set_tracking_uri(MLFLOW_TRACKING_APP_URI)
model_uri = f"models:/{MODEL_NAME}@{STAGE}"
model = mlflow.sklearn.load_model(model_uri)
if model is None:
    raise ValueError("Model not found.")
else:
    print("Model loaded.") 

app = FastAPI(
    title="API for Fraud Detection Project",
    description=description,
    version="1.0",
    contact={
        "name": "Olivier-52",
        "url": "https://huggingface.co/Olivier-52",
    },
    openapi_tags=tags_metadata,)

@app.get("/")
def index():
    """Return a message to the user.

    This endpoint does not take any parameters and returns a message
    to the user. It is used to test the API.

    Returns:
        str: A message to the user.
    """
    return "Hello world! Go to /docs to try the API."

class PredictionFeatures(BaseModel):
    category : str
    amt : float
    merch_fraud_level :str
    city_fraud_level : str
    state_fraud_level : str
    zip_fraud_level : str
    job_fraud_level : str
    age_fraud_level : str

@app.post("/predict", tags=["Predictions"])
def predict(features: PredictionFeatures):

    """
    Predict if a payment is fraud or not.

    This endpoint takes a set of features as input and returns a prediction
    of whether the payment is fraud or not.

    Parameters:
        features (PredictionFeatures): A set of features describing the payment

    Returns:
        dict: A dictionary containing the prediction of whether the payment is fraud or not
    """
    data = pd.DataFrame({
        "category" : [features.category],
        "amt" : [features.amt],
        "merch_fraud_level" : [features.merch_fraud_level],
        "city_fraud_level" : [features.city_fraud_level],
        "state_fraud_level" : [features.state_fraud_level],
        "zip_fraud_level" : [features.zip_fraud_level],
        "job_fraud_level" : [features.job_fraud_level],
        "age_fraud_level" : [features.age_fraud_level]
    })
    
    try:
        prediction = model.predict(data)[0]
        return {"prediction": int(prediction)}
    
    except Exception as e:
        return {"error": str(e)}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)