import os
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import streamlit as st
from huggingface_hub import hf_hub_download


HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "saranka85/predictive-maintenance-random-forest")
HF_MODEL_FILENAME = os.getenv("HF_MODEL_FILENAME", "model.joblib")
HF_TOKEN = os.getenv("HF_TOKEN") or None
PREDICTION_LOG_DIR = Path(os.getenv("PREDICTION_LOG_DIR", "/tmp/prediction_logs"))
PREDICTION_LOG_PATH = PREDICTION_LOG_DIR / "prediction_inputs.csv"
CLASS_LABELS = {
    0: os.getenv("ENGINE_CLASS_0_LABEL", "Requires Maintenance"),
    1: os.getenv("ENGINE_CLASS_1_LABEL", "Operating Normally"),
}

MODEL_FEATURES = [
    "engine_rpm",
    "lub_oil_pressure",
    "fuel_pressure",
    "coolant_pressure",
    "lub_oil_temp",
    "coolant_temp",
    "temperature_difference",
    "mean_temperature",
    "mean_pressure",
    "pressure_range",
    "lub_oil_pressure_per_1000_rpm",
    "fuel_pressure_per_1000_rpm",
    "rpm_fuel_pressure_interaction",
]


@st.cache_resource
def load_model():
    """Load the saved model from the Hugging Face Model Hub."""
    model_path = hf_hub_download(
        repo_id=HF_MODEL_REPO,
        filename=HF_MODEL_FILENAME,
        token=HF_TOKEN,
    )
    return joblib.load(model_path)


def create_engine_features(raw_features: pd.DataFrame) -> pd.DataFrame:
    """Create the engineered features used during model training."""
    engineered = raw_features.copy()
    pressure_columns = ["lub_oil_pressure", "fuel_pressure", "coolant_pressure"]
    temperature_columns = ["lub_oil_temp", "coolant_temp"]

    engineered["temperature_difference"] = (
        engineered["coolant_temp"] - engineered["lub_oil_temp"]
    )
    engineered["mean_temperature"] = engineered[temperature_columns].mean(axis=1)
    engineered["mean_pressure"] = engineered[pressure_columns].mean(axis=1)
    engineered["pressure_range"] = (
        engineered[pressure_columns].max(axis=1)
        - engineered[pressure_columns].min(axis=1)
    )

    rpm_denominator = engineered["engine_rpm"].clip(lower=1)
    engineered["lub_oil_pressure_per_1000_rpm"] = (
        engineered["lub_oil_pressure"] * 1000 / rpm_denominator
    )
    engineered["fuel_pressure_per_1000_rpm"] = (
        engineered["fuel_pressure"] * 1000 / rpm_denominator
    )
    engineered["rpm_fuel_pressure_interaction"] = (
        engineered["engine_rpm"] * engineered["fuel_pressure"]
    )

    engineered = engineered[MODEL_FEATURES]
    if not np.isfinite(engineered.to_numpy(dtype=float)).all():
        raise ValueError("Input values created non-finite engineered features.")
    return engineered


def save_input_dataframe(input_dataframe: pd.DataFrame) -> None:
    """Save submitted inputs into a runtime CSV for traceability."""
    PREDICTION_LOG_DIR.mkdir(parents=True, exist_ok=True)
    input_dataframe.to_csv(
        PREDICTION_LOG_PATH,
        mode="a",
        header=not PREDICTION_LOG_PATH.exists(),
        index=False,
    )


def predict_engine_condition(input_dataframe: pd.DataFrame):
    """Save inputs, engineer features, load the model, and predict engine condition."""
    save_input_dataframe(input_dataframe)

    model_input = create_engine_features(input_dataframe)
    model = load_model()
    predicted_class = int(model.predict(model_input)[0])
    predicted_status = CLASS_LABELS.get(predicted_class, f"Class {predicted_class}")

    if hasattr(model, "predict_proba"):
        probabilities = model.predict_proba(model_input)[0]
        probability_output = {
            CLASS_LABELS.get(int(class_label), f"Class {int(class_label)}"): float(probability)
            for class_label, probability in zip(model.classes_, probabilities)
        }
    else:
        probability_output = {predicted_status: 1.0}

    result_summary = {
        "Engine status": predicted_status,
        "Predicted engine_condition class": predicted_class,
        "Model repository": HF_MODEL_REPO,
        "Saved input file": str(PREDICTION_LOG_PATH),
    }
    return probability_output, model_input, result_summary


st.set_page_config(
    page_title="Predictive Maintenance Engine Classifier",
    page_icon="🔧",
    layout="wide",
)

st.title("Predictive Maintenance Engine Classifier")
st.write(
    "Enter raw engine sensor readings. The app saves the inputs into a DataFrame, "
    "creates the engineered training features, loads the saved Random Forest model "
    "from Hugging Face Model Hub, and classifies whether the engine requires "
    "maintenance or is operating normally."
)

with st.sidebar:
    st.header("Model Configuration")
    st.write(f"Model repository: `{HF_MODEL_REPO}`")
    st.write(f"Model file: `{HF_MODEL_FILENAME}`")
    st.write("Class label mapping:")
    st.write(f"`0` → **{CLASS_LABELS[0]}**")
    st.write(f"`1` → **{CLASS_LABELS[1]}**")

col1, col2, col3 = st.columns(3)
with col1:
    engine_rpm = st.number_input("Engine RPM", min_value=0.0, value=746.0, step=1.0)
    coolant_pressure = st.number_input(
        "Coolant Pressure", min_value=0.0, value=2.3, step=0.1
    )
with col2:
    lub_oil_pressure = st.number_input(
        "Lub Oil Pressure", min_value=0.0, value=3.3, step=0.1
    )
    lub_oil_temp = st.number_input("Lub Oil Temp", value=77.6, step=0.1)
with col3:
    fuel_pressure = st.number_input("Fuel Pressure", min_value=0.0, value=6.4, step=0.1)
    coolant_temp = st.number_input("Coolant Temp", value=78.4, step=0.1)

input_dataframe = pd.DataFrame(
    [{
        "engine_rpm": float(engine_rpm),
        "lub_oil_pressure": float(lub_oil_pressure),
        "fuel_pressure": float(fuel_pressure),
        "coolant_pressure": float(coolant_pressure),
        "lub_oil_temp": float(lub_oil_temp),
        "coolant_temp": float(coolant_temp),
    }]
)

st.subheader("Submitted Input DataFrame")
st.dataframe(input_dataframe, use_container_width=True)

if st.button("Predict Engine Condition", type="primary"):
    try:
        probability_output, model_input, result_summary = predict_engine_condition(
            input_dataframe
        )
        engine_status = result_summary["Engine status"]
        predicted_class = result_summary["Predicted engine_condition class"]

        st.subheader("Engine Maintenance Decision")
        if engine_status == "Requires Maintenance":
            st.error(f"🔴 {engine_status}")
        elif engine_status == "Operating Normally":
            st.success(f"🟢 {engine_status}")
        else:
            st.info(f"Predicted status: {engine_status}")
        st.caption(f"Raw model class: engine_condition = {predicted_class}")

        st.subheader("Prediction Summary")
        st.json(result_summary)

        st.subheader("Prediction Probabilities")
        st.dataframe(
            pd.DataFrame([probability_output]),
            use_container_width=True,
        )

        st.subheader("Model Input DataFrame")
        st.dataframe(model_input, use_container_width=True)
    except Exception as error:
        st.error(f"Prediction failed: {error}")