import os
import streamlit as st
from datasets import load_dataset
from openai import OpenAI

hyperbolic_api_key = os.getenv("HYPERBOLIC_API_KEY");

client = OpenAI(
	base_url="https://router.huggingface.co/hyperbolic",
    api_key=hyperbolic_api_key
)

# Load the dataset
dataset = load_dataset("andreska/Adrega62Manual", split="train")

# Function to read the content from the dataset
def read_dataset(dataset):
    text = []
    for item in dataset:
        text.append(item['text'])
    return "\n".join(text)

context = read_dataset(dataset)

# Inject custom CSS
st.markdown(
    """
    <style>
        .scrollable-div {
            height: 390px;
            width: 100%;
            overflow-y: auto;
            padding: 10px;
            border: 1px solid #ccc;
        }
    
       .block-container {
            padding-top: 3rem;
            padding-bottom: 0rem;
            padding-left: 5rem;
            padding-right: 5rem;
        }
        .stButton > button {
            height: 27px;
            background-color: #f85900;
            color: white;
            border: none;
            transition: all 0.3s ease;
        }
        .stButton > button:hover {
            background-color: #F97A33;
            color: white;
            border: none;
            opacity: 1;
        }
        .stButton > button:active {
            background-color: #B84200;
            color: white;
            opacity: 1;
        }
    </style>
    """,
    unsafe_allow_html=True
)

placeholder = st.empty()
# Define the placeholder globally (outside columns)
if st.session_state and 'conversation' in st.session_state:
    placeholder.markdown(f'<div class="scrollable-div">{st.session_state.conversation}</div>', unsafe_allow_html=True)
else:
    placeholder.markdown(f'<div class="scrollable-div"><p>Welcome! I am your Adrega AI assistant</p></div>', unsafe_allow_html=True)

def handle_submit():
    user_input = st.session_state.user_input

    if user_input:
        messages = {
            "role": "user",
            "content": user_input
        }

        completion = client.chat.completions.create(
        	model="Qwen/Qwen2.5-72B-Instruct", 
        	messages=messages, 
        	max_tokens=500,
        )

        try:
            # Send the request to the Hyperbolic API
            response = completion.choices[0].message
            response.raise_for_status()  # Raise an error for bad status codes

            answer = response.json().get("output", "No response received.")
            placeholder.markdown(f'<div class="scrollable-div"><p>{answer}</p></div>', unsafe_allow_html=True)

            st.session_state.conversation = f"<p>{answer}</p>"
            placeholder.markdown(f'<div class="scrollable-div">{st.session_state.conversation}</div>', unsafe_allow_html=True)

        except requests.exceptions.RequestException as e:
            error_message = f"An error occurred: {str(e)}"
            placeholder.markdown(f'<div class="scrollable-div"><p>{error_message}</p></div>', unsafe_allow_html=True)

st.text_input('Ask me a question', key='user_input', on_change=handle_submit)

if st.button("Ask"):
    handle_submit()