import streamlit as st
from tariff_scraper import TARIFF_URLS, scrape_tariff_data

def calculate_bill(units_consumed, tariff_data):
    """
    Calculate the bill based on units consumed and tariff rates.
    Assumes the tariff data contains rates per unit in a structured format.
    """
    for row in tariff_data:
        if "to" in row.lower():  # Check if the row contains range information
            try:
                parts = row.split("|")
                unit_range = parts[0].strip()
                rate = float(parts[-1].strip())  # Assuming rate is the last column
                
                # Parse the unit range (e.g., "1-100" or ">300")
                if "-" in unit_range:
                    low, high = map(int, unit_range.split("-"))
                    if low <= units_consumed <= high:
                        return units_consumed * rate
                elif ">" in unit_range:
                    threshold = int(unit_range.strip(">").strip())
                    if units_consumed > threshold:
                        return units_consumed * rate
            except Exception:
                pass
    return 0.0  # Default to 0 if no matching rate is found

def main():
    st.set_page_config(page_title="Electricity Bill Calculator", layout="wide")
    st.title("Electricity Bill Calculator")

    # Sidebar for settings
    with st.sidebar:
        st.header("Settings")

        # Company selection
        st.subheader("Select Company")
        selected_company = st.selectbox("Company", options=list(TARIFF_URLS.keys()))

        if st.button("Fetch Tariff Rates", key="fetch_tariff"):
            url = TARIFF_URLS[selected_company]
            st.session_state['tariff_data'] = scrape_tariff_data(url)
            if st.session_state['tariff_data']:
                st.success(f"Tariff data fetched for {selected_company}!")
            else:
                st.error("Failed to fetch tariff data. Please check the URL or try again.")

        # Appliances selection
        st.subheader("Add Appliances")
        appliances = {
            "Fan": 75,  # Power in watts
            "Bulb": 40,
            "Refrigerator": 150,
            "Air Conditioner": 1500,
            "Washing Machine": 500,
            "Microwave": 1200
        }

        # Allow users to select predefined or add custom appliances
        selected_appliances = st.multiselect("Select Appliances", options=list(appliances.keys()))
        custom_appliance_name = st.text_input("Custom Appliance Name")
        custom_appliance_power = st.number_input("Custom Appliance Power (Watts)", min_value=1, step=1)

        # Add custom appliances
        if custom_appliance_name and custom_appliance_power:
            appliances[custom_appliance_name] = custom_appliance_power
            st.success(f"Added custom appliance: {custom_appliance_name} ({custom_appliance_power}W)")

    # Main application for usage inputs and bill calculation
    st.header("Appliance Usage")
    st.write("Enter usage details for the selected appliances:")

    total_units = 0
    for appliance in selected_appliances:
        st.subheader(f"{appliance} ({appliances[appliance]}W)")
        num_units = st.number_input(f"How many {appliance}s?", min_value=1, step=1, key=f"{appliance}_units")
        usage_hours = st.number_input(f"Daily usage hours for {appliance}:", min_value=1, step=1, key=f"{appliance}_hours")
        monthly_units = (appliances[appliance] / 1000) * num_units * usage_hours * 30  # Convert watts to kWh
        total_units += monthly_units
        st.write(f"Estimated Monthly Units for {appliance}: {monthly_units:.2f} kWh")

    st.write("### Total Estimated Monthly Units Consumed:", f"{total_units:.2f} kWh")

    # Ensure tariff data is fetched before bill calculation
    if st.button("Calculate Bill"):
        if 'tariff_data' in st.session_state and st.session_state['tariff_data']:
            bill = calculate_bill(total_units, st.session_state['tariff_data'])
            if bill > 0:
                st.success(f"Estimated Monthly Bill: Rs. {bill:.2f}")
            else:
                st.error("Could not calculate the bill. Please check the tariff data.")
        else:
            st.error("Please fetch the tariff rates first using the 'Fetch Tariff Rates' button.")

if __name__ == "__main__":
    # Initialize session state to store tariff data
    if 'tariff_data' not in st.session_state:
        st.session_state['tariff_data'] = None
    main()