import streamlit as st
from pymongo import MongoClient
from datetime import datetime
import pandas as pd

# MongoDB connection details
MONGO_URI = "mongodb+srv://kmrlogistics:trucking2024@cluster0.zkghx.mongodb.net/"
DB_NAME = "trucking"
COLLECTION_NAME_DRIVERS = "drivers"
COLLECTION_NAME_TRUCKS = "trucks"
COLLECTION_NAME_OWNER_OPERATORS = "owner_operators"

# Establish a connection to MongoDB
client = MongoClient(MONGO_URI)
db = client[DB_NAME]
drivers_collection = db[COLLECTION_NAME_DRIVERS]
trucks_collection = db[COLLECTION_NAME_TRUCKS]
owner_operators_collection = db[COLLECTION_NAME_OWNER_OPERATORS]

# User authentication function
def login(username, password):
    """Check the provided username and password."""
    return username == "KMR LOGISTICS" and password == "kmr2024"

# Initialize Streamlit session state for login
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False

# Streamlit app layout
st.title("Kmr Logistics Management")

# Login section
st.sidebar.header("User Login")

if not st.session_state.logged_in:
    username = st.sidebar.text_input("Username")
    password = st.sidebar.text_input("Password", type="password")

    if st.sidebar.button("Login"):
        if login(username, password):
            st.session_state.logged_in = True
            st.sidebar.success("Logged in successfully!")
        else:
            st.sidebar.error("Invalid username or password.")
else:
    st.sidebar.success("Logged in as kmr logistics.")

# Show the main application only if the user is logged in
if st.session_state.logged_in:
    # Page navigation
    page = st.sidebar.radio("Select a page", [
        "Add Truck",
        "Delete Truck",
        "Add Driver",
        "Delete Driver",
        "Add Owner Operator",
        "Delete Owner Operator",
        "Assign Truck",
        "Add Load",
        "Update/Delete Load",
        "View Owner Operator Info",
        "View Information"
    ])

    if page == "Add Truck":
        st.header("Add Truck Information")
        truck_name = st.text_input("Truck Name")
        truck_plate_number = st.text_input("Truck Plate Number")

        if st.button("Save Truck Info"):
            if truck_name and truck_plate_number:
                # Check if truck exists
                existing_truck = trucks_collection.find_one({"truck_plate_number": truck_plate_number})
                if existing_truck:
                    st.warning(f"Truck with Plate Number {truck_plate_number} already exists.")
                else:
                    truck_data = {
                        "truck_name": truck_name,
                        "truck_plate_number": truck_plate_number
                    }
                    trucks_collection.insert_one(truck_data)
                    st.success(f"Truck {truck_name} added successfully!")
            else:
                st.error("Please fill in all fields.")

    elif page == "Delete Truck":
        st.header("Delete Truck")
        trucks = list(trucks_collection.find({}, {"truck_name": 1, "truck_plate_number": 1}))
        truck_names = [f"{t.get('truck_name', '')} ({t.get('truck_plate_number', '')})" for t in trucks]

        selected_truck = st.selectbox("Select Truck to Delete", truck_names)

        if st.button("Delete Truck"):
            if selected_truck:
                truck_plate_number = selected_truck.split(" (")[1][:-1]
                trucks_collection.delete_one({"truck_plate_number": truck_plate_number})
                st.success(f"Truck {selected_truck} deleted successfully!")
            else:
                st.error("Please select a truck to delete.")

    elif page == "Add Driver":
        st.header("Add Driver Information")
        first_name = st.text_input("First Name")
        last_name = st.text_input("Last Name")
        contact_info = st.text_input("Contact Info")
        pay_per_mile = st.number_input("Driver Pay per Mile", min_value=0.0, format="%f")

        if st.button("Save Driver Info"):
            if first_name and last_name and contact_info and pay_per_mile:
                full_name = f"{first_name} {last_name}"

                # Check if driver exists
                existing_driver = drivers_collection.find_one({"contact_info": contact_info})
                if existing_driver:
                    st.warning(f"Driver with Contact Info {contact_info} already exists.")
                else:
                    driver_data = {
                        "first_name": first_name,
                        "last_name": last_name,
                        "contact_info": contact_info,
                        "pay_per_mile": pay_per_mile,
                        "trucks": [],
                        "loads": []
                    }
                    drivers_collection.insert_one(driver_data)
                    st.success(f"Driver {full_name} added successfully!")
            else:
                st.error("Please fill in all fields.")

    elif page == "Delete Driver":
        st.header("Delete Driver")
        drivers = list(drivers_collection.find({}, {"first_name": 1, "last_name": 1}))
        driver_names = [f"{d.get('first_name', '')} {d.get('last_name', '')}" for d in drivers]

        selected_driver = st.selectbox("Select Driver to Delete", driver_names)

        if st.button("Delete Driver"):
            if selected_driver:
                first_name, last_name = selected_driver.split()
                drivers_collection.delete_one({"first_name": first_name, "last_name": last_name})
                st.success(f"Driver {selected_driver} deleted successfully!")
            else:
                st.error("Please select a driver to delete.")

    elif page == "Add Owner Operator":
        st.header("Add Owner Operator")
        name = st.text_input("Owner Operator Name")
        contact_info = st.text_input("Contact Number")
        company_percentage = st.number_input("Company Percentage", min_value=0.0, max_value=100.0, format="%f")

        if st.button("Save Owner Operator Info"):
            if name and contact_info and company_percentage:
                # Check if owner operator exists
                existing_owner_operator = owner_operators_collection.find_one({"contact_info": contact_info})
                if existing_owner_operator:
                    st.warning(f"Owner Operator with Contact Number {contact_info} already exists.")
                else:
                    owner_operator_data = {
                        "name": name,
                        "contact_info": contact_info,
                        "company_percentage": company_percentage,
                        "trucks": [],
                        "loads": []
                    }
                    owner_operators_collection.insert_one(owner_operator_data)
                    st.success(f"Owner Operator {name} added successfully!")
            else:
                st.error("Please fill in all fields.")

    elif page == "Delete Owner Operator":
        st.header("Delete Owner Operator")
        owner_operators = list(owner_operators_collection.find({}, {"name": 1}))
        owner_operator_names = [o.get('name', '') for o in owner_operators]

        selected_owner_operator = st.selectbox("Select Owner Operator to Delete", owner_operator_names)

        if st.button("Delete Owner Operator"):
            if selected_owner_operator:
                owner_operators_collection.delete_one({"name": selected_owner_operator})
                st.success(f"Owner Operator {selected_owner_operator} deleted successfully!")
            else:
                st.error("Please select an owner operator to delete.")

    elif page == "Assign Truck":
        st.header("Assign Truck to Driver or Owner Operator")
        drivers = list(drivers_collection.find({}, {"first_name": 1, "last_name": 1}))
        owner_operators = list(owner_operators_collection.find({}, {"name": 1}))

        driver_names = [f"Driver: {d.get('first_name', '')} {d.get('last_name', '')}" for d in drivers]
        owner_operator_names = [f"Owner Operator: {o.get('name', '')}" for o in owner_operators]

        all_names = driver_names + owner_operator_names

        selected_name = st.selectbox("Select Driver or Owner Operator", all_names)
        trucks = list(trucks_collection.find({}, {"truck_name": 1, "truck_plate_number": 1}))
        truck_names = [f"{t.get('truck_name', '')} ({t.get('truck_plate_number', '')})" for t in trucks]

        selected_truck_name = st.selectbox("Select Truck", truck_names)

        if st.button("Assign Truck"):
            if selected_name and selected_truck_name:
                selected_truck = next(t for t in trucks if f"{t.get('truck_name', '')} ({t.get('truck_plate_number', '')})" == selected_truck_name)
                truck_plate_number = selected_truck["truck_plate_number"]

                if selected_name.startswith("Driver:"):
                    selected_driver = next(d for d in drivers if f"Driver: {d.get('first_name', '')} {d.get('last_name', '')}" == selected_name)
                    # Update driver record to include truck
                    drivers_collection.update_one(
                        {"_id": selected_driver["_id"]},
                        {"$addToSet": {"trucks": truck_plate_number}}
                    )
                    st.success(f"Truck {selected_truck_name} assigned to Driver {selected_name.split(': ')[1]} successfully!")

                elif selected_name.startswith("Owner Operator:"):
                    selected_owner_operator = next(o for o in owner_operators if f"Owner Operator: {o.get('name', '')}" == selected_name)
                    # Update owner operator record to include truck
                    owner_operators_collection.update_one(
                        {"_id": selected_owner_operator["_id"]},
                        {"$addToSet": {"trucks": truck_plate_number}}
                    )
                    st.success(f"Truck {selected_truck_name} assigned to Owner Operator {selected_name.split(': ')[1]} successfully!")
            else:
                st.error("Please select both a driver/owner operator and a truck.")

    elif page == "Add Load":
        st.header("Add Load Information")

        # Fetch drivers, owner operators, and trucks
        drivers = list(drivers_collection.find({}, {"first_name": 1, "last_name": 1, "trucks": 1, "contact_info": 1, "pay_per_mile": 1}))
        owner_operators = list(owner_operators_collection.find({}, {"name": 1, "trucks": 1, "contact_info": 1, "company_percentage": 1}))

        driver_names = [f"Driver: {d.get('first_name', '')} {d.get('last_name', '')}" for d in drivers]
        owner_operator_names = [f"Owner Operator: {o.get('name', '')}" for o in owner_operators]

        all_names = driver_names + owner_operator_names

        selected_names = st.multiselect("Select Driver(s) or Owner Operator(s)", all_names, max_selections=2)

        if selected_names:
            involved_drivers = []
            involved_owner_operators = []
            truck_numbers = []

            for selected_name in selected_names:
                if selected_name.startswith("Driver:"):
                    selected_driver = next(d for d in drivers if f"Driver: {d.get('first_name', '')} {d.get('last_name', '')}" == selected_name)
                    involved_drivers.append(selected_driver)
                    truck_numbers.extend(selected_driver.get('trucks', []))
                elif selected_name.startswith("Owner Operator:"):
                    selected_owner_operator = next(o for o in owner_operators if f"Owner Operator: {o.get('name', '')}" == selected_name)
                    involved_owner_operators.append(selected_owner_operator)
                    truck_numbers.extend(selected_owner_operator.get('trucks', []))

            trucks_display = [
                f"{t['truck_name']} ({t['truck_plate_number']})"
                for t in trucks_collection.find({"truck_plate_number": {"$in": truck_numbers}})
            ]

            selected_truck_display = st.selectbox("Select Truck for Load", trucks_display)

            if selected_truck_display:
                selected_truck_no = selected_truck_display.split(" (")[1][:-1]

                load_id = st.text_input("Load ID")
                date_from = st.date_input("Date From", datetime.now())
                date_to = st.date_input("Date To", datetime.now())
                original_miles = st.number_input("Original Miles", min_value=0)
                bobtail = st.number_input("Bobtail Miles", min_value=0)
                payout = st.number_input("Payout", min_value=0.0, format="%f")
                fuel_cost = st.number_input("Fuel Cost", min_value=0.0, format="%f")

                actual_miles = original_miles - bobtail

                driver_payments = []

                if involved_drivers:
                    for driver in involved_drivers:
                        driver_pay_per_mile = driver.get("pay_per_mile", 0)
                        driver_pay = actual_miles * driver_pay_per_mile / len(involved_drivers)  # Split pay between drivers
                        driver_payments.append(driver_pay)

                if involved_owner_operators:
                    for owner_operator in involved_owner_operators:
                        company_percentage = owner_operator.get("company_percentage", 0)
                        pay_percent = (payout * company_percentage) / 100
                        driver_payments.append(pay_percent)

                # Determine if the load is shared by both drivers or is separate
                if len(involved_drivers) == 2:
                    common_load = set(involved_drivers[0].get('trucks', [])) & set(involved_drivers[1].get('trucks', []))
                    if common_load:
                        st.subheader("Shared Load Calculation")
                        st.write(f"Drivers {involved_drivers[0]['first_name']} {involved_drivers[0]['last_name']} and {involved_drivers[1]['first_name']} {involved_drivers[1]['last_name']} share the load.")
                        total_driver_pay = sum(driver_payments)
                        st.write(f"Total Pay for Drivers: ${total_driver_pay:.2f}")
                    else:
                        st.subheader("Separate Load Calculation")
                        st.write("Drivers do not share a load.")
                        st.write(f"Driver 1 Pay: ${driver_payments[0]:.2f}")
                        st.write(f"Driver 2 Pay: ${driver_payments[1]:.2f}")
                else:
                    st.subheader("Single Driver or Owner Operator Calculation")
                    total_driver_pay = sum(driver_payments)
                    st.write(f"Total Pay for Driver/Owner Operator: ${total_driver_pay:.2f}")

                if st.button("Save Load Info"):
                    # Check if the load ID already exists
                    existing_load = drivers_collection.find_one({"loads.load_id": load_id}) or owner_operators_collection.find_one({"loads.load_id": load_id})
                    if existing_load:
                        st.warning(f"Load ID {load_id} already exists. Please use a unique Load ID.")
                    else:
                        load_data = {
                            "truck_no": selected_truck_no,
                            "truck_name": selected_truck_display.split(" (")[0],
                            "load_id": load_id,
                            "date_from": date_from.strftime("%Y-%m-%d"),
                            "date_to": date_to.strftime("%Y-%m-%d"),
                            "original_miles": original_miles,
                            "bobtail": bobtail,
                            "actual_miles": actual_miles,
                            "payout": payout,
                            "fuel_cost": fuel_cost,
                            "driver_pay": total_driver_pay
                        }

                        # Assign the load to each driver involved
                        for driver in involved_drivers:
                            drivers_collection.update_one(
                                {"_id": driver["_id"]},
                                {"$push": {"loads": load_data}}
                            )
                        # Assign the load to each owner operator involved
                        for owner_operator in involved_owner_operators:
                            owner_operators_collection.update_one(
                                {"_id": owner_operator["_id"]},
                                {"$push": {"loads": load_data}}
                            )

                        st.success(f"Load information added successfully for the selected drivers/owner operators with Truck {selected_truck_display.split(' (')[0]}!")

    elif page == "Update/Delete Load":
        st.header("Update or Delete Load")

        load_id = st.text_input("Enter Load ID to Update or Delete")

        if load_id:
            # Search for the load across all drivers and owner operators
            driver_record = drivers_collection.find_one({"loads.load_id": load_id}, {"loads.$": 1, "first_name": 1, "last_name": 1})
            owner_operator_record = owner_operators_collection.find_one({"loads.load_id": load_id}, {"loads.$": 1, "name": 1})
            
            if driver_record:
                load_data = driver_record["loads"][0]
                entity_type = "driver"
                entity_name = f"{driver_record['first_name']} {driver_record['last_name']}"
            elif owner_operator_record:
                load_data = owner_operator_record["loads"][0]
                entity_type = "owner_operator"
                entity_name = owner_operator_record["name"]
            else:
                load_data = None

            if load_data:
                st.write(f"Load found for {entity_type.capitalize()}: {entity_name}")
                
                # Display load information
                load_data_display = {
                    "Truck No": load_data["truck_no"],
                    "Load ID": load_data["load_id"],
                    "Date From": load_data["date_from"],
                    "Date To": load_data["date_to"],
                    "Original Miles": load_data["original_miles"],
                    "Bobtail Miles": load_data["bobtail"],
                    "Payout": load_data["payout"],
                    "Fuel Cost": load_data["fuel_cost"],
                    "Driver/Owner Operator Pay": load_data["driver_pay"]
                }
                
                for key, value in load_data_display.items():
                    st.write(f"{key}: {value}")
                
                # Input fields for updating the load
                st.subheader("Update Load Information")
                date_from = st.date_input("Date From", datetime.strptime(load_data["date_from"], "%Y-%m-%d"))
                date_to = st.date_input("Date To", datetime.strptime(load_data["date_to"], "%Y-%m-%d"))
                original_miles = st.number_input("Original Miles", value=load_data["original_miles"], min_value=0)
                bobtail = st.number_input("Bobtail Miles", value=load_data["bobtail"], min_value=0)
                payout = st.number_input("Payout", value=load_data["payout"], min_value=0.0, format="%f")
                fuel_cost = st.number_input("Fuel Cost", value=load_data["fuel_cost"], min_value=0.0, format="%f")

                actual_miles = original_miles - bobtail

                if entity_type == "driver":
                    driver_pay = actual_miles * driver_record["loads"][0]["driver_pay"] / load_data["actual_miles"]

                elif entity_type == "owner_operator":
                    company_percentage = owner_operator_record["loads"][0].get("company_percentage", 0)
                    pay_percent = (payout * company_percentage) / 100
                    driver_pay = pay_percent

                if st.button("Update Load"):
                    # Update the load with new values
                    updated_load_data = {
                        "loads.$.date_from": date_from.strftime("%Y-%m-%d"),
                        "loads.$.date_to": date_to.strftime("%Y-%m-%d"),
                        "loads.$.original_miles": original_miles,
                        "loads.$.bobtail": bobtail,
                        "loads.$.actual_miles": actual_miles,
                        "loads.$.payout": payout,
                        "loads.$.fuel_cost": fuel_cost,
                        "loads.$.driver_pay": driver_pay
                    }

                    if entity_type == "driver":
                        drivers_collection.update_one(
                            {"_id": driver_record["_id"], "loads.load_id": load_id},
                            {"$set": updated_load_data}
                        )
                        st.success(f"Load ID {load_id} updated successfully!")
                    elif entity_type == "owner_operator":
                        owner_operators_collection.update_one(
                            {"_id": owner_operator_record["_id"], "loads.load_id": load_id},
                            {"$set": updated_load_data}
                        )
                        st.success(f"Load ID {load_id} updated successfully!")

                if st.button("Delete Load"):
                    # Remove the load from the driver's or owner operator's record
                    if entity_type == "driver":
                        drivers_collection.update_one(
                            {"_id": driver_record["_id"]},
                            {"$pull": {"loads": {"load_id": load_id}}}
                        )
                        st.success(f"Load ID {load_id} deleted successfully!")
                    elif entity_type == "owner_operator":
                        owner_operators_collection.update_one(
                            {"_id": owner_operator_record["_id"]},
                            {"$pull": {"loads": {"load_id": load_id}}}
                        )
                        st.success(f"Load ID {load_id} deleted successfully!")
            else:
                st.warning(f"No load found with Load ID {load_id}.")

    elif page == "View Owner Operator Info":
        st.header("View Owner Operator Information")

        owner_operators = list(owner_operators_collection.find({}))
        owner_operator_names = [o.get('name', '') for o in owner_operators]

        selected_owner_operator_name = st.selectbox("Select Owner Operator to View", owner_operator_names)

        if selected_owner_operator_name:
            owner_operator_record = owner_operators_collection.find_one({"name": selected_owner_operator_name})

            if owner_operator_record:
                st.write(f"**Name**: {owner_operator_record['name']}")
                st.write(f"**Contact Number**: {owner_operator_record['contact_info']}")
                st.write(f"**Company Percentage**: {owner_operator_record['company_percentage']}%")
                st.write(f"**Trucks Assigned**: {', '.join(owner_operator_record['trucks']) if owner_operator_record['trucks'] else 'None'}")

                # Date range input
                date_from = st.date_input("Date From", datetime.now())
                date_to = st.date_input("Date To", datetime.now())

                # Convert the date input from Streamlit to Pandas datetime
                date_from = pd.to_datetime(date_from)
                date_to = pd.to_datetime(date_to)

                # Display loads associated with this owner operator within the date range
                if "loads" in owner_operator_record and owner_operator_record["loads"]:
                    load_df = pd.DataFrame(owner_operator_record["loads"])
                    load_df["date_from"] = pd.to_datetime(load_df["date_from"])
                    load_df["date_to"] = pd.to_datetime(load_df["date_to"])

                    # Filter by date range
                    filtered_load_df = load_df[(load_df["date_from"] >= date_from) & (load_df["date_to"] <= date_to)]

                    if not filtered_load_df.empty:
                        st.subheader("Load Information")
                        st.dataframe(filtered_load_df)

                        # Total Calculations
                        total_payout = filtered_load_df["payout"].sum()
                        total_company_fee = filtered_load_df["driver_pay"].sum()
                        total_fuel_cost = filtered_load_df["fuel_cost"].sum()

                        st.subheader("Total Calculations")
                        st.write(f"Total Payout: ${total_payout:.2f}")
                        st.write(f"Total Company Fee: ${total_company_fee:.2f}")
                        st.write(f"Total Fuel Cost: ${total_fuel_cost:.2f}")
                    else:
                        st.write("No loads found in the selected date range.")
                else:
                    st.write("No loads associated with this Owner Operator.")

    elif page == "View Information":
        st.header("View Load Information by Truck Name")

        # Fetch truck details
        trucks = list(trucks_collection.find({}, {"truck_name": 1, "truck_plate_number": 1}))
        truck_names = [f"{t.get('truck_name', '')} ({t.get('truck_plate_number', '')})" for t in trucks]

        # Dropdown to select a truck
        selected_truck_name_view = st.selectbox("Select Truck Name to View Records", truck_names)

        if selected_truck_name_view:
            truck_record = trucks_collection.find_one({"truck_name": selected_truck_name_view.split(" ")[0]})
            if truck_record:
                truck_plate_number = truck_record["truck_plate_number"]

                # Date range inputs
                date_from = st.date_input("Date From", datetime.now())
                date_to = st.date_input("Date To", datetime.now())

                # Convert date input to strings for MongoDB query
                date_from_str = date_from.strftime("%Y-%m-%d")
                date_to_str = date_to.strftime("%Y-%m-%d")

                # Aggregate load records for the selected truck and date range from both drivers and owner_operators collections
                driver_loads = list(drivers_collection.aggregate([
                    {"$match": {"trucks": truck_plate_number}},
                    {"$unwind": "$loads"},
                    {"$match": {
                        "loads.truck_no": truck_plate_number,
                        "loads.date_from": {"$gte": date_from_str},
                        "loads.date_to": {"$lte": date_to_str}
                    }},
                    {"$group": {
                        "_id": "$loads.load_id",
                        "load_data": {"$first": "$loads"},
                        "drivers": {"$push": {"first_name": "$first_name", "last_name": "$last_name"}}
                    }}
                ]))

                # Check if there are loads shared between drivers
                if driver_loads:
                    # Prepare data for table display
                    table_data = []
                    for load in driver_loads:
                        load_data = load['load_data']
                        driver_names = [f"{d['first_name']} {d['last_name']}" for d in load['drivers']]
                        
                        table_data.append({
                            "Load ID": load['_id'],
                            "Truck No": load_data['truck_no'],
                            "Truck Name": load_data['truck_name'],
                            "Date From": load_data['date_from'],
                            "Date To": load_data['date_to'],
                            "Original Miles": load_data['original_miles'],
                            "Bobtail Miles": load_data['bobtail'],
                            "Actual Miles": load_data['actual_miles'],
                            "Payout": f"${load_data['payout']:.2f}",
                            "Fuel Cost": f"${load_data['fuel_cost']:.2f}",
                            "Driver Pay": f"${load_data['driver_pay']:.2f}",
                            "Drivers Involved": ', '.join(driver_names)
                        })

                    # Display the table using a DataFrame
                    df = pd.DataFrame(table_data)
                    st.subheader("Load Information")
                    st.dataframe(df)  # Use st.table(df) for a static table

                else:
                    st.warning(f"No load records found for Truck {selected_truck_name_view.split(' ')[0]} between {date_from} and {date_to}.")