import streamlit as st
from streamlit_ace import st_ace
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
import os
import subprocess
import black
from pylint import lint
from io import StringIO
import sys
import torch
from huggingface_hub import hf_hub_url, cached_download, HfApi
import re
from typing import List, Dict
from streamlit_jupyter import StreamlitPatcher, tqdm

# This line should be at the top of your script
StreamlitPatcher().jupyter()  # This patches Streamlit to work in Jupyter

# Access Hugging Face API key from secrets
hf_token = st.secrets["hf_token"]
if not hf_token:
    st.error("Hugging Face API key not found. Please make sure it is set in the secrets.")

HUGGING_FACE_REPO_URL = "https://huggingface.co/spaces/acecalisto3/0shotTest"
PROJECT_ROOT = "projects"
AGENT_DIRECTORY = "agents"
AVAILABLE_CODE_GENERATIVE_MODELS = ["bigcode/starcoder", "Salesforce/codegen-350M-mono", "microsoft/CodeGPT-small"]

# Global state to manage communication between Tool Box and Workspace Chat App
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'terminal_history' not in st.session_state:
    st.session_state.terminal_history = []
if 'workspace_projects' not in st.session_state:
    st.session_state.workspace_projects = {}
if 'available_agents' not in st.session_state:
    st.session_state.available_agents = []

# AI Guide Toggle
ai_guide_level = st.sidebar.radio("AI Guide Level", ["Full Assistance", "Partial Assistance", "No Assistance"])

class TextGenerationTool:
    def __init__(self, llm: str):
        self.llm = llm
        self.tokenizer = AutoTokenizer.from_pretrained(llm)
        self.model = AutoModelForCausalLM.from_pretrained(llm)

    def generate_text(self, prompt: str, max_length: int = 50) -> str:
        inputs = self.tokenizer(prompt, return_tensors="pt")
        outputs = self.model.generate(**inputs, max_length=max_length)
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

class AIAgent:
    def __init__(self, name: str, description: str, skills: List[str], llm: str):
        self.name = name
        self.description = description
        self.skills = skills
        self.text_gen_tool = TextGenerationTool(llm)
        self._hf_api = HfApi()  # Initialize HfApi here

    def generate_agent_response(self, prompt: str) -> str:
        return self.text_gen_tool.generate_text(prompt)

    def create_agent_prompt(self) -> str:
        skills_str = '\n'.join([f"* {skill}" for skill in self.skills])
        agent_prompt = f"""
As an elite expert developer, my name is {self.name}. I possess a comprehensive understanding of the following areas:
{skills_str}
I am confident that I can leverage my expertise to assist you in developing and deploying cutting-edge web applications. Please feel free to ask any questions or present any challenges you may encounter.
"""
        return agent_prompt

    def autonomous_build(self, chat_history: List[tuple[str, str]], workspace_projects: Dict[str, Dict], 
                        project_name: str, selected_model: str, hf_token: str) -> tuple[str, str]:
        summary = "Chat History:\n" + "\n".join([f"User: {u}\nAgent: {a}" for u, a in chat_history])
        summary += "\n\nWorkspace Projects:\n" + "\n".join([f"{p}: {details}" for p, details in workspace_projects.items()])
        next_step = "Based on the current state, the next logical step is to implement the main application logic."
        return summary, next_step

    def deploy_built_space_to_hf(self, project_name: str) -> str:
        space_content = generate_space_content(project_name)
        repository = self._hf_api.create_repo(
            repo_id=project_name, 
            private=True,
            token=hf_token,
            exist_ok=True,
            space_sdk="streamlit"
        )
        self._hf_api.upload_file(
            path_or_fileobj=space_content,
            path_in_repo="app.py",
            repo_id=project_name,
            repo_type="space",
            token=hf_token
        )
        return repository.name

    def has_valid_hf_token(self) -> bool:
        return self._hf_api.whoami(token=hf_token) is not None

def process_input(input_text: str) -> str:
    chatbot = pipeline("text-generation", model="microsoft/DialoGPT-medium", tokenizer="microsoft/DialoGPT-medium", clean_up_tokenization_spaces=True)
    response = chatbot(input_text, max_length=50, num_return_sequences=1)[0]['generated_text']
    return response

def run_code(code: str) -> str:
    try:
        result = subprocess.run(code, shell=True, capture_output=True, text=True)
        return result.stdout
    except Exception as e:
        return str(e)

def workspace_interface(project_name: str) -> str:
    project_path = os.path.join(PROJECT_ROOT, project_name)
    if not os.path.exists(project_path):
        os.makedirs(project_path)
        st.session_state.workspace_projects[project_name] = {'files': []}
        return f"Project '{project_name}' created successfully."
    else:
        return f"Project '{project_name}' already exists."

def add_code_to_workspace(project_name: str, code: str, file_name: str) -> str:
    project_path = os.path.join(PROJECT_ROOT, project_name)
    if not os.path.exists(project_path):
        return f"Project '{project_name}' does not exist."
    
    file_path = os.path.join(project_path, file_name)
    with open(file_path, "w") as file:
        file.write(code)
    st.session_state.workspace_projects[project_name]['files'].append(file_name)
    return f"Code added to '{file_name}' in project '{project_name}'."

def display_chat_history(chat_history: List[tuple[str, str]]) -> str:
    return "\n".join([f"User: {u}\nAgent: {a}" for u, a in chat_history])

def display_workspace_projects(workspace_projects: Dict[str, Dict]) -> str:
    return "\n".join([f"{p}: {details}" for p, details in workspace_projects.items()])

def generate_space_content(project_name: str) -> str:
    # Logic to generate the Streamlit app content based on project_name
    # ... (This is where you'll need to implement the actual code generation)
    return "import streamlit as st\nst.title('My Streamlit App')\nst.write('Hello, world!')"

# Function to display the AI Guide chat
def display_ai_guide_chat(chat_history: List[tuple[str, str]]):
    st.markdown("<div class='chat-history'>", unsafe_allow_html=True)
    for user_message, agent_message in chat_history:
        st.markdown(f"<div class='chat-message user'>{user_message}</div>", unsafe_allow_html=True)
        st.markdown(f"<div class='chat-message agent'>{agent_message}</div>", unsafe_allow_html=True)
    st.markdown("</div>", unsafe_allow_html=True)

# Load the CodeGPT tokenizer explicitly
code_generator_tokenizer = AutoTokenizer.from_pretrained("microsoft/CodeGPT-small-py", clean_up_tokenization_spaces=True)
# Load the CodeGPT model for code completion
code_generator = pipeline("text-generation", model="microsoft/CodeGPT-small-py", tokenizer=code_generator_tokenizer)

def analyze_code(code: str) -> List[str]:
    hints = []
    
    # Example pointer: Suggest using list comprehensions
    if re.search(r'for .* in .*:\n\s+.*\.append\(', code):
        hints.append("Consider using a list comprehension instead of a loop for appending to a list.")
    
    # Example pointer: Recommend using f-strings for string formatting
    if re.search(r'\".*\%s\"|\'.*\%s\'', code) or re.search(r'\".*\%d\"|\'.*\%d\'', code):
        hints.append("Consider using f-strings for cleaner and more efficient string formatting.")
    
    # Example pointer: Avoid using global variables
    if re.search(r'\bglobal\b', code):
        hints.append("Avoid using global variables. Consider passing parameters or using classes.")
    
    # Example pointer: Recommend using `with` statement for file operations
    if re.search(r'open\(.+\)', code) and not re.search(r'with open\(.+\)', code):
        hints.append("Consider using the `with` statement when opening files to ensure proper resource management.")
    
    return hints

def get_code_completion(prompt: str) -> str:
    # Generate code completion based on the current code input
    # Use max_new_tokens instead of max_length
    completions = code_generator(prompt, max_new_tokens=50, num_return_sequences=1) 
    return completions[0]['generated_text']

def lint_code(code: str) -> List[str]:
    # Capture pylint output
    pylint_output = StringIO()
    sys.stdout = pylint_output
    
    # Run pylint on the provided code
    pylint.lint.Run(['--from-stdin'], do_exit=False, input=code)
    
    # Reset stdout
    sys.stdout = sys.__stdout__
    
    # Extract pylint messages
    messages = pylint_output.getvalue().splitlines()
    
    return messages

def format_code(code: str) -> str:
    # Format code using Black
    formatted_code = black.format_str(code, mode=black.FileMode())
    return formatted_code

def main():
    st.title("Streamlit Workspace")

    # Load agents from the agent directory
    agent_files = [f for f in os.listdir(AGENT_DIRECTORY) if f.endswith(".py")]
    for agent_file in agent_files:
        agent_module = __import__(f"{AGENT_DIRECTORY}.{os.path.splitext(agent_file)[0]}")
        agent_class = getattr(agent_module, os.path.splitext(agent_file)[0])
        agent_instance = agent_class()
        st.session_state.available_agents.append(agent_instance)

    # Display the available agents
    st.subheader("Available Agents")
    for agent in st.session_state.available_agents:
        st.write(f"**{agent.name}**: {agent.description}")

    # Select an agent
    selected_agent = st.selectbox("Select an Agent", [agent.name for agent in st.session_state.available_agents])
    current_agent = next((agent for agent in st.session_state.available_agents if agent.name == selected_agent), None)

    # Display the agent's prompt
    if current_agent:
        st.subheader(f"{current_agent.name} Prompt")
        st.write(current_agent.create_agent_prompt())

    # Workspace Tab
    st.subheader("Workspace")
    workspace_tabs = st.tabs(["Chat", "Tool Box", "Projects"])

    with workspace_tabs[0]:
        # Chat Tab
        st.subheader("Chat with your Agent")
        user_input = st.text_input("Enter your message:")

        if user_input:
            st.session_state.chat_history.append((user_input, current_agent.generate_agent_response(user_input)))
            user_input = ""  # Clear the input field

        # Display chat history
        st.markdown(display_chat_history(st.session_state.chat_history))

        # AI Guide
        if ai_guide_level != "No Assistance":
            st.subheader("AI Guide")
            guide_chat_history = []
            if ai_guide_level == "Full Assistance":
                guide_chat_history.append((
                    "I'm building a Streamlit app to display data from a CSV file.",
                    "Great! Let's start by creating a new project in the workspace."
                ))
                guide_chat_history.append((
                    "Create a new project called 'data_app'.",
                    "Okay, I've created the project 'data_app'. What would you like to name the main file?"
                ))
                guide_chat_history.append((
                    "Name it 'app.py'.",
                    "Alright, I've added an empty 'app.py' file to the 'data_app' project. Now, let's add some code to read the CSV file."
                ))
                guide_chat_history.append((
                    "Add the following code to 'app.py':\n```python\nimport pandas as pd\nimport streamlit as st\n\ndata = pd.read_csv('data.csv')\nst.write(data)\n```",
                    "Excellent! Now you can run this code to see the data from your CSV file in the Streamlit app."
                ))
            elif ai_guide_level == "Partial Assistance":
                guide_chat_history.append((
                    "How can I read data from a CSV file in Streamlit?",
                    "You can use the `pandas` library to read the CSV file and then use `streamlit.write()` to display it."
                ))
            display_ai_guide_chat(guide_chat_history)

    with workspace_tabs[1]:
        # Tool Box Tab
        st.subheader("Tool Box")
        tool_tabs = st.tabs(["Code Editor", "Terminal", "Code Analysis"])

        with tool_tabs[0]:
            # Code Editor Tab
            st.subheader("Code Editor")
            code_editor = st_ace(
                placeholder="Write your code here...",
                height=300,
                theme="monokai",
                key="code_editor",
                language="python",
                auto_update=True
            )

            st.button("Run Code", on_click=lambda: st.write(run_code(code_editor)))

            # Code Completion
            st.subheader("Code Completion")
            completion_prompt = st.text_area("Enter code for completion:")
            if completion_prompt:
                completed_code = get_code_completion(completion_prompt)
                st.write(f"**Completion:** {completed_code}")

        with tool_tabs[1]:
            # Terminal Tab
            st.subheader("Terminal")
            terminal_input = st.text_input("Enter a command:")

            if terminal_input:
                st.session_state.terminal_history.append(terminal_input)
                st.write(run_code(terminal_input))
                terminal_input = ""  # Clear the input field

            # Display terminal history
            st.markdown("\n".join(st.session_state.terminal_history))

        with tool_tabs[2]:
            # Code Analysis Tab
            st.subheader("Code Analysis")
            code_to_analyze = st.text_area("Enter code to analyze:")
            if code_to_analyze:
                # Analyze code
                analysis_results = analyze_code(code_to_analyze)
                if analysis_results:
                    st.write("**Code Analysis Results:**")
                    for hint in analysis_results:
                        st.write(f"- {hint}")
                else:
                    st.write("No code analysis suggestions found.")

                # Lint code
                lint_results = lint_code(code_to_analyze)
                if lint_results:
                    st.write("**Linting Results:**")
                    for message in lint_results:
                        st.write(f"- {message}")
                else:
                    st.write("No linting issues found.")

                # Format code
                formatted_code = format_code(code_to_analyze)
                st.write("**Formatted Code:**")
                st.code(formatted_code, language="python")

    with workspace_tabs[2]:
        # Projects Tab
        st.subheader("Projects")
        project_name = st.text_input("Enter project name:")
        if st.button("Create Project"):
            st.write(workspace_interface(project_name))

        # Display existing projects
        st.markdown(display_workspace_projects(st.session_state.workspace_projects))

        # Add code to a project
        selected_project = st.selectbox("Select a project", list(st.session_state.workspace_projects.keys()))
        code_to_add = st.text_area("Enter code to add:")
        file_name = st.text_input("Enter file name:")
        if st.button("Add Code"):
            st.write(add_code_to_workspace(selected_project, code_to_add, file_name))

if __name__ == "__main__":
    main()