import streamlit as st
import requests
import os

# API key from environment variable
API_KEY = os.environ.get("NEBIUS_API_KEY")

if not API_KEY:
    st.error("API key not found. Please set the `NEBIUS_API_KEY` environment variable.")

# Function to call Nebius API
def generate_response(prompt, api_key):
    api_url = "https://api.studio.nebius.ai/v1/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {
        "model": "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
        "messages": [
            {"role": "system", "content": "You are an AI that generates short, creative titles based on the user's input."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.6,
        "max_tokens": 512,
        "top_p": 0.9,
        "top_k": 50
    }
    
    response = requests.post(api_url, headers=headers, json=payload)
    
    if response.status_code == 200:
        return response.json()
    else:
        st.error(f"Error: {response.status_code}, {response.text}")
        return None

# Custom CSS for centering
st.markdown(
    """
    <style>
    .title-container {
        text-align: center;
        margin-bottom: 20px;
    }
    
    </style>
    """,
    unsafe_allow_html=True
)

# Centered title
#st.markdown('<div class="title-container"><h1>AI Title Generator</h1></div>', unsafe_allow_html=True)

# Input bar for user prompt
user_input = st.text_area(
    label="Input Text Area",  # Required but can be empty
    placeholder="Type or Pate Your Input......."
)



if st.button("Generate", use_container_width=True):
    if user_input.strip():
        with st.spinner("Generating... Please wait!"):
            result = generate_response(user_input, API_KEY)
        
        if result:
            try:
                # Extracting generated titles
                assistant_message = result["choices"][0]["message"]["content"]
                
                # Enhanced Output with Markdown
                st.markdown(
                    f"""
                    <div style="background-color:#000; padding:15px; border-radius:8px;">
                    <pre style="color:#000; font-family:monospace; white-space:pre-wrap;">{assistant_message}</pre>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
            except KeyError as e:
                st.error(f"Unexpected response format: {e}")
    else:
        st.warning("Please provide input before clicking Generate.")
st.markdown('</div>', unsafe_allow_html=True)