import streamlit as st
from groq import Groq
from langchain_groq import ChatGroq
from langchain import hub
from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader
from langchain.document_loaders import WebBaseLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
import bs4
import torch

# Initialize the Groq client
client = Groq(api_key="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")

def generate_response_groq(context, query):
    """Generate response using Groq API."""
    prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="llama3-8b-8192",
    )
    response = chat_completion.choices[0].message.content
    return response

def extract_text_from_url(url):
    """Extract text from URL using WebBaseLoader."""
    loader = WebBaseLoader(
        web_paths=(url,),
        bs_kwargs=dict(
            parse_only=bs4.SoupStrainer()
        ),
    )
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = text_splitter.split_documents(docs)
    text = ""
    for split in splits:
        text += split.page_content
    return text

# Set the page layout to wide for better UI space
st.set_page_config(page_title="IESCO Tariff Guide", layout="wide")

# UI
st.title("IESCO Tariff Guide")
st.markdown("<div class='main-content'>", unsafe_allow_html=True)

url = "https://iesco.com.pk/index.php/customer-services/tariff-guide#"
document_text = extract_text_from_url(url)
st.text_area("Extracted Text", document_text, height=200)

query = st.text_input("Enter your query")

if st.button("Submit"):
    if query:
        with st.spinner("Generating response..."):
            response = generate_response_groq(document_text, query)
            st.write("**Response:**")
            st.write(response)
    else:
        st.error("Please enter a query.")

st.markdown("</div>", unsafe_allow_html=True)