import os
import streamlit as st
from app_functions import Get_DialoGPT_Response, Get_DistilGPT_Response, Get_MedGPT_Response  # Importing functions

st.set_page_config(
    page_title="Healthwise AI Assistant",
    page_icon="🚑",
    layout="centered",
    initial_sidebar_state="auto"
)

st.header("Healthwise AI Assistant 🚑")

# Input fields
input_text = st.text_input("Enter your prompt here:")

col1, col2, col3 = st.columns(3)

with col1:
    no_words = st.text_input("Max no.of words for better response")

with col2:
    user_type = st.selectbox(
        'Answer for',
        ('Professional', 'Practitioner', 'General'),
        index=0
    )

with col3:
    model_type = st.selectbox(
        'Model Type',
        ('DialoGPT', 'DistilGPT', 'MedGPT'),
        index=0
    )

submit = st.button("Generate Response")

# Generate response
if submit:
    if not input_text or not no_words or not user_type:
        st.error("Please fill in all fields to generate a response.")
    else:
        if model_type == 'DialoGPT':
            response = Get_DialoGPT_Response(input_text, no_words, user_type)
        elif model_type == 'DistilGPT':
            response = Get_DistilGPT_Response(input_text, no_words, user_type)
        else:
            response = Get_MedGPT_Response(input_text, no_words, user_type)
        
        st.write(response)