# Initialize Replit LLM Pipeline
try:
    nlp_pipeline = pipeline("text2text-generation", model=model_name)
    st.success("Model loaded successfully!")
except Exception as e:
    st.error(f"Error loading model: {e}. Please verify the model name or your internet connection.")
    st.stop()

# Input Section
st.subheader("Input Your Query")
user_input = st.text_area("Enter your query or task description", height=150)

# Submit Button
if st.button("Generate Output"):
    if user_input.strip() == "":
        st.warning("Please enter a valid input.")
    else:
        with st.spinner("Processing..."):
            try:
                # Generate response using Replit LLM
                output = nlp_pipeline(user_input)
                response = output[0]["generated_text"]
                st.subheader("AI Response")
                st.write(response)
            except Exception as e:
                st.error(f"An error occurred: {e}")

# Handling Code Optimization:
if task_type == "Code Optimization":
    st.info("Paste your Python code for optimization recommendations.")
    user_code = st.text_area("Paste your Python code", height=200)
    if st.button("Optimize Code"):
        if user_code.strip() == "":
            st.warning("Please paste valid Python code to optimize.")
        else:
            with st.spinner("Analyzing and optimizing..."):
                try:
                    optimization_prompt = f"Optimize the following Python code:\n\n{user_code}"
                    output = nlp_pipeline(optimization_prompt)
                    optimized_code = output[0]["generated_text"]
                    st.subheader("Optimized Code")
                    st.code(optimized_code)
                except Exception as e:
                    st.error(f"An error occurred while optimizing code: {e}")