import streamlit as st
import os
from openai import OpenAI

# Set up NVIDIA API client
client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=os.environ.get("NVIDIA_API_KEY")
)
# Streamlit UI
st.title("AI-Powered Text Generation App")
st.write("Interact with an AI model to generate text based on your inputs.")

# Response specification features
st.markdown("## 🛠️ Response Specification Features")
st.markdown("*The expanders below are parameters that you can adjust to customize the AI response.*")

with st.expander("🎨 *Temperature (Creativity Control)*"):
    st.write("""
    This parameter controls the *creativity* of the AI's responses:
    - *0.0*: Always the same response (deterministic).  
    - *0.1 - 0.3*: Mostly factual and repetitive.  
    - *0.4 - 0.7*: Balanced between coherence and creativity.  
    - *0.8 - 1.0*: Highly creative but less predictable.
    """)

with st.expander("📏 *Max Tokens (Response Length)*"):
    st.write("Defines the maximum number of words/subwords in the response.")

with st.expander("🎯 *Top-p (Nucleus Sampling)*"):
    st.write("""
    Controls word diversity by sampling from top-probability tokens:
    - **High top_p + Low temperature** → More factual, structured responses.
    - **High top_p + High temperature** → More diverse, unexpected responses.
    """)

with st.expander("🔄 *Number of Responses*"):
    st.write("Specifies how many response variations the AI should generate.")

with st.expander("✅ *Fact-Checking*"):
    st.write("""
    - If *enabled*, AI prioritizes factual accuracy.  
    - If *disabled*, AI prioritizes creativity.
    """)

st.markdown("""
### 🔎 *Summary*  
- temperature → Adjusts *creativity vs accuracy*.  
- max_tokens → Defines *response length*.  
- top_p → Fine-tunes *word diversity*.  
- fact_check → Ensures *factual correctness* (but may reduce fluency).  
- num_responses → Generates *different variations* of the same prompt.
""")

# Function to query the AI model (based on your friend's code)
def query_ai_model(prompt, model="meta/llama-3.1-405b-instruct", temperature=0.7, max_tokens=512, top_p=0.9, fact_check=False, num_responses=1):
    responses = []
    
    try:
        if fact_check:
            prompt = "Ensure factual accuracy. " + prompt

        for _ in range(num_responses):  # Response loop for multiple responses
            completion = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens
            )
            response = completion.choices[0].message.content
            responses.append(response)

    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
    
    return responses  # Return a list of responses

# Input Fields for Streamlit UI
user_input = st.text_area("Your Prompt:", placeholder="Type something...")

# Dropdown Menus
output_format = st.selectbox("Select Output Format:", ["Story", "Poem", "Article", "Code"])
tone_style = st.selectbox("Select Tone/Style:", ["Formal", "Informal", "Humorous", "Technical"])

# Sliders
creativity_level = st.slider("Creativity Level:", min_value=0.0, max_value=1.0, value=0.7, step=0.1)
max_length = st.slider("Max Length (tokens):", min_value=100, max_value=1024, value=512, step=50)

# Numeric Inputs
num_responses = st.number_input("Number of Responses:", min_value=1, max_value=5, value=1, step=1)

# Checkboxes
enable_creativity = st.checkbox("Enable Creative Mode", value=True)
fact_checking = st.checkbox("Enable Fact-Checking")

# Button to generate response
if st.button("Generate Answer"):
    if user_input.strip():
        with st.spinner("Generating response..."):
            full_prompt = f"Format: {output_format}\nTone: {tone_style}\nPrompt: {user_input}"
            ai_responses = query_ai_model(
                full_prompt,
                temperature=creativity_level if enable_creativity else 0.2,
                max_tokens=max_length,
                top_p=0.9 if enable_creativity else 0.7,
                fact_check=fact_checking,
                num_responses=num_responses
            )

        st.success("AI Responses:")
        for i, response in enumerate(ai_responses, 1):
            st.markdown(f"### Response {i}")
            st.write(response)
    else:
        st.warning("Please enter a prompt before clicking the button.")