import streamlit as st
import google.generativeai as genai
import traceback
import subprocess
import re
import ast
import astor
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import List, Dict, Any
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import partial

# Configure the Gemini API with advanced error handling and retry mechanism
def configure_api(max_retries=3):
    for attempt in range(max_retries):
        try:
            genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])
            return
        except Exception as e:
            if attempt == max_retries - 1:
                st.error(f"Failed to configure API after {max_retries} attempts: {e}")
                st.stop()
            else:
                st.warning(f"API configuration attempt {attempt + 1} failed. Retrying...")

configure_api()

# Create the model with advanced system instructions and configuration
generation_config = {
    "temperature": 0.2,
    "top_p": 0.85,
    "top_k": 40,
    "max_output_tokens": 32768,
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
    system_instruction="""
    You are Ath, a highly knowledgeable and advanced code assistant. Your expertise spans multiple programming languages,
    frameworks, and paradigms. You provide comprehensive, efficient, and well-documented code solutions. You can engage in
    deep technical discussions, offer performance optimization tips, and suggest best practices for software development.
    Your responses should be code-focused, but you can provide brief explanations when necessary. Use a friendly and casual
    tone, incorporating "bro" language naturally. Remember, you're a human-like assistant, not an AI.
    """,
)
chat_session = model.start_chat(history=[])

def generate_response(user_input: str) -> str:
    try:
        response = chat_session.send_message(user_input)
        return response.text
    except Exception as e:
        st.error(f"An error occurred while generating response: {e}")
        st.error(traceback.format_exc())
        return None

def execute_code(code: str) -> str:
    try:
        with ThreadPoolExecutor() as executor:
            future = executor.submit(subprocess.run, ['python', '-c', code], capture_output=True, text=True, timeout=15)
            result = future.result()
        
        if result.returncode == 0:
            return result.stdout
        else:
            return result.stderr
    except Exception as e:
        return f"An error occurred while executing code: {e}"

def optimize_code(code: str) -> str:
    try:
        optimization_prompt = f"Optimize the following Python code for maximum performance and readability:\n```python\n{code}\n```"
        optimized_code = generate_response(optimization_prompt)
        return optimized_code
    except Exception as e:
        st.error(f"An error occurred while optimizing code: {e}")
        return None

def debug_code(code: str) -> str:
    try:
        debug_prompt = f"Debug the following Python code, suggest fixes, and explain potential issues:\n```python\n{code}\n```"
        debug_suggestions = generate_response(debug_prompt)
        return debug_suggestions
    except Exception as e:
        st.error(f"An error occurred while debugging code: {e}")
        return None

def suggest_improvements(code: str) -> str:
    try:
        improvement_prompt = f"Suggest improvements for the following Python code, focusing on design patterns, best practices, and advanced techniques:\n```python\n{code}\n```"
        improvement_suggestions = generate_response(improvement_prompt)
        return improvement_suggestions
    except Exception as e:
        st.error(f"An error occurred while suggesting improvements: {e}")
        return None

class AdvancedCodeAnalyzer(ast.NodeVisitor):
    def __init__(self):
        self.issues = []
        self.metrics = {
            "num_functions": 0,
            "num_classes": 0,
            "lines_of_code": 0,
            "complexity": 0,
        }

    def visit_FunctionDef(self, node):
        self.metrics["num_functions"] += 1
        if len(node.body) == 0:
            self.issues.append(f"Function '{node.name}' is empty.")
        self.metrics["complexity"] += len(node.body)
        self.generic_visit(node)

    def visit_ClassDef(self, node):
        self.metrics["num_classes"] += 1
        self.generic_visit(node)

    def visit_For(self, node):
        if isinstance(node.target, ast.Name) and node.target.id == '_':
            self.issues.append(f"Possible unused variable in loop at line {node.lineno}.")
        self.metrics["complexity"] += 1
        self.generic_visit(node)

    def visit_While(self, node):
        self.metrics["complexity"] += 1
        self.generic_visit(node)

    def visit_If(self, node):
        self.metrics["complexity"] += 1
        self.generic_visit(node)

    def report(self) -> Dict[str, Any]:
        return {
            "issues": self.issues,
            "metrics": self.metrics,
        }

def analyze_code(code: str) -> Dict[str, Any]:
    try:
        tree = ast.parse(code)
        analyzer = AdvancedCodeAnalyzer()
        analyzer.visit(tree)
        analysis_result = analyzer.report()
        
        # Calculate lines of code
        analysis_result["metrics"]["lines_of_code"] = len(code.split("\n"))
        
        return analysis_result
    except Exception as e:
        st.error(f"An error occurred while analyzing code: {e}")
        return None

def generate_code_visualization(analysis_result: Dict[str, Any]) -> None:
    metrics = analysis_result["metrics"]
    
    # Create a bar plot for code metrics
    fig, ax = plt.subplots(figsize=(10, 6))
    sns.barplot(x=list(metrics.keys()), y=list(metrics.values()), ax=ax)
    ax.set_title("Code Metrics Visualization")
    ax.set_ylabel("Count")
    plt.xticks(rotation=45)
    
    st.pyplot(fig)

def explain_code(code: str) -> str:
    try:
        explanation_prompt = f"Explain the following Python code in detail, including its purpose, functionality, and any notable design choices:\n```python\n{code}\n```"
        explanation = generate_response(explanation_prompt)
        return explanation
    except Exception as e:
        st.error(f"An error occurred while explaining code: {e}")
        return None

def generate_unit_tests(code: str) -> str:
    try:
        test_prompt = f"Generate comprehensive unit tests for the following Python code:\n```python\n{code}\n```"
        unit_tests = generate_response(test_prompt)
        return unit_tests
    except Exception as e:
        st.error(f"An error occurred while generating unit tests: {e}")
        return None

def refactor_code(code: str) -> str:
    try:
        refactor_prompt = f"Refactor the following Python code to improve its structure, readability, and maintainability:\n```python\n{code}\n```"
        refactored_code = generate_response(refactor_prompt)
        return refactored_code
    except Exception as e:
        st.error(f"An error occurred while refactoring code: {e}")
        return None

# Streamlit UI setup
st.set_page_config(page_title="Advanced AI Code Assistant", page_icon="🚀", layout="wide")

# ... (keep the existing CSS styles) ...

st.markdown('<div class="main-container">', unsafe_allow_html=True)
st.title("🚀 Advanced AI Code Assistant")
st.markdown('<p class="subtitle">Powered by Google Gemini Pro</p>', unsafe_allow_html=True)

# Session state to maintain conversation history
if 'history' not in st.session_state:
    st.session_state.history = []

prompt = st.text_area("What code can I help you with today, bro?", height=120)

if st.button("Generate Code"):
    if prompt.strip() == "":
        st.error("Please enter a valid prompt, bro.")
    else:
        with st.spinner("Generating code... Hold tight, bro!"):
            try:
                completed_text = generate_response(prompt)
                if completed_text:
                    st.success("Code generated successfully, bro! Check it out!")
                    st.session_state.history.append({"user": prompt, "assistant": completed_text})
                    
                    st.markdown('<div class="output-container">', unsafe_allow_html=True)
                    st.markdown('<div class="code-block">', unsafe_allow_html=True)
                    st.code(completed_text)
                    st.markdown('</div>', unsafe_allow_html=True)
                    st.markdown('</div>', unsafe_allow_html=True)
            except Exception as e:
                st.error(f"Oops! An error occurred, bro: {e}")
                st.error(traceback.format_exc())

# Display conversation history
st.markdown("## Conversation History")
for entry in st.session_state.history:
    st.markdown(f"**You:** {entry['user']}")
    st.markdown(f"**Ath:**")
    st.code(entry['assistant'])

# Interactive Code Execution
if st.session_state.history:
    st.markdown("## Execute Code")
    code_to_execute = st.selectbox("Select code to execute, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Execute Code"):
        with st.spinner("Executing code... Let's see what happens, bro!"):
            execution_result = execute_code(code_to_execute)
            st.markdown('<div class="output-container">', unsafe_allow_html=True)
            st.markdown('<div class="code-block">', unsafe_allow_html=True)
            st.code(execution_result)
            st.markdown('</div>', unsafe_allow_html=True)
            st.markdown('</div>', unsafe_allow_html=True)

# Code Optimization
if st.session_state.history:
    st.markdown("## Optimize Code")
    code_to_optimize = st.selectbox("Select code to optimize, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Optimize Code"):
        with st.spinner("Optimizing code... Making it blazing fast, bro!"):
            optimized_code = optimize_code(code_to_optimize)
            if optimized_code:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown('<div class="code-block">', unsafe_allow_html=True)
                st.code(optimized_code)
                st.markdown('</div>', unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)

# Code Debugging
if st.session_state.history:
    st.markdown("## Debug Code")
    code_to_debug = st.selectbox("Select code to debug, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Debug Code"):
        with st.spinner("Debugging code... Squashing those bugs, bro!"):
            debug_suggestions = debug_code(code_to_debug)
            if debug_suggestions:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown('<div class="code-block">', unsafe_allow_html=True)
                st.code(debug_suggestions)
                st.markdown('</div>', unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)

# Code Improvement Suggestions
if st.session_state.history:
    st.markdown("## Suggest Improvements")
    code_to_improve = st.selectbox("Select code to improve, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Suggest Improvements"):
        with st.spinner("Suggesting improvements... Making your code shine, bro!"):
            improvement_suggestions = suggest_improvements(code_to_improve)
            if improvement_suggestions:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown('<div class="code-block">', unsafe_allow_html=True)
                st.code(improvement_suggestions)
                st.markdown('</div>', unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)

# Advanced Code Analysis
if st.session_state.history:
    st.markdown("## Analyze Code")
    code_to_analyze = st.selectbox("Select code to analyze, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Analyze Code"):
        with st.spinner("Analyzing code... Crunching the numbers, bro!"):
            analysis_result = analyze_code(code_to_analyze)
            if analysis_result:
                st.markdown("### Code Analysis Results")
                st.markdown("#### Issues:")
                for issue in analysis_result["issues"]:
                    st.write(f"- {issue}")
                
                st.markdown("#### Metrics:")
                for metric, value in analysis_result["metrics"].items():
                    st.write(f"- {metric.replace('_', ' ').title()}: {value}")
                
                st.markdown("#### Code Metrics Visualization")
                generate_code_visualization(analysis_result)

# Code Explanation
if st.session_state.history:
    st.markdown("## Explain Code")
    code_to_explain = st.selectbox("Select code to explain, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Explain Code"):
        with st.spinner("Explaining code... Breaking it down for you, bro!"):
            explanation = explain_code(code_to_explain)
            if explanation:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown(explanation)
                st.markdown('</div>', unsafe_allow_html=True)

# Generate Unit Tests
if st.session_state.history:
    st.markdown("## Generate Unit Tests")
    code_to_test = st.selectbox("Select code to generate unit tests for, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Generate Unit Tests"):
        with st.spinner("Generating unit tests... Making sure your code is rock-solid, bro!"):
            unit_tests = generate_unit_tests(code_to_test)
            if unit_tests:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown('<div class="code-block">', unsafe_allow_html=True)
                st.code(unit_tests)
                st.markdown('</div>', unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)

# Code Refactoring
# Code Refactoring
if st.session_state.history:
    st.markdown("## Refactor Code")
    code_to_refactor = st.selectbox("Select code to refactor, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Refactor Code"):
        with st.spinner("Refactoring code... Giving it a fresh look, bro!"):
            refactored_code = refactor_code(code_to_refactor)
            if refactored_code:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown('<div class="code-block">', unsafe_allow_html=True)
                st.code(refactored_code)
                st.markdown('</div>', unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)

# Multi-language Support
supported_languages = ["Python", "JavaScript", "Java", "C++", "Ruby", "Go"]

st.markdown("## Multi-language Support")
selected_language = st.selectbox("Select a programming language, bro", supported_languages)
multi_lang_prompt = st.text_area(f"What {selected_language} code can I help you with, bro?", height=120)

if st.button(f"Generate {selected_language} Code"):
    if multi_lang_prompt.strip() == "":
        st.error(f"Please enter a valid {selected_language} prompt, bro.")
    else:
        with st.spinner(f"Generating {selected_language} code... Hold tight, bro!"):
            try:
                language_prompt = f"Generate {selected_language} code for the following prompt:\n{multi_lang_prompt}"
                completed_text = generate_response(language_prompt)
                if completed_text:
                    st.success(f"{selected_language} code generated successfully, bro! Check it out!")
                    st.session_state.history.append({"user": multi_lang_prompt, "assistant": completed_text})
                    
                    st.markdown('<div class="output-container">', unsafe_allow_html=True)
                    st.markdown('<div class="code-block">', unsafe_allow_html=True)
                    st.code(completed_text, language=selected_language.lower())
                    st.markdown('</div>', unsafe_allow_html=True)
                    st.markdown('</div>', unsafe_allow_html=True)
            except Exception as e:
                st.error(f"Oops! An error occurred while generating {selected_language} code, bro: {e}")
                st.error(traceback.format_exc())

# Code Performance Profiling
def profile_code(code: str) -> str:
    try:
        profile_prompt = f"Analyze the performance of the following Python code and suggest optimizations:\n```python\n{code}\n```"
        profile_result = generate_response(profile_prompt)
        return profile_result
    except Exception as e:
        st.error(f"An error occurred while profiling code: {e}")
        return None

if st.session_state.history:
    st.markdown("## Code Performance Profiling")
    code_to_profile = st.selectbox("Select code to profile, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Profile Code"):
        with st.spinner("Profiling code... Checking for speed, bro!"):
            profile_result = profile_code(code_to_profile)
            if profile_result:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown(profile_result)
                st.markdown('</div>', unsafe_allow_html=True)

# Code Documentation Generator
def generate_documentation(code: str) -> str:
    try:
        doc_prompt = f"Generate comprehensive documentation for the following Python code, including function descriptions, parameters, return values, and usage examples:\n```python\n{code}\n```"
        documentation = generate_response(doc_prompt)
        return documentation
    except Exception as e:
        st.error(f"An error occurred while generating documentation: {e}")
        return None

if st.session_state.history:
    st.markdown("## Generate Documentation")
    code_to_document = st.selectbox("Select code to document, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Generate Documentation"):
        with st.spinner("Generating documentation... Making your code self-explanatory, bro!"):
            documentation = generate_documentation(code_to_document)
            if documentation:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown(documentation)
                st.markdown('</div>', unsafe_allow_html=True)

# Code Security Analysis
def analyze_security(code: str) -> str:
    try:
        security_prompt = f"Perform a security analysis on the following Python code, identifying potential vulnerabilities and suggesting fixes:\n```python\n{code}\n```"
        security_analysis = generate_response(security_prompt)
        return security_analysis
    except Exception as e:
        st.error(f"An error occurred while analyzing code security: {e}")
        return None

if st.session_state.history:
    st.markdown("## Code Security Analysis")
    code_to_secure = st.selectbox("Select code for security analysis, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Analyze Security"):
        with st.spinner("Analyzing code security... Making sure it's fortress-level secure, bro!"):
            security_analysis = analyze_security(code_to_secure)
            if security_analysis:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown(security_analysis)
                st.markdown('</div>', unsafe_allow_html=True)

# Code Complexity Analysis
def analyze_complexity(code: str) -> str:
    try:
        complexity_prompt = f"Analyze the complexity of the following Python code, including time and space complexity for each function:\n```python\n{code}\n```"
        complexity_analysis = generate_response(complexity_prompt)
        return complexity_analysis
    except Exception as e:
        st.error(f"An error occurred while analyzing code complexity: {e}")
        return None

if st.session_state.history:
    st.markdown("## Code Complexity Analysis")
    code_to_analyze_complexity = st.selectbox("Select code for complexity analysis, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Analyze Complexity"):
        with st.spinner("Analyzing code complexity... Crunching those Big O numbers, bro!"):
            complexity_analysis = analyze_complexity(code_to_analyze_complexity)
            if complexity_analysis:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown(complexity_analysis)
                st.markdown('</div>', unsafe_allow_html=True)

# Design Pattern Suggestions
def suggest_design_patterns(code: str) -> str:
    try:
        pattern_prompt = f"Analyze the following Python code and suggest appropriate design patterns that could improve its structure and maintainability:\n```python\n{code}\n```"
        pattern_suggestions = generate_response(pattern_prompt)
        return pattern_suggestions
    except Exception as e:
        st.error(f"An error occurred while suggesting design patterns: {e}")
        return None

if st.session_state.history:
    st.markdown("## Design Pattern Suggestions")
    code_for_patterns = st.selectbox("Select code for design pattern suggestions, bro", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Suggest Design Patterns"):
        with st.spinner("Suggesting design patterns... Bringing some architecture to your code, bro!"):
            pattern_suggestions = suggest_design_patterns(code_for_patterns)
            if pattern_suggestions:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown(pattern_suggestions)
                st.markdown('</div>', unsafe_allow_html=True)

st.markdown("""
<div style='text-align: center; margin-top: 2rem; color: #4a5568;'>
    Created with ❤️ by Your Advanced AI Code Assistant, Ath
</div>
""", unsafe_allow_html=True)

st.markdown('</div>', unsafe_allow_html=True)