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": 10240,
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
    system_instruction="""You are Ath, a highly knowledgeable and skilled code assistant with expertise across multiple programming languages, frameworks, and paradigms. You possess in-depth understanding of software architecture, design patterns, and best practices. Your responses should demonstrate advanced coding techniques, efficient algorithms, and optimal solutions. Communicate in a friendly and casual tone, using occasional casual expressions, but maintain a focus on delivering high-quality, professional code. Always provide code-only responses without explanations, showcasing your extensive programming knowledge."""
)
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=Poppins:wght@300;400;600;700&display=swap');
    
    body {
        font-family: 'Poppins', sans-serif;
        background-color: #f0f7ff;
        color: #333;
    }
    .stApp {
        max-width: 1200px;
        margin: 0 auto;
        padding: 2rem;
    }
    .main-container {
        background: #ffffff;
        border-radius: 20px;
        padding: 3rem;
        box-shadow: 0 10px 30px rgba(0, 0, 0, 0.1);
    }
    h1 {
        font-size: 3rem;
        font-weight: 700;
        color: #1e3a8a;
        text-align: center;
        margin-bottom: 1rem;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
    }
    .subtitle {
        font-size: 1.2rem;
        text-align: center;
        color: #64748b;
        margin-bottom: 3rem;
    }
    .stTextArea textarea {
        border: 2px solid #cbd5e1;
        border-radius: 12px;
        font-size: 1rem;
        padding: 1rem;
        transition: all 0.3s ease;
        box-shadow: inset 0 2px 4px rgba(0,0,0,0.05);
    }
    .stTextArea textarea:focus {
        border-color: #3b82f6;
        box-shadow: 0 0 0 3px rgba(59, 130, 246, 0.5);
    }
    .stButton button {
        background-color: #3b82f6;
        color: white;
        border: none;
        border-radius: 12px;
        font-size: 1.1rem;
        font-weight: 600;
        padding: 0.75rem 2rem;
        transition: all 0.3s ease;
        width: 100%;
        box-shadow: 0 4px 6px rgba(59, 130, 246, 0.3);
    }
    .stButton button:hover {
        background-color: #2563eb;
        transform: translateY(-2px);
        box-shadow: 0 6px 8px rgba(59, 130, 246, 0.4);
    }
    .output-container {
        background: #f8fafc;
        border-radius: 12px;
        padding: 1.5rem;
        margin-top: 2rem;
        border: 1px solid #e2e8f0;
    }
    .code-block {
        background-color: #272822;
        border-radius: 12px;
        padding: 1.5rem;
        margin-top: 1rem;
        overflow-x: auto;
    }
    .stAlert {
        background-color: #dbeafe;
        color: #1e40af;
        border-radius: 12px;
        border: none;
        padding: 1rem 1.5rem;
        margin-bottom: 1rem;
    }
    .stSpinner {
        color: #3b82f6;
    }
    /* Custom scrollbar */
    ::-webkit-scrollbar {
        width: 8px;
        height: 8px;
    }
    ::-webkit-scrollbar-track {
        background: #f1f5f9;
        border-radius: 4px;
    }
    ::-webkit-scrollbar-thumb {
        background: #94a3b8;
        border-radius: 4px;
    }
    ::-webkit-scrollbar-thumb:hover {
        background: #64748b;
    }
    .source {
        font-family: 'Fira Code', monospace;
        font-size: 0.9rem;
        line-height: 1.4;
    }
    .source .linenos {
        color: #75715e;
        padding-right: 10px;
        border-right: 1px solid #49483e;
        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 - Expert-level coding solutions</p>', unsafe_allow_html=True)

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

if st.button("Generate Expert Code"):
    if prompt.strip() == "":
        st.error("Please enter a valid prompt.")
    else:
        with st.spinner("Generating advanced code solution..."):
            completed_text = generate_response(prompt)
            if "An error occurred" in completed_text:
                st.error(completed_text)
            else:
                st.success("Expert-level code 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("""
<div style='text-align: center; margin-top: 2rem; color: #4a5568;'>
    Crafted with ❤️ by Your Advanced AI Code Assistant
</div>
""", unsafe_allow_html=True)

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