import gradio as gr
import pandas as pd
import numpy as np
import os
import sys

sys.path.append('.')

try:
    from model_loader import load_engine_model
    print("✅ Model loader imported!")
except ImportError as e:
    print(f"Import error for model_loader: {e}")
    class MockModel:
        def predict(self, X):
            return {'prediction': 0, 'probability': 0.50, 'condition': 'Normal'}
        def predict_proba(self, X):
            return np.array([[0.50, 0.50]])
    def load_engine_model(model_name="dhani10/engine-condition-model"):
        print("⚠️ Using Mock Model.")
        return MockModel()

# Default model repo (adjust if needed)
MODEL_REPO = os.getenv("MODEL_REPO", "dhani10/engine-condition-model")
model = load_engine_model(MODEL_REPO)
print(f"Model loaded: {type(model)}")

def predict_condition(engine_rpm, oil_pressure, fuel_pressure, coolant_pressure, oil_temp, coolant_temp):
    input_data = {
        'Engine rpm': engine_rpm,
        'Lub oil pressure': oil_pressure,
        'Fuel pressure': fuel_pressure,
        'Coolant pressure': coolant_pressure,
        'lub oil temp': oil_temp,
        'Coolant temp': coolant_temp
    }
    result = model.predict(input_data)
    color = "green" if result.get('condition','Error') == 'Normal' else "red"
    emoji = "✅" if result.get('condition','Error') == 'Normal' else "⚠️"
    proba = result.get('probability', 0.0)
    return f"<h2 style='color:{color};'>{emoji} {result.get('condition','Error')}</h2><p>Confidence: {proba:.2%}</p>"

with gr.Blocks(title="Engine Maintenance Prediction") as demo:
    gr.Markdown("# 🚀 Engine Maintenance Prediction")
    gr.Markdown("Enter sensor readings to predict engine condition.")

    with gr.Row():
        with gr.Column():
            engine_rpm       = gr.Number(label="Engine RPM", value=791.0, step=1.0)
            oil_pressure     = gr.Number(label="Lub Oil Pressure (bar)", value=3.3, step=0.01)
            fuel_pressure    = gr.Number(label="Fuel Pressure (bar)", value=6.6, step=0.01)
            coolant_pressure = gr.Number(label="Coolant Pressure (bar)", value=2.3, step=0.01)
            oil_temp         = gr.Number(label="Lub Oil Temp (°C)", value=77.6, step=0.1)
            coolant_temp     = gr.Number(label="Coolant Temp (°C)", value=78.4, step=0.1)
            btn = gr.Button("Predict", variant="primary")
        with gr.Column():
            output = gr.HTML(label="Prediction Result")

    btn.click(
        predict_condition,
        inputs=[engine_rpm, oil_pressure, fuel_pressure, coolant_pressure, oil_temp, coolant_temp],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)