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

# MongoDB and environment setup
import certifi
ca = certifi.where()
from dotenv import load_dotenv 
load_dotenv()

mongo_db_url = os.getenv("MONGO_DB_URL")

import pymongo 
from src.exception.exception import DeliveryTimeException
from src.logging.logger import logging
from src.pipeline.training_pipeline import TrainingPipeline
from src.utils.main_utils.utils import load_object
from src.utils.ml_utils.model.estimator import DeliveryPredictionModel

# MongoDB client setup
try:
    client = pymongo.MongoClient(mongo_db_url, tlsCAFile=ca)
    from src.constants.training_pipeline import DATA_INGESTION_COLLECTION_NAME
    from src.constants.training_pipeline import DATA_INGESTION_DATABASE_NAME
    database = client[DATA_INGESTION_DATABASE_NAME]
    collection = database[DATA_INGESTION_COLLECTION_NAME]
except Exception as e:
    logging.warning(f"MongoDB connection failed: {e}. Continuing without database.")
    client = None

# Feature list
TOP_12_FEATURES = [
    'multiple_deliveries', 'Road_traffic_density', 'Vehicle_condition', 'Delivery_person_Ratings',
    'distance_deliveries', 'Weather_conditions', 'Festival', 'distance_traffic', 'distance',
    'Delivery_person_Age', 'prep_traffic', 'City'
]

# Mapping dictionaries for dropdowns
WEATHER_CONDITIONS = {'Sunny': 0, 'Cloudy': 1, 'Fog': 2, 'Sandstorms': 3, 'Stormy': 4, 'Windy': 5}
TRAFFIC_DENSITY = {'Low': 0, 'Medium': 1, 'High': 2, 'Jam': 3}
VEHICLE_CONDITION = {'Poor': 0, 'Good': 1, 'Excellent': 2}
FESTIVAL = {'No': 0, 'Yes': 1}
CITY = {'Urban': 0, 'Semi-Urban': 1, 'Metropolitan': 2}

# Load model and scaler
try:
    model = joblib.load('final_model/model.pkl')
    preprocessor = joblib.load('final_model/preprocessor.pkl')
    scaler = preprocessor['scaler']
    expected_features = preprocessor['feature_names']
    logging.info("Model and scaler loaded successfully")
except Exception as e:
    logging.error(f"Error loading model or scaler: {e}")
    raise e

def train_model():
    """Function to train the model"""
    try:
        train_pipeline = TrainingPipeline()
        train_pipeline.run_pipeline()
        return "Training completed successfully!"
    except Exception as e:
        logging.error(f"Error in training pipeline: {e}")
        return f"Training failed: {str(e)}"

def predict_delivery_time(
    multiple_deliveries,
    road_traffic_density,
    vehicle_condition,
    delivery_person_ratings,
    distance_deliveries,
    weather_conditions,
    festival,
    distance_traffic,
    distance,
    delivery_person_age,
    prep_traffic,
    city
):
    """
    Predict delivery time based on input features
    """
    try:
        # Validate and convert inputs
        input_data = {
            'multiple_deliveries': float(multiple_deliveries),
            'Road_traffic_density': TRAFFIC_DENSITY[road_traffic_density],
            'Vehicle_condition': VEHICLE_CONDITION[vehicle_condition],
            'Delivery_person_Ratings': float(delivery_person_ratings),
            'distance_deliveries': float(distance_deliveries),
            'Weather_conditions': WEATHER_CONDITIONS[weather_conditions],
            'Festival': FESTIVAL[festival],
            'distance_traffic': float(distance_traffic),
            'distance': float(distance),
            'Delivery_person_Age': float(delivery_person_age),
            'prep_traffic': float(prep_traffic),
            'City': CITY[city]
        }

        # Create DataFrame and ensure correct feature order
        df = pd.DataFrame([input_data])
        df = df[expected_features]

        # Scale features
        features_scaled = scaler.transform(df)
        
        # Make prediction
        if hasattr(model, 'predict'):
            prediction = model.predict(features_scaled)
        else:
            prediction = model.model.predict(features_scaled)
        
        # Convert to scalar
        if isinstance(prediction, np.ndarray):
            prediction = float(prediction[0])
        else:
            prediction = float(prediction)

        logging.info(f"Prediction: {prediction}")
        
        return f"🕒 Estimated Delivery Time: **{round(prediction, 2)} minutes** ({round(prediction/60, 2)} hours)"

    except Exception as e:
        logging.error(f"Error in prediction: {e}")
        return f"❌ Error: {str(e)}"

# Create Gradio Interface
with gr.Blocks(title="Delivery Time Prediction", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🚚 Delivery Time Prediction System
        ### Predict estimated delivery time based on various factors
        Enter the delivery details below to get an estimated delivery time.
        """
    )
    
    with gr.Tab("Prediction"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📦 Delivery Details")
                multiple_deliveries = gr.Number(
                    label="Multiple Deliveries",
                    value=1,
                    minimum=1,
                    maximum=10,
                    info="Number of deliveries in the order"
                )
                distance = gr.Number(
                    label="Distance (km)",
                    value=5.0,
                    minimum=0,
                    info="Total delivery distance"
                )
                distance_deliveries = gr.Number(
                    label="Distance for Deliveries (km)",
                    value=5.0,
                    minimum=0
                )
                distance_traffic = gr.Number(
                    label="Distance in Traffic (km)",
                    value=2.0,
                    minimum=0
                )
                prep_traffic = gr.Number(
                    label="Preparation Time in Traffic (min)",
                    value=10.0,
                    minimum=0
                )
                
            with gr.Column():
                gr.Markdown("### 👤 Delivery Person Details")
                delivery_person_age = gr.Number(
                    label="Delivery Person Age",
                    value=25,
                    minimum=18,
                    maximum=65
                )
                delivery_person_ratings = gr.Slider(
                    label="Delivery Person Rating",
                    minimum=1.0,
                    maximum=5.0,
                    value=4.5,
                    step=0.1
                )
                vehicle_condition = gr.Dropdown(
                    label="Vehicle Condition",
                    choices=list(VEHICLE_CONDITION.keys()),
                    value="Good"
                )
                
            with gr.Column():
                gr.Markdown("### 🌍 Environmental Conditions")
                weather_conditions = gr.Dropdown(
                    label="Weather Conditions",
                    choices=list(WEATHER_CONDITIONS.keys()),
                    value="Sunny"
                )
                road_traffic_density = gr.Dropdown(
                    label="Road Traffic Density",
                    choices=list(TRAFFIC_DENSITY.keys()),
                    value="Medium"
                )
                city = gr.Dropdown(
                    label="City Type",
                    choices=list(CITY.keys()),
                    value="Urban"
                )
                festival = gr.Radio(
                    label="Festival Day?",
                    choices=list(FESTIVAL.keys()),
                    value="No"
                )
        
        with gr.Row():
            predict_btn = gr.Button("🔮 Predict Delivery Time", variant="primary", size="lg")
            clear_btn = gr.ClearButton(components=[
                multiple_deliveries, road_traffic_density, vehicle_condition,
                delivery_person_ratings, distance_deliveries, weather_conditions,
                festival, distance_traffic, distance, delivery_person_age,
                prep_traffic, city
            ], value="🔄 Clear")
        
        output = gr.Markdown(label="Prediction Result")
        
        predict_btn.click(
            fn=predict_delivery_time,
            inputs=[
                multiple_deliveries, road_traffic_density, vehicle_condition,
                delivery_person_ratings, distance_deliveries, weather_conditions,
                festival, distance_traffic, distance, delivery_person_age,
                prep_traffic, city
            ],
            outputs=output
        )
    
    with gr.Tab("Training"):
        gr.Markdown(
            """
            ## 🎓 Model Training
            Click the button below to retrain the model with the latest data from MongoDB.
            """
        )
        train_btn = gr.Button("🚀 Start Training", variant="primary")
        train_output = gr.Textbox(label="Training Status", lines=3)
        
        train_btn.click(
            fn=train_model,
            inputs=[],
            outputs=train_output
        )
    
    with gr.Tab("About"):
        gr.Markdown(
            """
            ## 📊 About This Application
            
            This delivery time prediction system uses machine learning to estimate delivery times based on:
            
            - **Delivery Details**: Number of deliveries, distances
            - **Delivery Person**: Age, ratings, vehicle condition
            - **Environment**: Weather, traffic density, city type, festival days
            
            ### Features Used
            The model considers 12 key features to make accurate predictions:
            1. Multiple Deliveries
            2. Road Traffic Density
            3. Vehicle Condition
            4. Delivery Person Ratings
            5. Distance for Deliveries
            6. Weather Conditions
            7. Festival Day
            8. Distance in Traffic
            9. Total Distance
            10. Delivery Person Age
            11. Preparation Time in Traffic
            12. City Type
            
            ### How to Use
            1. Enter all the required delivery details in the **Prediction** tab
            2. Click the "Predict Delivery Time" button
            3. View the estimated delivery time in minutes and hours
            
            ### Model Training
            Use the **Training** tab to retrain the model with new data from the database.
            """
        )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,  # Default port for Hugging Face Spaces
        share=False
    )