import streamlit as st
import pdfplumber
import re


# Function to extract text from PDF
def extract_text_from_pdf(pdf_path):
    try:
        with pdfplumber.open(pdf_path) as pdf:
            text = "".join([page.extract_text() for page in pdf.pages])
        return text
    except Exception as e:
        return str(e)


# Function to extract numeric rates
def extract_rates(text):
    rates = re.findall(r"\b\d+(?:\.\d+)?\b", text)  # Regex for numbers, including decimals
    return rates


# Function to calculate electricity bill
def calculate_bill(units, rate_per_unit):
    return units * rate_per_unit


# Streamlit App
def main():
    st.title("Tariff Rate Extractor and Electricity Bill Calculator")
    st.sidebar.title("Options")

    app_mode = st.sidebar.selectbox(
        "Choose the mode:",
        ["Extract Rates from PDF", "Calculate Electricity Bill"]
    )

    if app_mode == "Extract Rates from PDF":
        st.header("Upload a PDF to Extract Rates")
        pdf_file = st.file_uploader("Upload PDF", type=["pdf"])

        if pdf_file is not None:
            with open(f"/tmp/{pdf_file.name}", "wb") as f:
                f.write(pdf_file.getbuffer())

            # Extract text and rates
            text = extract_text_from_pdf(f"/tmp/{pdf_file.name}")
            rates = extract_rates(text)

            st.subheader("Extracted Numeric Rates")
            if rates:
                st.write(rates)
            else:
                st.write("No numeric rates found in the document.")

    elif app_mode == "Calculate Electricity Bill":
        st.header("Electricity Bill Calculator")

        units = st.number_input(
            "Enter the number of units consumed (kWh):", min_value=0.0, step=0.1
        )
        rate_per_unit = st.number_input(
            "Enter the rate per unit (e.g., 5.5):", min_value=0.0, step=0.1
        )

        if st.button("Calculate Bill"):
            total_bill = calculate_bill(units, rate_per_unit)
            st.success(f"Your total electricity bill is: ₹{total_bill:.2f}")


if __name__ == "__main__":
    main()