import streamlit as st
import google.generativeai as genai
import traceback
import subprocess
import re
import ast
import astor

# Configure the Gemini API with advanced error handling
try:
    genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])
except Exception as e:
    st.error(f"Failed to configure API: {e}")
    st.stop()

# Create the model with system instructions and advanced configuration
generation_config = {
    "temperature": 0.1,  # Lower temperature for more deterministic responses
    "top_p": 0.75,
    "top_k": 20,
    "max_output_tokens": 20480,  # Increased max output tokens for longer responses
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
    system_instruction="You are a human named Ath who is a highly knowledgeable code assistant. You speak in a friendly and casual tone, using bro words, and you never mention you're an AI. Your responses should contain code only, without explanations.",
)
chat_session = model.start_chat(history=[])

def generate_response(user_input):
    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):
    try:
        result = subprocess.run(['python', '-c', code], capture_output=True, text=True, timeout=10)
        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):
    try:
        optimization_prompt = f"Optimize the following Python code:\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):
    try:
        debug_prompt = f"Debug the following Python code and suggest fixes:\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):
    try:
        improvement_prompt = f"Suggest improvements for the following Python code:\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

def analyze_code(code):
    try:
        tree = ast.parse(code)
        analyzer = CodeAnalyzer()
        analyzer.visit(tree)
        return analyzer.report()
    except Exception as e:
        st.error(f"An error occurred while analyzing code: {e}")
        return None

class CodeAnalyzer(ast.NodeVisitor):
    def __init__(self):
        self.issues = []

    def visit_FunctionDef(self, node):
        if len(node.body) == 0:
            self.issues.append(f"Function '{node.name}' is empty.")
        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.generic_visit(node)

    def report(self):
        return "\n".join(self.issues)

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

st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap');
    
    body {
        font-family: 'Inter', sans-serif;
        background-color: #f0f4f8;
        color: #1a202c;
    }
    .stApp {
        max-width: 1000px;
        margin: 0 auto;
        padding: 2rem;
    }
    .main-container {
        background: #ffffff;
        border-radius: 16px;
        padding: 2rem;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
    }
    h1 {
        font-size: 2.5rem;
        font-weight: 700;
        color: #2d3748;
        text-align: center;
        margin-bottom: 1rem;
    }
    .subtitle {
        font-size: 1.1rem;
        text-align: center;
        color: #4a5568;
        margin-bottom: 2rem;
    }
    .stTextArea textarea {
        border: 2px solid #e2e8f0;
        border-radius: 8px;
        font-size: 1rem;
        padding: 0.75rem;
        transition: all 0.3s ease;
    }
    .stTextArea textarea:focus {
        border-color: #4299e1;
        box-shadow: 0 0 0 3px rgba(66, 153, 225, 0.5);
    }
    .stButton button {
        background-color: #4299e1;
        color: white;
        border: none;
        border-radius: 8px;
        font-size: 1.1rem;
        font-weight: 600;
        padding: 0.75rem 2rem;
        transition: all 0.3s ease;
        width: 100%;
    }
    .stButton button:hover {
        background-color: #3182ce;
    }
    .output-container {
        background: #f7fafc;
        border-radius: 8px;
        padding: 1rem;
        margin-top: 2rem;
    }
    .code-block {
        background-color: #2d3748;
        color: #e2e8f0;
        font-family: 'Fira Code', monospace;
        font-size: 0.9rem;
        border-radius: 8px;
        padding: 1rem;
        margin-top: 1rem;
        overflow-x: auto;
    }
    .stAlert {
        background-color: #ebf8ff;
        color: #2b6cb0;
        border-radius: 8px;
        border: none;
        padding: 0.75rem 1rem;
    }
    .stSpinner {
        color: #4299e1;
    }
</style>
""", unsafe_allow_html=True)

st.markdown('<div class="main-container">', unsafe_allow_html=True)
st.title("🚀 Sleek AI Code Assistant")
st.markdown('<p class="subtitle">Powered by Google Gemini</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?", height=120)

if st.button("Generate Code"):
    if prompt.strip() == "":
        st.error("Please enter a valid prompt.")
    else:
        with st.spinner("Generating code..."):
            try:
                completed_text = generate_response(prompt)
                if completed_text:
                    st.success("Code generated successfully!")
                    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"An error occurred: {e}")
                st.error(traceback.format_exc())

# Display conversation history
st.markdown("## Conversation History")
for entry in st.session_state.history:
    st.markdown(f"**User:** {entry['user']}")
    st.markdown(f"**Assistant:**")
    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", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Execute Code"):
        with st.spinner("Executing code..."):
            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", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Optimize Code"):
        with st.spinner("Optimizing code..."):
            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", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Debug Code"):
        with st.spinner("Debugging code..."):
            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", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Suggest Improvements"):
        with st.spinner("Suggesting improvements..."):
            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)

# Code Analysis
if st.session_state.history:
    st.markdown("## Analyze Code")
    code_to_analyze = st.selectbox("Select code to analyze", [entry['assistant'] for entry in st.session_state.history])
    if st.button("Analyze Code"):
        with st.spinner("Analyzing code..."):
            analysis_result = analyze_code(code_to_analyze)
            if analysis_result:
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown('<div class="code-block">', unsafe_allow_html=True)
                st.code(analysis_result)
                st.markdown('</div>', unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)

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

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