import os
import streamlit as st
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
import faiss
from groq import Groq
import json
import re

# Set up Groq client
client = Groq(api_key=os.environ.get('GroqApi'))

# Load the embedding model
model = SentenceTransformer('all-MiniLM-L6-v2')  # Open-source embedding model

# Initialize FAISS
dimension = 384  # Embedding dimension for 'all-MiniLM-L6-v2'
index = faiss.IndexFlatL2(dimension)

# Streamlit app
st.title("Electricity Bill Calculation App")

# File upload
uploaded_file = st.file_uploader("Upload your electricity bill PDF", type="pdf")

if uploaded_file:
    # Extract text from PDF
    reader = PdfReader(uploaded_file)
    text = " ".join([page.extract_text() for page in reader.pages])

    # Tokenize and chunk text
    sentences = text.split(". ")  # Simple sentence splitting
    embeddings = model.encode(sentences)

    # Store embeddings in FAISS
    faiss.normalize_L2(embeddings)
    index.add(embeddings)

    # Extract company and user type
    company = None
    if "LESCO" in text.upper():
        company = "LESCO"
    elif "FESCO" in text.upper():
        company = "FESCO"

    user_type = None
    if "PROTECTED" in text.upper():
        user_type = "Protected"
    elif "UNPROTECTED" in text.upper():
        user_type = "Unprotected"

    st.write(f"Detected Company: {company}")
    st.write(f"Detected User Type: {user_type}")

    if company and user_type:
        # Appliance usage input
        st.subheader("Appliance Usage Details")
        num_appliances = st.number_input("Number of appliances", min_value=1, max_value=20, step=1)
        
        appliance_data = []
        for i in range(num_appliances):
            st.write(f"Appliance {i + 1}")
            name = st.text_input(f"Name of Appliance {i + 1}", key=f"appliance_{i}")
            power = st.number_input(f"Power (Watts) of {name}", key=f"power_{i}")
            hours = st.number_input(f"Usage hours per day for {name}", key=f"hours_{i}")
            if name and power and hours:
                appliance_data.append({"name": name, "power": power, "hours": hours})

        if st.button("Calculate Bill"):
            # Calculate total units
            total_units = sum([(appliance["power"] * appliance["hours"] * 30) / 1000 for appliance in appliance_data])
            st.write(f"Total Units Consumed: {total_units:.2f} kWh")

            # Get tariff rate from Groq
            query_content = {
                "company": company,
                "user_type": user_type,
                "units": total_units,
            }
            chat_completion = client.chat.completions.create(
                messages=[{"role": "user", "content": json.dumps(query_content)}],
                model="llama3-8b-8192",
            )

            # Extract numeric value from the response
            response_content = chat_completion.choices[0].message.content.strip()

            try:
                # Attempt to parse the response as JSON or extract the numeric value
                if response_content.startswith('{') and response_content.endswith('}'):
                    # Parse JSON if valid
                    response_data = json.loads(response_content)
                    tariff_rate = response_data.get('tariff_rate', 0.0)  # Replace 'tariff_rate' with correct key
                else:
                    # Extract numeric value using regex
                    match = re.search(r"[-+]?\d*\.\d+|\d+", response_content)
                    tariff_rate = float(match.group()) if match else 0.0

                # Calculate and display the bill
                total_bill = tariff_rate * total_units
                st.write(f"Tariff Rate: {tariff_rate} PKR/kWh")
                st.write(f"Total Bill: {total_bill:.2f} PKR")

            except ValueError as e:
                st.error(f"Error parsing Groq response: {e}")
                st.write("Groq response content:")
                st.write(response_content)