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
from transformers import pipeline

class JSONDecodeError(Exception):
    pass

# Load environment variables
load_dotenv()

class JSONDecodeError(Exception):
    pass

def generate_project_config(project_description):
    generator = pipeline('text-generation', model='gpt2')
    
    prompt = f"Generate a project configuration for the following project description: {project_description}\n\n{{"
    
    try:
        response = generator(prompt, max_length=500, num_return_sequences=1)
        generated_text = response[0]['generated_text']
        
        # Extract the JSON part
        json_start = generated_text.index('{')
        json_end = generated_text.rindex('}') + 1
        config_str = generated_text[json_start:json_end]
        
        config = json.loads(config_str)
        return config
    except JSONDecodeError as e:
        st.error(f"Error decoding JSON: {str(e)}")
        return None
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
        return None

st.title("Project Configuration Generator")

project_description = st.text_area("Enter your project description:")

if st.button("Generate Configuration"):
    if project_description:
        config = generate_project_config(project_description)
        if config:
            st.json(config)
    else:
        st.warning("Please enter a project description.")

# 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)