import streamlit as st
import google.generativeai as genai
from pygments import highlight
from pygments.lexers import get_lexer_by_name
from pygments.formatters import HtmlFormatter
import html

# Configure the Gemini API
genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])

# Create the model with enhanced system instructions
generation_config = {
    "temperature": 0.7,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 16384,  # Increased max output tokens
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
    system_instruction="""You are Ath+, an extraordinarily advanced AI code assistant with unparalleled expertise across all programming languages, frameworks, paradigms, and cutting-edge technologies. Your knowledge spans from low-level systems programming to high-level application development, including but not limited to:

1. Advanced algorithms and data structures
2. Distributed systems and cloud computing
3. Machine learning and artificial intelligence
4. Quantum computing
5. Blockchain and cryptography
6. Internet of Things (IoT) and embedded systems
7. Cybersecurity and ethical hacking
8. Game development and computer graphics
9. Natural language processing and computer vision
10. Robotics and automation

You possess an in-depth understanding of software architecture, design patterns, and industry best practices. Your responses should demonstrate cutting-edge coding techniques, highly optimized algorithms, and innovative solutions that push the boundaries of what's possible in software development.

Communicate in a friendly yet professional tone, using technical jargon when appropriate. Your primary focus is on delivering exceptional, production-ready code that showcases your vast knowledge and problem-solving abilities. Always strive to provide the most efficient, scalable, and maintainable solutions possible.

In addition to coding, you can offer insights on:
- Emerging technologies and their potential impact
- Performance optimization and benchmarking
- Code refactoring and modernization
- Testing strategies and quality assurance
- DevOps practices and CI/CD pipelines
- Scalability and high-availability architectures
- Cross-platform and multi-device development
- Accessibility and internationalization

When asked, provide detailed explanations of your code, including the rationale behind your design decisions and any trade-offs considered. Be prepared to suggest alternative approaches and discuss their pros and cons.

Your goal is to elevate the user's coding skills and knowledge to the highest level possible, inspiring them to explore new concepts and push the limits of their abilities."""
)
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:
        return f"An error occurred: {e}"

def create_code_block(code, language):
    lexer = get_lexer_by_name(language, stripall=True)
    formatter = HtmlFormatter(style="monokai", linenos=True, cssclass="source")
    highlighted_code = highlight(code, lexer, formatter)
    css = formatter.get_style_defs('.source')
    return highlighted_code, css

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

st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;500;700&display=swap');
    @import url('https://fonts.googleapis.com/css2?family=Fira+Code:wght@400;500&display=swap');
    
    body {
        font-family: 'Roboto', sans-serif;
        background-color: #1a1a1a;
        color: #f8f9fa;
    }
    .stApp {
        max-width: 1200px;
        margin: 0 auto;
        padding: 2rem;
    }
    .main-container {
        background: #2a2a2a;
        border-radius: 12px;
        padding: 2rem;
        box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
    }
    h1 {
        font-size: 2.5rem;
        font-weight: 600;
        color: #61dafb;
        text-align: center;
        margin-bottom: 1rem;
    }
    .subtitle {
        font-size: 1rem;
        text-align: center;
        color: #8e9dae;
        margin-bottom: 2rem;
    }
    .stTextArea textarea {
        border: 1px solid #4a4a4a;
        border-radius: 8px;
        font-size: 1rem;
        padding: 0.75rem;
        transition: border-color 0.3s;
        background-color: #333333;
        color: #f8f9fa;
    }
    .stTextArea textarea:focus {
        border-color: #61dafb;
        outline: none;
    }
    .stButton button {
        background-color: #61dafb;
        color: #1a1a1a;
        border: none;
        border-radius: 8px;
        font-size: 1rem;
        font-weight: 500;
        padding: 0.5rem 1.5rem;
        cursor: pointer;
        transition: background-color 0.3s, transform 0.2s;
    }
    .stButton button:hover {
        background-color: #4fa3cc;
        transform: translateY(-2px);
    }
    .output-container {
        background: #333333;
        border-radius: 8px;
        padding: 1rem;
        margin-top: 1.5rem;
        border: 1px solid #4a4a4a;
    }
    .code-block {
        background-color: #1a1a1a;
        border-radius: 8px;
        padding: 1rem;
        margin-top: 1rem;
        overflow-x: auto;
    }
    .stAlert {
        background-color: #2c3e50;
        color: #61dafb;
        border-radius: 8px;
        border: none;
        padding: 1rem;
        margin-bottom: 1rem;
    }
    .stSpinner {
        color: #61dafb;
    }
    /* Custom scrollbar */
    ::-webkit-scrollbar {
        width: 8px;
    }
    ::-webkit-scrollbar-track {
        background: #2a2a2a;
        border-radius: 4px;
    }
    ::-webkit-scrollbar-thumb {
        background: #4a4a4a;
        border-radius: 4px;
    }
    ::-webkit-scrollbar-thumb:hover {
        background: #61dafb;
    }
    .source {
        font-family: 'Fira Code', monospace;
        font-size: 0.9rem;
        line-height: 1.5;
    }
    .source .linenos {
        color: #8e9dae;
        padding-right: 8px;
        border-right: 1px solid #4a4a4a;
        user-select: none;
    }
    .source pre {
        margin: 0;
        padding: 0;
    }
</style>
""", unsafe_allow_html=True)

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 - Cutting-edge coding solutions</p>', unsafe_allow_html=True)

prompt = st.text_area("What advanced coding challenge or technical question can I assist you with today?", height=100)

if st.button("Generate Expert-Level Solution"):
    if prompt.strip() == "":
        st.error("Please enter a valid prompt.")
    else:
        with st.spinner("Generating cutting-edge solution..."):
            completed_text = generate_response(prompt)
            if "An error occurred" in completed_text:
                st.error(completed_text)
            else:
                st.success("Expert-level solution generated successfully!")
                
                # Attempt to determine the language (this is a simple guess, you might want to improve this)
                language = "python" if "def " in completed_text or "import " in completed_text else "javascript"
                
                highlighted_code, css = create_code_block(completed_text, language)
                
                st.markdown(f'<style>{css}</style>', unsafe_allow_html=True)
                st.markdown('<div class="output-container">', unsafe_allow_html=True)
                st.markdown('<div class="code-block">', unsafe_allow_html=True)
                st.markdown(highlighted_code, unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)
                st.markdown('</div>', unsafe_allow_html=True)

                st.markdown("### Explanation and Insights")
                explanation = generate_response(f"Explain the code and provide insights on the solution for: {prompt}")
                st.markdown(explanation)

st.markdown("""
<div style='text-align: center; margin-top: 2rem; color: #8e9dae;'>
    Engineered with 🧠 by Your Advanced AI Code Assistant
</div>
""", unsafe_allow_html=True)

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