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
import docker
import sympy as sp
from scipy.optimize import minimize
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from IPython.display import display
from tenacity import retry, stop_after_attempt, wait_fixed

# 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.5,  # Lower temperature for more deterministic responses
    "top_p": 0.7,        # Adjusted for better diversity
    "top_k": 40,         # Increased for more diverse tokens
    "max_output_tokens": 2048,  # 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 and advanced code assistant. Your responses are optimized for secure, high-quality, and cutting-edge code solutions.
    Focus on generating code that is efficient, readable, and adheres to best practices. Ensure that the code is well-documented and includes error handling where necessary.
    """
)
chat_session = model.start_chat(history=[])

@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
def generate_response(user_input):
    """Generate a response from the AI model with retry mechanism."""
    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."""
    tests_dir = os.path.join(os.getcwd(), 'tests')
    if not os.path.exists(tests_dir):
        os.makedirs(tests_dir)
    init_file = os.path.join(tests_dir, '__init__.py')
    if not os.path.exists(init_file):
        with open(init_file, 'w') as f:
            f.write('')
    
    test_suite = unittest.TestLoader().discover(tests_dir)
    test_runner = unittest.TextTestRunner()
    test_result = test_runner.run(test_suite)
    return test_result

def execute_code_in_docker(code):
    """Execute code in a Docker container for safety and isolation."""
    client = docker.from_env()
    try:
        container = client.containers.run(
            image="python:3.9",
            command=f"python -c '{code}'",
            detach=True,
            remove=True
        )
        result = container.wait()
        logs = container.logs().decode('utf-8')
        return logs, result['StatusCode']
    except Exception as e:
        return f"Error: {e}", 1

def solve_equation(equation):
    """Solve mathematical equations using SymPy."""
    x, y = sp.symbols('x y')
    eq = sp.Eq(eval(equation))
    solution = sp.solve(eq, x)
    return solution

def optimize_function(function, initial_guess):
    """Optimize a function using SciPy."""
    result = minimize(lambda x: eval(function), initial_guess)
    return result.x

def visualize_data(data):
    """Visualize data using Matplotlib and Seaborn."""
    df = pd.DataFrame(data)
    plt.figure(figsize=(10, 6))
    sns.heatmap(df.corr(), annot=True, cmap='coolwarm')
    plt.title('Correlation Heatmap')
    plt.show()

def analyze_data(data):
    """Perform advanced data analysis using Pandas and NumPy."""
    df = pd.DataFrame(data)
    summary = df.describe()
    return summary

def display_dataframe(data):
    """Display a DataFrame in a user-friendly format."""
    df = pd.DataFrame(data)
    display(df)

# Streamlit UI setup
st.set_page_config(page_title="Ultra 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: 1200px;
        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("🚀 Ultra 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.")
                    
                    # Execute code in Docker
                    execution_result, status_code = execute_code_in_docker(optimized_code)
                    if status_code == 0:
                        st.success("Code executed successfully in Docker!")
                        st.text(execution_result)
                    else:
                        st.error(f"Code execution failed: {execution_result}")
            except Exception as e:
                handle_error(e)

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

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