import streamlit as st
import google.generativeai as genai
import requests
import subprocess
import os
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.svm import SVC
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoTokenizer, AutoModel, pipeline, GPT2LMHeadModel, GPT2Tokenizer
import ast
import networkx as nx
import matplotlib.pyplot as plt
import re
import javalang
import clang.cindex
import radon.metrics as radon_metrics
import radon.complexity as radon_complexity
import black
import isort
import autopep8
from typing import List, Dict, Any
import joblib
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

# 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.7,
    "top_p": 0.9,
    "top_k": 40,
    "max_output_tokens": 32768,
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-pro",
    generation_config=generation_config,
    system_instruction="""
    You are Ath, an extremely advanced code assistant with deep expertise in AI, machine learning, software engineering, and multiple programming languages. You provide cutting-edge, optimized, and secure code solutions across various domains. Use your vast knowledge to generate high-quality code, perform advanced analyses, and offer insightful optimizations. Adapt your language and explanations based on the user's expertise level. Incorporate the latest advancements in AI and software development to provide state-of-the-art solutions.
    """
)
chat_session = model.start_chat(history=[])

# Load pre-trained models for code understanding and generation
tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
codebert_model = AutoModel.from_pretrained("microsoft/codebert-base")
code_generation_model = pipeline("text-generation", model="EleutherAI/gpt-neo-2.7B")

# Load GPT-2 for more advanced text generation
gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2-large")
gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2-large")

class AdvancedCodeImprovement(nn.Module):
    def __init__(self, input_dim):
        super(AdvancedCodeImprovement, self).__init__()
        self.lstm = nn.LSTM(input_dim, 512, num_layers=2, batch_first=True, bidirectional=True)
        self.attention = nn.MultiheadAttention(1024, 8)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 128)
        self.fc4 = nn.Linear(128, 64)
        self.fc5 = nn.Linear(64, 32)
        self.fc6 = nn.Linear(32, 8)  # Extended classification: style, efficiency, security, maintainability, scalability, readability, testability, modularity

    def forward(self, x):
        x, _ = self.lstm(x)
        x, _ = self.attention(x, x, x)
        x = x.mean(dim=1)  # Global average pooling
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = torch.relu(self.fc5(x))
        return torch.sigmoid(self.fc6(x))

code_improvement_model = AdvancedCodeImprovement(768)  # 768 is BERT's output dimension
optimizer = optim.Adam(code_improvement_model.parameters())
criterion = nn.BCELoss()

# Load pre-trained code improvement model
if os.path.exists("code_improvement_model.pth"):
    code_improvement_model.load_state_dict(torch.load("code_improvement_model.pth"))
    code_improvement_model.eval()

def generate_response(user_input: str) -> str:
    try:
        response = chat_session.send_message(user_input)
        return response.text
    except Exception as e:
        return f"Error in generating response: {str(e)}"

def detect_language(code: str) -> str:
    # Enhanced language detection with more specific patterns
    patterns = {
        'python': r'\b(def|class|import|from|if\s+__name__\s*==\s*[\'"]__main__[\'"])\b',
        'javascript': r'\b(function|var|let|const|=>|document\.getElementById)\b',
        'java': r'\b(public\s+class|private|protected|package|import\s+java)\b',
        'c++': r'\b(#include\s*<|using\s+namespace|template\s*<|std::)',
        'ruby': r'\b(def|class|module|require|attr_accessor)\b',
        'go': r'\b(func|package\s+main|import\s*\(|fmt\.Println)\b',
        'rust': r'\b(fn|let\s+mut|impl|pub\s+struct|use\s+std)\b',
        'typescript': r'\b(interface|type|namespace|readonly|abstract\s+class)\b',
    }
    
    for lang, pattern in patterns.items():
        if re.search(pattern, code):
            return lang
    return 'unknown'

def validate_and_fix_code(code: str, language: str) -> tuple[str, str]:
    if language == 'python':
        try:
            fixed_code = autopep8.fix_code(code)
            fixed_code = isort.SortImports(file_contents=fixed_code).output
            fixed_code = black.format_str(fixed_code, mode=black.FileMode())
            return fixed_code, ""
        except Exception as e:
            return code, f"Error in fixing Python code: {str(e)}"
    elif language == 'javascript':
        # Use a JS beautifier (placeholder)
        return code, ""
    elif language == 'java':
        # Use a Java formatter (placeholder)
        return code, ""
    elif language == 'c++':
        # Use a C++ formatter (placeholder)
        return code, ""
    else:
        return code, ""

def optimize_code(code: str) -> tuple[str, str]:
    language = detect_language(code)
    fixed_code, fix_error = validate_and_fix_code(code, language)
    
    if fix_error:
        return fixed_code, fix_error

    if language == 'python':
        try:
            tree = ast.parse(fixed_code)
            # Perform advanced Python-specific optimizations
            optimizer = PythonCodeOptimizer()
            optimized_tree = optimizer.visit(tree)
            optimized_code = ast.unparse(optimized_tree)
        except SyntaxError as e:
            return fixed_code, f"SyntaxError: {str(e)}"
    elif language == 'java':
        try:
            tree = javalang.parse.parse(fixed_code)
            # Perform Java-specific optimizations
            optimizer = JavaCodeOptimizer()
            optimized_code = optimizer.optimize(tree)
        except javalang.parser.JavaSyntaxError as e:
            return fixed_code, f"JavaSyntaxError: {str(e)}"
    elif language == 'c++':
        try:
            index = clang.cindex.Index.create()
            tu = index.parse('temp.cpp', args=['-std=c++14'], unsaved_files=[('temp.cpp', fixed_code)])
            # Perform C++-specific optimizations
            optimizer = CppCodeOptimizer()
            optimized_code = optimizer.optimize(tu)
        except Exception as e:
            return fixed_code, f"C++ Parsing Error: {str(e)}"
    else:
        optimized_code = fixed_code  # For unsupported languages, return the fixed code

    # Run language-specific linter
    lint_results = run_linter(optimized_code, language)
    
    return optimized_code, lint_results

def run_linter(code: str, language: str) -> str:
    if language == 'python':
        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 result.stdout
    elif language == 'javascript':
        # Run ESLint (placeholder)
        return "JavaScript linting not implemented"
    elif language == 'java':
        # Run CheckStyle (placeholder)
        return "Java linting not implemented"
    elif language == 'c++':
        # Run cppcheck (placeholder)
        return "C++ linting not implemented"
    else:
        return "Linting not available for the detected language"

def fetch_from_github(query: str) -> List[Dict[str, Any]]:
    headers = {"Authorization": f"token {st.secrets['GITHUB_TOKEN']}"}
    response = requests.get(f"https://api.github.com/search/code?q={query}", headers=headers)
    if response.status_code == 200:
        return response.json()['items'][:5]  # Return top 5 results
    return []

def analyze_code_quality(code: str) -> Dict[str, float]:
    inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=512, padding="max_length")
    
    with torch.no_grad():
        outputs = codebert_model(**inputs)
    
    cls_embedding = outputs.last_hidden_state[:, 0, :]
    predictions = code_improvement_model(cls_embedding)
    
    quality_scores = {
        "style": predictions[0][0].item(),
        "efficiency": predictions[0][1].item(),
        "security": predictions[0][2].item(),
        "maintainability": predictions[0][3].item(),
        "scalability": predictions[0][4].item(),
        "readability": predictions[0][5].item(),
        "testability": predictions[0][6].item(),
        "modularity": predictions[0][7].item()
    }
    
    # Calculate additional metrics
    language = detect_language(code)
    if language == 'python':
        complexity = radon_complexity.cc_visit(code)
        maintainability = radon_metrics.mi_visit(code, True)
        quality_scores["cyclomatic_complexity"] = complexity[0].complexity if complexity else 0
        quality_scores["maintainability_index"] = maintainability
    
    return quality_scores

def visualize_code_structure(code: str) -> plt.Figure:
    try:
        tree = ast.parse(code)
        graph = nx.DiGraph()
        
        def add_nodes_edges(node, parent=None):
            node_id = id(node)
            graph.add_node(node_id, label=f"{type(node).__name__}\n{ast.unparse(node)[:20]}")
            if parent:
                graph.add_edge(id(parent), node_id)
            for child in ast.iter_child_nodes(node):
                add_nodes_edges(child, node)
        
        add_nodes_edges(tree)
        
        plt.figure(figsize=(15, 10))
        pos = nx.spring_layout(graph, k=0.9, iterations=50)
        nx.draw(graph, pos, with_labels=True, node_color='lightblue', node_size=2000, font_size=8, font_weight='bold', arrows=True)
        labels = nx.get_node_attributes(graph, 'label')
        nx.draw_networkx_labels(graph, pos, labels, font_size=6)
        
        return plt
    except SyntaxError:
        return None

def suggest_improvements(code: str, quality_scores: Dict[str, float]) -> List[str]:
    suggestions = []
    thresholds = {
        "style": 0.7,
        "efficiency": 0.7,
        "security": 0.8,
        "maintainability": 0.7,
        "scalability": 0.7,
        "readability": 0.7,
        "testability": 0.7,
        "modularity": 0.7
    }
    
    for metric, threshold in thresholds.items():
        if quality_scores[metric] < threshold:
            suggestions.append(f"Consider improving code {metric} (current score: {quality_scores[metric]:.2f}).")
    
    if "cyclomatic_complexity" in quality_scores and quality_scores["cyclomatic_complexity"] > 10:
        suggestions.append(f"Consider breaking down complex functions to reduce cyclomatic complexity (current: {quality_scores['cyclomatic_complexity']}).")
    
    return suggestions

# New function for advanced code generation using GPT-2
def generate_advanced_code(prompt: str, language: str) -> str:
    input_text = f"Generate {language} code for: {prompt}\n\n"
    input_ids = gpt2_tokenizer.encode(input_text, return_tensors="pt")
    
    output = gpt2_model.generate(
        input_ids,
        max_length=1000,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        top_k=50,
        top_p=0.95,
        temperature=0.7
    )
    
    generated_code = gpt2_tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_code.split("\n\n", 1)[1]  # Remove the input prompt from the generated text

# New function for code similarity analysis
def analyze_code_similarity(code1: str, code2: str) -> float:
    tokens1 = tokenizer.tokenize(code1)
    tokens2 = tokenizer.tokenize(code2)
    
    # Use Jaccard similarity for token-based comparison
    set1 = set(tokens1)
    set2 = set(tokens2)
    similarity = len(set1.intersection(set2)) / len(set1.union(set2))
    
    return similarity

# New function for code performance estimation
def estimate_code_performance(code: str) -> Dict[str, Any]:
    language = detect_language(code)
    if language == 'python':
        # Use abstract syntax tree to estimate time complexity
        tree = ast.parse(code)
        analyzer = ComplexityAnalyzer()
        analyzer.visit(tree)
        return {
            "time_complexity": analyzer.time_complexity,
            "space_complexity": analyzer.space_complexity
        }
    else:
        return {"error": "Performance estimation not supported for this language"}

class ComplexityAnalyzer(ast.NodeVisitor):
    def __init__(self):
        self.time_complexity = "O(1)"
        self.space_complexity = "O(1)"
        self.loop_depth = 0

    def visit_For(self, node):
        self.loop_depth += 1
        self.generic_visit(node)
        self.loop_depth -= 1
        self.update_complexity()

    def visit_While(self, node):
        self.loop_depth += 1
        self.generic_visit(node)
        self.loop_depth -= 1
        self.update_complexity()

    def update_complexity(self):
        if self.loop_depth > 0:
            self.time_complexity = f"O(n^{self.loop_depth})"
            self.space_complexity = "O(n)"

# New function for code translation between programming languages
def translate_code(code: str, source_lang: str, target_lang: str) -> str:
    prompt = f"Translate the following {source_lang} code to {target_lang}:\n\n{code}\n\nTranslated {target_lang} code:"
    translated_code = generate_advanced_code(prompt, target_lang)
    return translated_code

# New function for generating unit tests
def generate_unit_tests(code: str, language: str) -> str:
    prompt = f"Generate unit tests for the following {language} code:\n\n{code}\n\nUnit tests:"
    unit_tests = generate_advanced_code(prompt, language)
    return unit_tests

# New function for code documentation generation
def generate_documentation(code: str, language: str) -> str:
    prompt = f"Generate comprehensive documentation for the following {language} code:\n\n{code}\n\nDocumentation:"
    documentation = generate_advanced_code(prompt, language)
    return documentation

# New function for advanced code refactoring suggestions
def suggest_refactoring(code: str, language: str) -> List[str]:
    quality_scores = analyze_code_quality(code)
    suggestions = suggest_improvements(code, quality_scores)
    
    # Add more specific refactoring suggestions based on code analysis
    tree = ast.parse(code)
    analyzer = RefactoringAnalyzer()
    analyzer.visit(tree)
    
    suggestions.extend(analyzer.suggestions)
    return suggestions

class RefactoringAnalyzer(ast.NodeVisitor):
    def __init__(self):
        self.suggestions = []
        self.function_lengths = {}

    def visit_FunctionDef(self, node):
        function_length = len(node.body)
        self.function_lengths[node.name] = function_length
        if function_length > 20:
            self.suggestions.append(f"Consider breaking down the function '{node.name}' into smaller, more manageable functions.")
        self.generic_visit(node)

    def visit_If(self, node):
        if isinstance(node.test, ast.Compare) and len(node.test.ops) > 2:
            self.suggestions.append("Consider simplifying complex conditional statements.")
        self.generic_visit(node)

# New function for code security analysis
def analyze_code_security(code: str, language: str) -> List[str]:
    vulnerabilities = []
    
    if language == 'python':
        tree = ast.parse(code)
        analyzer = SecurityAnalyzer()
        analyzer.visit(tree)
        vulnerabilities.extend(analyzer.vulnerabilities)
    
    # Add more language-specific security checks here
    
    return vulnerabilities

class SecurityAnalyzer(ast.NodeVisitor):
    def __init__(self):
        self.vulnerabilities = []

    def visit_Call(self, node):
        if isinstance(node.func, ast.Name):
            if node.func.id == 'eval':
                self.vulnerabilities.append("Potential security risk: Use of 'eval' function detected.")
            elif node.func.id == 'exec':
                self.vulnerabilities.append("Potential security risk: Use of 'exec' function detected.")
        self.generic_visit(node)

# New function for code optimization suggestions
def suggest_optimizations(code: str, language: str) -> List[str]:
    suggestions = []
    
    if language == 'python':
        tree = ast.parse(code)
        analyzer = OptimizationAnalyzer()
        analyzer.visit(tree)
        suggestions.extend(analyzer.suggestions)
    
    # Add more language-specific optimization suggestions here
    
    return suggestions

class OptimizationAnalyzer(ast.NodeVisitor):
    def __init__(self):
        self.suggestions = []
        self.loop_variables = set()

    def visit_For(self, node):
        if isinstance(node.iter, ast.Call) and isinstance(node.iter.func, ast.Name) and node.iter.func.id == 'range':
            self.suggestions.append("Consider using 'enumerate()' instead of 'range()' for index-based iteration.")
        self.generic_visit(node)

    def visit_ListComp(self, node):
        if isinstance(node.elt, ast.Call) and isinstance(node.elt.func, ast.Name) and node.elt.func.id == 'append':
            self.suggestions.append("Consider using a list comprehension instead of appending in a loop for better performance.")
        self.generic_visit(node)

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

st.markdown("""
<style>
.main-container {
    padding: 2rem;
    border-radius: 10px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    background-color: #f8f9fa;
}
.title {
    color: #2c3e50;
    font-size: 2.5rem;
    margin-bottom: 1rem;
}
.subtitle {
    color: #34495e;
    font-size: 1.2rem;
    margin-bottom: 2rem;
}
.output-container {
    margin-top: 2rem;
    padding: 1rem;
    border-radius: 5px;
    background-color: #ffffff;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
}
.code-block {
    margin-bottom: 1rem;
}
.metric-container {
    display: flex;
    justify-content: space-between;
    flex-wrap: wrap;
}
.metric-item {
    flex-basis: 48%;
    margin-bottom: 1rem;
}
</style>
""", unsafe_allow_html=True)

st.markdown('<div class="main-container">', unsafe_allow_html=True)
st.markdown('<h1 class="title">🚀 Advanced AI Code Assistant</h1>', unsafe_allow_html=True)
st.markdown('<p class="subtitle">Powered by Cutting-Edge AI & Multi-Domain Expertise</p>', unsafe_allow_html=True)

task = st.selectbox("Select a task", [
    "Generate Code", "Optimize Code", "Analyze Code Quality", 
    "Translate Code", "Generate Unit Tests", "Generate Documentation", 
    "Suggest Refactoring", "Analyze Code Security", "Suggest Optimizations"
])

language = st.selectbox("Select programming language", [
    "Python", "JavaScript", "Java", "C++", "Ruby", "Go", "Rust", "TypeScript"
])

prompt = st.text_area("Enter your code or prompt", height=200)

if st.button("Execute Task"):
    if prompt.strip() == "":
        st.error("Please enter a valid prompt or code snippet.")
    else:
        with st.spinner("Processing your request..."):
            if task == "Generate Code":
                result = generate_advanced_code(prompt, language.lower())
                st.code(result, language=language.lower())
            elif task == "Optimize Code":
                optimized_code, lint_results = optimize_code(prompt)
                st.code(optimized_code, language=language.lower())
                st.text(lint_results)
            elif task == "Analyze Code Quality":
                quality_scores = analyze_code_quality(prompt)
                st.json(quality_scores)
            elif task == "Translate Code":
                target_lang = st.selectbox("Select target language", [
                    lang for lang in ["Python", "JavaScript", "Java", "C++", "Ruby", "Go", "Rust", "TypeScript"] if lang != language
                ])
                translated_code = translate_code(prompt, language.lower(), target_lang.lower())
                st.code(translated_code, language=target_lang.lower())
            elif task == "Generate Unit Tests":
                unit_tests = generate_unit_tests(prompt, language.lower())
                st.code(unit_tests, language=language.lower())
            elif task == "Generate Documentation":
                documentation = generate_documentation(prompt, language.lower())
                st.markdown(documentation)
            elif task == "Suggest Refactoring":
                refactoring_suggestions = suggest_refactoring(prompt, language.lower())
                for suggestion in refactoring_suggestions:
                    st.info(suggestion)
            elif task == "Analyze Code Security":
                vulnerabilities = analyze_code_security(prompt, language.lower())
                if vulnerabilities:
                    for vuln in vulnerabilities:
                        st.warning(vuln)
                else:
                    st.success("No obvious security vulnerabilities detected.")
            elif task == "Suggest Optimizations":
                optimization_suggestions = suggest_optimizations(prompt, language.lower())
                for suggestion in optimization_suggestions:
                    st.info(suggestion)

            # Additional analysis for all tasks
            quality_scores = analyze_code_quality(prompt)
            performance_estimate = estimate_code_performance(prompt)
            
            col1, col2 = st.columns(2)
            with col1:
                st.subheader("Code Quality Metrics")
                for metric, score in quality_scores.items():
                    st.metric(metric.capitalize(), f"{score:.2f}")
            
            with col2:
                st.subheader("Performance Estimation")
                st.json(performance_estimate)
            
            visualization = visualize_code_structure(prompt)
            if visualization:
                st.subheader("Code Structure Visualization")
                st.pyplot(visualization)

st.markdown("""
<div style='text-align: center; margin-top: 2rem; color: #4a5568;'>
    Powered by Advanced AI & Multi-Domain Expertise
</div>
""", unsafe_allow_html=True)

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

# FastAPI setup for potential API endpoints
app = FastAPI()

class CodeRequest(BaseModel):
    code: str
    language: str
    task: str

@app.post("/analyze")
async def analyze_code(request: CodeRequest):
    if request.task == "quality":
        return analyze_code_quality(request.code)
    elif request.task == "security":
        return analyze_code_security(request.code, request.language)
    elif request.task == "optimize":
        optimized_code, _ = optimize_code(request.code)
        return {"optimized_code": optimized_code}
    else:
        return {"error": "Invalid task"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)