app.py

import streamlit as st
import sqlite3
import hashlib
import os
import google.generativeai as genai
import zipfile
from git import Repo
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import requests
from fpdf import FPDF
import markdown
import tempfile
import time


# Database setup
DB_FILE = "users.db"

def create_user_table():
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS users (
            username TEXT PRIMARY KEY,
            password TEXT
        )
    """)
    conn.commit()
    conn.close()

def add_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    try:
        cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, hashed_password))
        conn.commit()
    except sqlite3.IntegrityError:
        st.error("Username already exists. Please choose a different username.")
    conn.close()

def authenticate_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    cursor.execute("SELECT * FROM users WHERE username = ? AND password = ?", (username, hashed_password))
    user = cursor.fetchone()
    conn.close()
    return user

def initialize_session_state():
    if "authenticated" not in st.session_state:
        st.session_state.authenticated = False
    if "username" not in st.session_state:
        st.session_state.username = None
    if "page" not in st.session_state:
        st.session_state.page = "login"
    if "current_project" not in st.session_state:
        st.session_state.current_project = None
    if "project_uploaded" not in st.session_state:
        st.session_state.project_uploaded = False

def main():
    st.title("SimplifAI")

    # Initialize session state
    initialize_session_state()

    # Initialize database
    create_user_table()

    # Page routing logic
    if st.session_state.page == "login":
        login_page()
    elif st.session_state.page == "project_staging":  # Updated from "workspace"
        project_staging_page()  # Updated function name
    elif st.session_state.page == "project_view":
        project_view_page()
    elif st.session_state.page == "generate_documentation":
        generate_documentation_page()
    elif st.session_state.page == "saved_documentation":
        saved_documentation_page()

def login_page():
    st.image("SimplifAI Logo.jpeg", use_container_width=True)
    
    st.subheader("Please Log In or Register to Continue")
    auth_mode = st.radio("Choose an Option", ["Log In", "Register"], horizontal=True)

    if auth_mode == "Log In":
        st.subheader("Log In")
        username = st.text_input("Username", key="login_username")
        password = st.text_input("Password", type="password", key="login_password")

        # Handle single-click login
        if st.button("Log In"):
            if authenticate_user(username, password):
                st.session_state.authenticated = True
                st.session_state.username = username
                st.session_state.page = "project_staging"  # Updated from "workspace"
                st.rerun()
            else:
                st.error("Invalid username or password. Please try again.")

    elif auth_mode == "Register":
        st.subheader("Register")
        username = st.text_input("Create Username", key="register_username")
        password = st.text_input("Create Password", type="password", key="register_password")

        # Handle single-click registration
        if st.button("Register"):
            if username and password:
                add_user(username, password)
                st.success("Account created successfully! You can now log in.")
            else:
                st.error("Please fill in all fields.")

def logout_user():
    """Logs out the user safely by clearing session state and navigating to the login page."""
    st.session_state.authenticated = False
    st.session_state.username = None
    st.session_state.page = "login"
    st.session_state.current_project = None  # Clear project data
    st.session_state.project_uploaded = False  # Reset project uploaded flag


def project_staging_page():
    # Sidebar with logout button
    st.sidebar.image("SimplifAI Logo Long.jpeg", use_container_width=True)
    st.sidebar.title(f"Hello, {st.session_state.username}!")
    if st.sidebar.button("Log Out"):
        logout_user()
        st.rerun()

    # User's folder for projects
    user_folder = os.path.join("user_projects", st.session_state.username)
    os.makedirs(user_folder, exist_ok=True)

    # Refresh project list dynamically
    projects = [d for d in os.listdir(user_folder) if os.path.isdir(os.path.join(user_folder, d))]

    # Display "Projects" dropdown
    selected_project = st.sidebar.selectbox("Projects", ["Select a project"] + projects)

    if selected_project != "Select a project":
        st.session_state.current_project = selected_project
        st.session_state.page = "project_view"
        st.rerun()

    # Display success message if a project was uploaded
    if st.session_state.project_uploaded:
        st.success(f"Project '{st.session_state.current_project}' uploaded successfully!")
        st.session_state.project_uploaded = False  # Reset after showing the message

    # Main content area
    st.subheader("Project Staging")  # Updated header
    st.write("You can create a new project by uploading files or folders, or by cloning a GitHub repository.")

    # Existing upload functionality...

    # User action selection
    action = st.radio("Choose an action", ["Upload Files or Folders", "Clone GitHub Repository"], horizontal=True)

    project_name = st.text_input("Enter a project name")

    if action == "Upload Files or Folders":
        st.subheader("Upload Files or Folders")
        uploaded_files = st.file_uploader(
            "Upload one or more files or a .zip archive for folders", accept_multiple_files=True
        )

        if uploaded_files and project_name:
            if st.button("Upload Project"):
                project_folder = os.path.join(user_folder, project_name)
                os.makedirs(project_folder, exist_ok=True)

                for uploaded_file in uploaded_files:
                    # Save uploaded .zip files or regular files
                    file_path = os.path.join(project_folder, uploaded_file.name)

                    with open(file_path, "wb") as f:
                        f.write(uploaded_file.getbuffer())

                    # If a .zip file is uploaded, extract its contents
                    if uploaded_file.name.endswith(".zip"):
                        try:
                            with zipfile.ZipFile(file_path, "r") as zip_ref:
                                zip_ref.extractall(project_folder)
                            os.remove(file_path)  # Remove the .zip file after extraction
                            st.success(f"Folder from {uploaded_file.name} extracted successfully!")
                        except zipfile.BadZipFile:
                            st.error(f"File {uploaded_file.name} is not a valid .zip file.")
                    else:
                        st.success(f"File {uploaded_file.name} saved successfully!")

                # Update session state and trigger a rerun
                st.session_state.current_project = project_name
                st.session_state.project_uploaded = True
                st.rerun()

    elif action == "Clone GitHub Repository":
        st.subheader("Clone GitHub Repository")
        repo_url = st.text_input("Enter the GitHub repository URL")

        if repo_url and project_name:
            if st.button("Upload Project"):
                project_folder = os.path.join(user_folder, project_name)
                os.makedirs(project_folder, exist_ok=True)

                try:
                    Repo.clone_from(repo_url, project_folder)
                    # Update session state and trigger a rerun
                    st.session_state.current_project = project_name
                    st.session_state.project_uploaded = True
                    st.rerun()
                except Exception as e:
                    st.error(f"Failed to clone repository: {e}")

#------------------------------------------------------------------------------------------------------------------------------------------------------------------------


# Configure Gemini API
gemini = os.getenv("GEMINI")
genai.configure(api_key=gemini)
model = genai.GenerativeModel("gemini-1.5-flash")

def read_project_files(project_path):
    """Reads all files in the project directory and its subdirectories."""
    file_paths = []
    for root, _, files in os.walk(project_path):
        for file in files:
            # Skip .git files or folders from GitHub clones
            if ".git" not in root:
                file_paths.append(os.path.join(root, file))
    return file_paths

def read_files(file_paths):
    """Reads content from a list of file paths."""
    file_contents = {}
    for file_path in file_paths:
        if os.path.exists(file_path):
            try:
                # Attempt to read the file as UTF-8
                with open(file_path, 'r', encoding='utf-8') as file:
                    file_contents[file_path] = file.read()
            except UnicodeDecodeError:
                print(f"Skipping binary or non-UTF-8 file: {file_path}")
        else:
            print(f"File not found: {file_path}")
    return file_contents


def generate_prompt(file_contents, functionality_description):
    """Generates a prompt for Gemini to analyze the files."""
    prompt = "Analyze the following code files to identify all functions required to implement the functionality: "
    prompt += f"'{functionality_description}'.\n\n"

    for file_path, content in file_contents.items():
        prompt += f"File: {os.path.basename(file_path)}\n{content}\n\n"

    prompt += "For each relevant function, provide:\n"
    prompt += "1. Which file the function is found in.\n"
    prompt += "2. The function name.\n"
    prompt += "3. Dependencies on other functions or modules.\n"

    prompt += """
            Return your output in the following format providing no commentary:
            Project Summary:
                <A summary of the project>
            Functionality:
                <a more precise version of the user defined funcionality>
            Functions:
                <for each file containing relevant functions>
                    <file name of the file that contains relevant functions>:
                        -<relevant function header from the current file>:
                            -Function Dependencies:<function dependencies seperated by commas>
    """

    return prompt

def identify_required_functions(project_path, functionality_description):
    """Identifies required functions for a specified functionality."""
    # Gather all file paths in the project directory
    file_paths = read_project_files(project_path)

    # Read file contents
    file_contents = read_files(file_paths)

    # Generate a refined prompt for Gemini
    prompt = generate_prompt(file_contents, functionality_description)

    # Call the Gemini model
    response = model.generate_content(prompt)

    # Process and return the response
    return response.text

def generate_qwen_documentation(function_list, project_path, functionality_description):
    """
    Iteratively generates documentation for functions using Qwen and ensures all are fully documented.
    Args:
        function_list (list): Initial list of functions related to the specified functionality.
        project_path (str): Path to the project directory.
        functionality_description (str): Description of the functionality.
    Returns:
        dict: A dictionary with Qwen-generated documentation for each function.
    """
    completed_documentation = {}
    remaining_functions = function_list

    while remaining_functions:
        # Generate documentation for remaining functions using Qwen
        qwen_outputs = query_qwen_for_documentation(remaining_functions, project_path)

        # Clean and validate Qwen's output using Gemini
        completed, truncated = clean_and_validate_qwen_output(qwen_outputs, remaining_functions)

        # Update documentation and determine remaining functions
        completed_documentation.update(completed)
        remaining_functions = truncated  # Retry only truncated or missing functions

    return completed_documentation

def query_qwen_for_documentation(function_list, project_path):
    """
    Queries Hugging Face API to generate documentation for functions.
    Args:
        function_list (list): List of functions for documentation.
        project_path (str): Path to the project directory.
    Returns:
        dict: Hugging Face API's output for each function.
    """

    passkey = os.getenv("QWEN")  # Hugging Face API key from environment variables
    api_url = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-14B-Instruct"
    headers = {"Authorization": f"Bearer {passkey}"}

    outputs = {}
    for file_name, function in function_list:
        prompt = f"Generate detailed documentation for the following function in file `{file_name}`:\n{function}"
        payload = {"inputs": prompt, "parameters": {"max_length": 1024, "truncate": True}}

        # Send the request to the Hugging Face API
        try:
            response = requests.post(api_url, headers=headers, json=payload)
            response.raise_for_status()
            result = response.json()

            # Extract generated text from the API response
            if "generated_text" in result:
                outputs[(file_name, function)] = {"prompt": prompt, "output": result["generated_text"]}
            else:
                st.error(f"No generated text for function: {function}. Response: {result}")

        except requests.exceptions.RequestException as e:
            st.error(f"Error querying Hugging Face API for function {function}: {e}")
            outputs[(file_name, function)] = {"prompt": prompt, "output": None}

    return outputs


def extract_cleaned_gemini_output(gemini_output):
    """
    Removes the 'Functions' section and any irrelevant content from the Gemini output.
    Args:
        gemini_output (str): The raw output returned by Gemini.
    Returns:
        str: Cleaned output without the 'Functions' section.
    """
    lines = gemini_output.splitlines()
    cleaned_output = []
    skip_section = False

    for line in lines:
        line = line.strip()

        # Detect the start of the "Functions:" section and skip it
        if line.startswith("Functions:"):
            skip_section = True

        # Detect the end of the "Functions:" section and resume processing
        if skip_section and (line.startswith("Tasks:") or line.startswith("Project Summary:")):
            skip_section = False

        # Append lines outside the skipped section
        if not skip_section:
            cleaned_output.append(line)

    # Return the cleaned output with blank lines filtered out
    return "\n".join(line for line in cleaned_output if line)


def split_into_chunks(content, chunk_size=1000):
    """Splits large content into smaller chunks."""
    return [content[i:i + chunk_size] for i in range(0, len(content), chunk_size)]



def generate_detailed_documentation(file_contents, functionality_description):
    """
    Generates detailed documentation using Gemini directly.
    Args:
        file_contents (dict): A dictionary with file paths as keys and their content as values.
        functionality_description (str): A description of the functionality to document.
    Returns:
        str: The generated documentation.
    """
    prompt = f"""
    The following code files are provided. Analyze their contents and generate comprehensive documentation.
    Functionality description: '{functionality_description}'
    Tasks:
    1. Generate a project summary:
        '
        Project Summary:
        <Include project description and library or module dependencies>
        '
    2. Refine the user-defined functionality:
        '
        Functionality Summary:
        <Provide an enhanced description of user-specified functionality>
        '
    3. Describe the functionality flow:
        '
        Functionality Flow of Control:
        <In point form, explain the FLOW OF CONTROL for ONLY the enhanced description of the user-specified functionality, NOT the whole project. 
        Make sure there is no empty line between this content and 'Functionality Flow of Control'>
        '
    4. Generate detailed documentation ONLY for each function INVOLVED in the enhanced user-specified functionality:
        '
        Function Documentation:
        For each file that contains functions INVOLVED in the enhanced user-specified functiionality: (Do not include this line, take it as an instruction)
        File: <file name and extension>
        For each function in the file:
        - '<function name>':
          - Summary: <Description of the function's purpose>
          - Inputs: <Details of inputs and their types>
          - Outputs: <Details of outputs and their types>
          - Dependencies: <Dependencies on other modules/functions>
          - Data structures: <Details of data structures used>
          - Algorithmic Details: <Description of the algorithm used>
          - Error Handling: <Description of how the function handles errors>
          - Assumptions: <Any assumptions the function makes>
          - Example Usage: <Example demonstrating usage>
        '
    Please return ONLY the required documentation in the specified format, and do not write duplicate documentations.
    
    Code files:
    """
    for file_path, content in file_contents.items():
        prompt += f"\nFile: {os.path.basename(file_path)}\n{content}\n"

    response = model.generate_content(prompt)
    return process_gemini_output(response.text.strip())

def process_gemini_output(output):
    """
    Processes the raw output from Gemini to format it appropriately.

    - Removes ``` at the top and bottom of the output.
    - Removes all pairs of **.
    - Replaces ` used as a quote with '.
    - Replaces leading '*' with '-' for lists.

    Args:
        output (str): The raw Gemini output.

    Returns:
        str: The processed and cleaned output.
    """
    # Remove leading and trailing ```
    if output.startswith("```") or output.endswith("'''"):
        output = output[3:-3].strip()

    # Remove all pairs of **
    output = output.replace("**", "")

    # Replace ` with '
    output = output.replace("`", "'")

    # Replace leading '*' with '-' for list items
    lines = output.splitlines()
    processed_lines = []
    for line in lines:
        stripped_line = line.strip()
        if stripped_line.startswith("*"):
            # Only replace if the '*' is at the start or followed by whitespace
            processed_lines.append(line.replace("*", "-", 1))
        else:
            processed_lines.append(line)

    return "\n".join(processed_lines)


def generate_documentation_page():
    # Sidebar with "Log Out" and "Back to Project" buttons
    st.sidebar.image("SimplifAI Logo Long.jpeg", use_container_width=True)
    st.sidebar.title(f"Project: {st.session_state.current_project}")
    if st.sidebar.button("Back to Project"):
        st.session_state.page = "project_view"
        st.rerun()
    if st.sidebar.button("Log Out"):
        logout_user()
        st.rerun()

    st.subheader(f"Generate Documentation for {st.session_state.current_project}")
    st.write("Enter the functionality or parts of the project for which you'd like to generate documentation.")

    # Prompt user for functionality description
    functionality = st.text_area(
        "Describe the functionality",
        placeholder="e.g., Explain the function of the file `main.py`",
    )

    # Button to start analyzing functionality
    if st.button("Analyze"):
        if functionality.strip():
            with st.spinner("Analyzing project files... Please wait."):

                # Get the path of the current project
                user_folder = os.path.join("user_projects", st.session_state.username)
                project_folder = os.path.join(user_folder, st.session_state.current_project)

                if os.path.exists(project_folder):
                    try:
                        # Gather all file paths in the project directory
                        file_paths = read_project_files(project_folder)

                        # Read file contents
                        file_contents = read_files(file_paths)

                        # Generate documentation using Gemini
                        documentation = generate_detailed_documentation(file_contents, functionality)

                        # Save the documentation in session state for exporting and viewing
                        st.session_state.generated_documentation = documentation

                        # Display the final documentation
                        st.success("Documentation generated successfully!")
                        st.text_area("Generated Documentation", documentation, height=600)
                    except Exception as e:
                        st.error(f"An error occurred: {e}")
                else:
                    st.error("Project folder not found. Ensure the GitHub repository was cloned successfully.")
        else:
            st.error("Please enter the functionality to analyze.")

    # Add export/download buttons if documentation is available
    if "generated_documentation" in st.session_state and st.session_state.generated_documentation:
        documentation = st.session_state.generated_documentation

        # Define paths for PDF and Markdown files
        user_folder = os.path.join("user_projects", st.session_state.username)
        os.makedirs(user_folder, exist_ok=True)  # Ensure the user folder exists

        pdf_path = os.path.join(user_folder, f"{st.session_state.current_project}_Documentation.pdf")
        markdown_path = os.path.join(user_folder, f"{st.session_state.current_project}_Documentation.md")

        # Show appropriate buttons based on file generation status
        pdf_generated = os.path.exists(pdf_path)
        markdown_generated = os.path.exists(markdown_path)

        # Generate and replace buttons dynamically
        if not pdf_generated:
            if st.button("Generate PDF"):
                try:
                    generate_pdf(documentation, pdf_path)
                    st.session_state.pdf_generated = True  # Set a flag for PDF generation
                    st.rerun()  # Trigger a rerun to show the download button
                except Exception as e:
                    st.error(f"Failed to generate PDF: {e}")
        else:
            try:
                with open(pdf_path, "rb") as pdf_file:
                    st.download_button(
                        label="Download PDF",
                        data=pdf_file.read(),
                        file_name=f"{st.session_state.current_project}_Documentation.pdf",
                        mime="application/pdf",
                    )
            except FileNotFoundError:
                st.write("Click 'Generate PDF' to be able to download and store the documentation as a PDF")

        if not markdown_generated:
            if st.button("Generate Markdown File"):
                try:
                    generate_markdown_file(documentation, markdown_path)
                    st.session_state.markdown_generated = True  # Set a flag for Markdown generation
                    st.rerun()  # Trigger a rerun to show the download button
                except Exception as e:
                    st.error(f"Failed to generate Markdown file: {e}")
        else:
            try:
                with open(markdown_path, "rb") as markdown_file:
                    st.download_button(
                        label="Download Markdown File",
                        data=markdown_file.read(),
                        file_name=f"{st.session_state.current_project}_Documentation.md",
                        mime="text/markdown",
                    )
            except FileNotFoundError:
                st.write("Click 'Generate Markdown File' to be able to download and store the documentation as a markdown file")






# Helper function to generate PDF
def generate_pdf(documentation, pdf_path):
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
    pdf.set_font("Courier", size=12)  # Maintain the IDE-like font style

    # Add headers and content
    for line in documentation.splitlines():
        try:
            if line.startswith("- '") or line.startswith("File:"):
                pdf.set_font("Courier", style="B", size=12)  # Bold for specific lines
                pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1'))
            elif line.startswith("Project Summary:") or line.startswith("Functionality Summary:") or \
               line.startswith("Functionality Flow:") or line.startswith("Function Documentation:"):
                pdf.set_font("Courier", style="B", size=14)  # Bold larger headers
                pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1'))
            else:
                pdf.set_font("Courier", size=12)  # Regular for other lines
                pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1'))
        except UnicodeEncodeError as e:
            pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1'))

    # Save the PDF
    pdf.output(pdf_path)


# Helper function to generate Markdown file
def generate_markdown_file(documentation, output_path):
    """
    Generates a markdown file from the documentation with structured formatting.
    Args:
        documentation (str): The documentation content to format and save.
        output_path (str): Path to save the generated markdown file.
    """
    formatted_lines = []
    for line in documentation.splitlines():
        # Add bold formatting to specific lines (e.g., function names)
        if line.startswith("- '") and line.endswith("':"):
            formatted_lines.append(f"**{line.strip()}**")
        # Add section headers with markdown syntax
        elif line.startswith("Project Summary:") or line.startswith("Functionality Summary:") or \
             line.startswith("Functionality Flow:") or line.startswith("Function Documentation:"):
            formatted_lines.append(f"# {line.strip()}")
        else:
            # Regular text as-is
            formatted_lines.append(line)

    # Write the formatted lines to a markdown file
    with open(output_path, "w") as f:
        f.write("\n".join(formatted_lines))



#------------------------------------------------------------------------------------------------------------------------------------------------------------------------
def saved_documentation_page():
    # Sidebar with "Back to Project" and "Log Out" buttons
    st.sidebar.image("SimplifAI Logo Long.jpeg", use_container_width=True)
    st.sidebar.title(f"Project: {st.session_state.current_project}")
    if st.sidebar.button("Back to Project"):
        st.session_state.page = "project_view"
        st.rerun()
    if st.sidebar.button("Log Out"):
        logout_user()
        st.rerun()

    st.subheader(f"Saved Documentation for {st.session_state.current_project}")
    st.write("Below is a list of saved documentation for this project. Click to download.")

    # Define paths for generated files
    user_folder = os.path.join("user_projects", st.session_state.username)
    pdf_path = os.path.join(user_folder, f"{st.session_state.current_project}_Documentation.pdf")
    markdown_path = os.path.join(user_folder, f"{st.session_state.current_project}_Documentation.md")

    # Check if files exist
    pdf_exists = os.path.exists(pdf_path)
    markdown_exists = os.path.exists(markdown_path)

    # Display available documentation
    if not pdf_exists and not markdown_exists:
        st.info("No documentation has been generated yet. Please go to the 'Generate Documentation' page.")
    else:
        if pdf_exists:
            with open(pdf_path, "rb") as pdf_file:
                st.download_button(
                    label="Download PDF Documentation",
                    data=pdf_file.read(),
                    file_name=f"{st.session_state.current_project}_Documentation.pdf",
                    mime="application/pdf",
                )

        if markdown_exists:
            with open(markdown_path, "rb") as markdown_file:
                st.download_button(
                    label="Download Markdown Documentation",
                    data=markdown_file.read(),
                    file_name=f"{st.session_state.current_project}_Documentation.md",
                    mime="text/markdown",
                )


def project_view_page():
    # Sidebar with logout and return buttons
    st.sidebar.image("SimplifAI Logo Long.jpeg", use_container_width=True)
    st.sidebar.title(f"Project: {st.session_state.current_project}")
    if st.sidebar.button("Back to Project Staging"):
        st.session_state.page = "project_staging"
        st.rerun()
    if st.sidebar.button("Log Out"):
        logout_user()
        st.rerun()

    # Main content for project page
    st.subheader(f"Project: {st.session_state.current_project}")
    st.write("Manage your project and explore its files.")

    # Calculate project metrics
    user_folder = os.path.join("user_projects", st.session_state.username)
    project_folder = os.path.join(user_folder, st.session_state.current_project)

    # Count number of files
    num_files = sum([len(files) for _, _, files in os.walk(project_folder)])

    # Count total lines of code in all text files
    num_lines = 0
    for root, _, files in os.walk(project_folder):
        for file in files:
            file_path = os.path.join(root, file)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    num_lines += sum(1 for _ in f)
            except (UnicodeDecodeError, IsADirectoryError):
                continue

    # Display metrics side by side
    col1, col2 = st.columns(2)
    with col1:
        st.metric(label="Number of Files", value=num_files)
    with col2:
        st.metric(label="Lines of Code", value=num_lines)

    # Buttons for documentation functionality displayed side by side
    col3, col4, col5 = st.columns(3)
    with col3:
        if st.button("Generate Documentation"):
            st.session_state.page = "generate_documentation"
            st.rerun()
    with col4:
        if st.button("Saved Documentation"):
            st.session_state.page = "saved_documentation"
            st.rerun()
    with col5:
        if st.button("Show File Structure"):
            st.session_state.show_file_structure = not st.session_state.show_file_structure

    # Display file structure if toggled
    if "show_file_structure" not in st.session_state:
        st.session_state.show_file_structure = False

    if st.session_state.show_file_structure:
        def generate_tree(folder_path, prefix=""):
            """Generates a tree-like structure for the given folder."""
            items = sorted(os.listdir(folder_path))
            tree = ""
            for index, item in enumerate(items):
                item_path = os.path.join(folder_path, item)
                connector = "└── " if index == len(items) - 1 else "├── "
                tree += f"{prefix}{connector}{item}\n"
                if os.path.isdir(item_path):
                    subtree_prefix = prefix + ("    " if index == len(items) - 1 else "│   ")
                    tree += generate_tree(item_path, subtree_prefix)
            return tree

        # Display the tree structure
        st.write("File structure:")
        tree_structure = f".\n{generate_tree(project_folder)}"
        st.text(tree_structure)

if __name__ == "__main__":
    main()