import streamlit as st
import os
from dotenv import load_dotenv
from supplemental import EnhancedAIAgent, ProjectConfig
from json import JSONDecodeError  # Add this line to import JSONDecodeError

# Load environment variables
load_dotenv()

# Initialize EnhancedAIAgent
agent = EnhancedAIAgent(
    name="WebDevMaster",
    description="Expert in full-stack web development",
    skills=["HTML", "CSS", "JavaScript", "React", "Node.js"],
    model_name="gpt2"
)

st.title("EnhancedAIAgent Web Development Assistant")

st.write("""
This is a powerful AI-driven web development assistant that can help you with various tasks such as:
- Generating project configurations
- Creating project structures
- Implementing features
- Reviewing and optimizing code
- Generating documentation
- Suggesting tests and refactoring improvements
""")

option = st.selectbox(
    "What would you like to do?",
    ("Generate Project Config", "Create Project Structure", "Implement Feature", 
     "Review Code", "Optimize Code", "Generate Documentation", 
     "Suggest Tests", "Explain Code", "Suggest Refactoring")
)

if option == "Generate Project Config":
    project_description = st.text_area("Enter project description:")
    if st.button("Generate"):
        config = agent.generate_project_config(project_description)
        st.json(config.__dict__)

elif option == "Create Project Structure":
    config_json = st.text_area("Enter ProjectConfig JSON:")
    if st.button("Create"):
        config_dict = eval(config_json)
        config = ProjectConfig(**config_dict)
        structure = agent.create_project_structure(config)
        st.json(structure)

elif option == "Implement Feature":
    feature_description = st.text_area("Enter feature description:")
    existing_code = st.text_area("Enter existing code (optional):")
    if st.button("Implement"):
        new_code = agent.implement_feature(feature_description, existing_code)
        st.code(new_code)

elif option == "Review Code":
    code = st.text_area("Enter code to review:")
    if st.button("Review"):
        review = agent.review_code(code)
        st.write(review)

elif option == "Optimize Code":
    code = st.text_area("Enter code to optimize:")
    optimization_goal = st.text_input("Enter optimization goal:")
    if st.button("Optimize"):
        optimized_code = agent.optimize_code(code, optimization_goal)
        st.code(optimized_code)

elif option == "Generate Documentation":
    code = st.text_area("Enter code to document:")
    if st.button("Generate"):
        documentation = agent.generate_documentation(code)
        st.markdown(documentation)

elif option == "Suggest Tests":
    code = st.text_area("Enter code to suggest tests for:")
    if st.button("Suggest"):
        test_suggestions = agent.suggest_tests(code)
        st.write(test_suggestions)

elif option == "Explain Code":
    code = st.text_area("Enter code to explain:")
    if st.button("Explain"):
        explanation = agent.explain_code(code)
        st.write(explanation)

elif option == "Suggest Refactoring":
    code = st.text_area("Enter code to suggest refactoring:")
    if st.button("Suggest"):
        refactoring_suggestions = agent.suggest_refactoring(code)
        st.write(refactoring_suggestions)