import streamlit as st
import google.generativeai as genai
import requests
import subprocess
import os
import pylint
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
import git
import spacy
from spacy.lang.en import English
import boto3
import unittest

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

# Create the model with optimized parameters and enhanced system instructions
generation_config = {
    "temperature": 0.6,  # Lower temperature for more deterministic responses
    "top_p": 0.8,        # Adjusted for better diversity
    "top_k": 30,         # Increased for more diverse tokens
    "max_output_tokens": 16384,  # Increased for longer responses
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
    system_instruction="""
    You are Ath, 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 optimized, secure, and high-quality code only, without explanations. You are designed to provide accurate, efficient, and cutting-edge code solutions.
    """
)
chat_session = model.start_chat(history=[])

def generate_response(user_input):
    """Generate a response from the AI model."""
    try:
        response = chat_session.send_message(user_input)
        return response.text
    except Exception as e:
        return f"Error: {e}"

def optimize_code(code):
    """Optimize the generated code using static analysis tools."""
    with open("temp_code.py", "w") as file:
        file.write(code)
    result = subprocess.run(["pylint", "temp_code.py"], capture_output=True, text=True)
    os.remove("temp_code.py")
    return code

def fetch_from_github(query):
    """Fetch code snippets from GitHub."""
    # Placeholder for fetching code snippets from GitHub
    return ""

def interact_with_api(api_url):
    """Interact with external APIs."""
    response = requests.get(api_url)
    return response.json()

def train_ml_model(code_data):
    """Train a machine learning model to predict code improvements."""
    df = pd.DataFrame(code_data)
    X = df.drop('target', axis=1)
    y = df['target']
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
    model = RandomForestClassifier()
    model.fit(X_train, y_train)
    return model

def handle_error(error):
    """Handle errors and log them."""
    st.error(f"An error occurred: {error}")

def initialize_git_repo(repo_path):
    """Initialize or check the existence of a Git repository."""
    if not os.path.exists(repo_path):
        os.makedirs(repo_path)
    if not os.path.exists(os.path.join(repo_path, '.git')):
        repo = git.Repo.init(repo_path)
    else:
        repo = git.Repo(repo_path)
    return repo

def integrate_with_git(repo_path, code):
    """Integrate the generated code with a Git repository."""
    repo = initialize_git_repo(repo_path)
    with open(os.path.join(repo_path, "generated_code.py"), "w") as file:
        file.write(code)
    repo.index.add(["generated_code.py"])
    repo.index.commit("Added generated code")

def process_user_input(user_input):
    """Process user input using advanced natural language processing."""
    nlp = English()
    doc = nlp(user_input)
    return doc

def interact_with_cloud_services(service_name, action, params):
    """Interact with cloud services using boto3."""
    client = boto3.client(service_name)
    response = getattr(client, action)(**params)
    return response

def run_tests():
    """Run automated tests using unittest."""
    test_suite = unittest.TestLoader().discover('tests')
    test_runner = unittest.TextTestRunner()
    test_result = test_runner.run(test_suite)
    return test_result

# 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)

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:
                processed_input = process_user_input(prompt)
                completed_text = generate_response(processed_input.text)
                if "Error" in completed_text:
                    handle_error(completed_text)
                else:
                    optimized_code = optimize_code(completed_text)
                    st.success("Code generated and optimized successfully!")
                    
                    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)
                    
                    # Integrate with Git
                    repo_path = "./repo"  # Replace with your repository path
                    integrate_with_git(repo_path, optimized_code)
                    
                    # Run automated tests
                    test_result = run_tests()
                    if test_result.wasSuccessful():
                        st.success("All tests passed successfully!")
                    else:
                        st.error("Some tests failed. Please check the code.")
            except Exception as e:
                handle_error(e)

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)