import os
import sys
import asyncio
import aiohttp
import json
import logging
import hashlib
import re
import gc
from datetime import datetime
from typing import List, Dict, Any
from collections import deque

import torch
import torch.nn.functional as F
from transformers import GPTJForCausalLM, AutoTokenizer, AlbertModel, BartForConditionalGeneration
from keybert import KeyBERT
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
    QTextEdit, QLabel, QProgressBar, QFileDialog, QMessageBox, QTabWidget, QSplitter
)
from PyQt6.QtCore import Qt, QRunnable, QThreadPool, QObject, pyqtSignal, pyqtSlot, QTimer, QEvent
from PyQt6.QtGui import QTextCursor, QColor, QSyntaxHighlighter, QTextCharFormat
from PyQt6.QtCore import QRegularExpression

import aiosqlite
from bs4 import BeautifulSoup
import psutil
from sentence_transformers import SentenceTransformer
from accelerate import Accelerator
import aiofiles
from aiolimiter import AsyncLimiter
import spacy
from rich.console import Console
from rich.logging import RichHandler
import bitsandbytes as bnb
from rake_nltk import Rake
import nltk
from groq import Groq
import ray
from pygments import highlight
from pygments.lexers import PythonLexer, RustLexer
from pygments.formatters import HtmlFormatter

console = Console()
logging.basicConfig(
    level="INFO",
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(console=console)]
)
logger = logging.getLogger("rich")

from dotenv import load_dotenv
load_dotenv()

GPT_J_MODEL = r'C:\solana_rust_bot\gpt-j-6B'
ALBERT_MODEL = 'albert-base-v2'
BART_MODEL = 'facebook/bart-large'
KNOWLEDGE_DB_PATH = 'solana_knowledge_base.db'
MEMORY_PATH = 'memory.json'
GROQ_API_KEY = os.getenv('GROQ_API_KEY')

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")

accelerator = Accelerator()

MAX_GROQ_CALLS_PER_MINUTE = 30
groq_limiter = AsyncLimiter(MAX_GROQ_CALLS_PER_MINUTE, 60)

nlp = spacy.load("en_core_web_sm")

rake = Rake()

nltk.download('punkt', quiet=True)
nltk.download('stopwords', quiet=True)

sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')

groq_client = Groq(api_key=GROQ_API_KEY)

ray.init(ignore_reinit_error=True)

async def load_models():
    global gpt_j_model, gpt_j_tokenizer, albert_model, albert_tokenizer, bart_model, bart_tokenizer, keybert_model
    try:
        print("Loading models...")
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        # Load GPT-J 6B
        gpt_j_tokenizer = AutoTokenizer.from_pretrained(GPT_J_MODEL)
        gpt_j_model = GPTJForCausalLM.from_pretrained(GPT_J_MODEL, low_cpu_mem_usage=True).to(device)
        
        # Load ALBERT
        albert_tokenizer = AutoTokenizer.from_pretrained(ALBERT_MODEL)
        albert_model = AlbertModel.from_pretrained(ALBERT_MODEL).to(device)
        
        # Load BART
        bart_tokenizer = AutoTokenizer.from_pretrained(BART_MODEL)
        bart_model = BartForConditionalGeneration.from_pretrained(BART_MODEL).to(device)
        
        # Load KeyBERT
        keybert_model = KeyBERT()
        
        print("All models loaded successfully.")
        return True
    except Exception as e:
        print(f"Error loading models: {e}")
        return False

class SolanaKnowledgeDB:
    def __init__(self, db_path: str):
        self.db_path = db_path

    async def init_db(self):
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute('''
                    CREATE TABLE IF NOT EXISTS knowledge_base
                        (id TEXT PRIMARY KEY, content TEXT, category TEXT, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)
                ''')
                await db.commit()
        except Exception as e:
            logger.error(f"Failed to initialize database: {e}")

    async def add_knowledge(self, content: str, category: str):
        content_hash = hashlib.sha256(content.encode()).hexdigest()
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    'INSERT OR REPLACE INTO knowledge_base (id, content, category) VALUES (?, ?, ?)',
                    (content_hash, content, category)
                )
                await db.commit()
        except Exception as e:
            logger.error(f"Failed to add knowledge: {e}")

    async def search_knowledge(self, query: str) -> List[Dict[str, Any]]:
        query_embedding = sentence_transformer.encode(query)
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute('SELECT * FROM knowledge_base')
            rows = await cursor.fetchall()
            results = []
            for row in rows:
                content_embedding = sentence_transformer.encode(row[1])
                similarity = F.cosine_similarity(
                    torch.tensor(query_embedding).unsqueeze(0),
                    torch.tensor(content_embedding).unsqueeze(0)
                ).item()
                results.append((similarity, {
                    'id': row[0],
                    'content': row[1],
                    'category': row[2],
                    'timestamp': row[3]
                }))
            results.sort(key=lambda x: x[0], reverse=True)
            return [item[1] for item in results[:5]]

    async def close_connection(self):
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute('VACUUM')
                await db.commit()
            logger.info("Database connection closed and optimized.")
        except Exception as e:
            logger.error(f"Error closing database connection: {e}")

class WebCrawler:
    def __init__(self):
        self.session = None

    async def create_session(self):
        if self.session is None:
            self.session = aiohttp.ClientSession()

    async def close_session(self):
        if self.session:
            await self.session.close()
            self.session = None

    async def crawl_sourcegraph(self, query):
        await self.create_session()
        url = f"https://sourcegraph.com/search?q={query}&patternType=literal"
        try:
            async with self.session.get(url, timeout=30) as response:
                if response.status == 200:
                    html = await response.text()
                    soup = BeautifulSoup(html, 'html.parser')
                    results = soup.find_all('div', class_='result-container')
                    return [result.get_text() for result in results]
                else:
                    return []
        except asyncio.TimeoutError:
            return ["Error: Request timed out"]
        except Exception as e:
            return [f"Error: {str(e)}"]

    async def crawl_solana_docs(self, topic):
        await self.create_session()
        url = f"https://docs.solana.com/search?q={topic}"
        try:
            async with self.session.get(url, timeout=30) as response:
                if response.status == 200:
                    html = await response.text()
                    soup = BeautifulSoup(html, 'html.parser')
                    results = soup.find_all('div', class_='search-result-item')
                    return [result.get_text() for result in results]
                else:
                    return []
        except asyncio.TimeoutError:
            return ["Error: Request timed out"]
        except Exception as e:
            return [f"Error: {str(e)}"]

class KnowledgeUpdater:
    def __init__(self, knowledge_db: SolanaKnowledgeDB, web_crawler: WebCrawler):
        self.knowledge_db = knowledge_db
        self.web_crawler = web_crawler
        self.topics = [
            "Solana API", "Solana trading", "Solana RPC", "Solana websockets",
            "Helius API", "Jupiter V6", "Raydium DEX", "SPL tokens",
            "Solana Rust programming", "Solana transaction structure",
            "Solana program development", "Solana account model",
            "Solana staking", "Solana wallet integration", "Solana NFTs",
            "Solana DeFi protocols", "Solana security best practices",
            "Solana performance optimization", "Solana cluster management",
            "Solana ecosystem tools and frameworks"
        ]

    async def update_knowledge_base(self):
        for topic in self.topics:
            logger.info(f"Updating knowledge base for topic: {topic}")
            groq_info = await self.get_groq_information(topic)
            sourcegraph_info = await self.web_crawler.crawl_sourcegraph(topic)
            solana_docs_info = await self.web_crawler.crawl_solana_docs(topic)

            combined_info = (
                f"GROQ Info: {groq_info}\n\n"
                f"Sourcegraph Info: {' '.join(sourcegraph_info)}\n\n"
                f"Solana Docs Info: {' '.join(solana_docs_info)}"
            )
            await self.knowledge_db.add_knowledge(combined_info, topic)

            logger.info(f"Knowledge base updated for topic: {topic}")

    async def get_groq_information(self, topic):
        prompt = (
            f"Provide detailed, technical information about {topic} in the context of Solana blockchain development and trading. "
            f"Include code examples and best practices where applicable."
        )
        try:
            async with groq_limiter:
                response = await groq_client.chat.completions.create(
                    model="mixtral-8x7b-32768",
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=2000
                )
            return response.choices[0].message.content
        except Exception as e:
            logger.error(f"Error getting information from GROQ: {e}")
            return ""

class WorkerSignals(QObject):
    finished = pyqtSignal()
    result = pyqtSignal(object)
    error = pyqtSignal(str)

class AsyncWorker(QRunnable):
    def __init__(self, fn, *args, **kwargs):
        super().__init__()
        self.fn = fn
        self.args = args
        self.kwargs = kwargs
        self.signals = WorkerSignals()

    @pyqtSlot()
    def run(self):
        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            result = loop.run_until_complete(self.fn(*self.args, **self.kwargs))
            self.signals.result.emit(result)
        except Exception as e:
            self.signals.error.emit(str(e))
        finally:
            self.signals.finished.emit()

class CodeHighlighter(QSyntaxHighlighter):
    def __init__(self, parent, language):
        super().__init__(parent)
        self.language = language
        self.highlighting_rules = []

        keyword_format = QTextCharFormat()
        keyword_format.setForeground(Qt.GlobalColor.darkBlue)
        keyword_format.setFontWeight(700)

        keyword_patterns = [
            "\\bconst\\b", "\\blet\\b", "\\bfn\\b", "\\bpub\\b", "\\buse\\b", "\\bmod\\b",
            "\\bstruct\\b", "\\benum\\b", "\\bself\\b", "\\bmut\\b", "\\bimpl\\b", "\\btrait\\b"
        ]

        for pattern in keyword_patterns:
            self.highlighting_rules.append((QRegularExpression(pattern), keyword_format))

        class_format = QTextCharFormat()
        class_format.setFontWeight(700)
        class_format.setForeground(Qt.GlobalColor.darkMagenta)
        self.highlighting_rules.append((QRegularExpression("\\bQ[A-Za-z]+\\b"), class_format))
        self.highlighting_rules.append((QRegularExpression("\\b[A-Z][A-Za-z]*\\b"), class_format))

        quotation_format = QTextCharFormat()
        quotation_format.setForeground(Qt.GlobalColor.darkGreen)
        self.highlighting_rules.append((QRegularExpression("\".*\""), quotation_format))

        function_format = QTextCharFormat()
        function_format.setFontItalic(True)
        function_format.setForeground(Qt.GlobalColor.blue)
        self.highlighting_rules.append((QRegularExpression("\\b[A-Za-z0-9_]+(?=\\()"), function_format))

        comment_format = QTextCharFormat()
        comment_format.setForeground(Qt.GlobalColor.darkGreen)
        self.highlighting_rules.append((QRegularExpression("//[^\n]*"), comment_format))

    def highlightBlock(self, text):
        for pattern, format in self.highlighting_rules:
            expression = QRegularExpression(pattern)
            it = expression.globalMatch(text)
            while it.hasNext():
                match = it.next()
                self.setFormat(match.capturedStart(), match.capturedLength(), format)

class App(QMainWindow):
    MAX_CHAT_MEMORY = 100
    MAX_TOKENS_LIMIT = 4000

    def __init__(self):
        super().__init__()
        self.setWindowTitle('Ultimate Solana AI Assistant V3 - GPT-J 6B Edition')
        self.setGeometry(100, 100, 1200, 800)

        self.threadpool = QThreadPool()
        self.knowledge_db = SolanaKnowledgeDB(KNOWLEDGE_DB_PATH)
        self.web_crawler = WebCrawler()
        self.knowledge_updater = KnowledgeUpdater(self.knowledge_db, self.web_crawler)
        self.chat_memory = self.load_chat_memory()
        self.current_token_count = 0

        self.model_loaded = False
        self.init_ui()
        self.init_knowledge_db()
        self.start_model_loading()
        self.start_knowledge_update_timer()

        app = QApplication.instance()
        app = QApplication.instance()
        app.aboutToQuit.connect(self.graceful_shutdown)

    def load_chat_memory(self):
        try:
            with open(MEMORY_PATH, 'r') as f:
                return deque(json.load(f), maxlen=self.MAX_CHAT_MEMORY)
        except FileNotFoundError:
            return deque(maxlen=self.MAX_CHAT_MEMORY)

    def save_chat_memory(self):
        with open(MEMORY_PATH, 'w') as f:
            json.dump(list(self.chat_memory), f)

    def init_knowledge_db(self):
        worker = AsyncWorker(self.knowledge_db.init_db)
        worker.signals.error.connect(self.handle_db_init_error)
        self.threadpool.start(worker)

    def handle_db_init_error(self, error_message):
        QMessageBox.critical(
            self,
            "Database Error",
            f"Failed to initialize the knowledge database: {error_message}"
        )
        logger.error(f"Database initialization error: {error_message}")

    def start_model_loading(self):
        worker = AsyncWorker(load_models)
        worker.signals.result.connect(self.on_models_loaded)
        worker.signals.error.connect(self.handle_model_load_error)
        self.threadpool.start(worker)

    def on_models_loaded(self, success):
        if success:
            self.model_loaded = True
            logger.info("All models loaded successfully.")
        else:
            QMessageBox.critical(
                self,
                "Model Load Error",
                "Failed to load one or more AI models."
            )

    def handle_model_load_error(self, error_message):
        QMessageBox.critical(
            self,
            "Model Load Error",
            f"An error occurred while loading the models: {error_message}"
        )
        logger.error(f"Model loading error: {error_message}")

    def start_knowledge_update_timer(self):
        self.knowledge_update_timer = QTimer(self)
        self.knowledge_update_timer.timeout.connect(self.update_knowledge_base)
        self.knowledge_update_timer.start(3600000)  # Run every hour

    def update_knowledge_base(self):
        worker = AsyncWorker(self.knowledge_updater.update_knowledge_base)
        self.threadpool.start(worker)

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        self.tabs = QTabWidget()
        main_layout.addWidget(self.tabs)

        # AI Chat Tab
        chat_tab = QWidget()
        chat_layout = QVBoxLayout(chat_tab)
        self.tabs.addTab(chat_tab, "AI Chat")

        self.chat_display = QTextEdit()
        self.chat_display.setReadOnly(True)
        chat_layout.addWidget(self.chat_display)

        input_layout = QHBoxLayout()
        self.chat_input = QTextEdit()
        self.chat_input.setFixedHeight(50)
        self.chat_input.setPlaceholderText("Type your message here...")
        self.chat_input.installEventFilter(self)
        input_layout.addWidget(self.chat_input)

        send_btn = QPushButton('Send')
        send_btn.clicked.connect(self.send_chat)
        input_layout.addWidget(send_btn)

        chat_layout.addLayout(input_layout)

        self.clear_chat_btn = QPushButton('Clear Chat')
        self.clear_chat_btn.clicked.connect(self.clear_chat)
        chat_layout.addWidget(self.clear_chat_btn)

        # Code Editor Tab
        code_tab = QWidget()
        code_layout = QVBoxLayout(code_tab)
        self.tabs.addTab(code_tab, "Code Editor")

        self.code_editor = QTextEdit()
        self.code_editor.setPlaceholderText("Enter your Solana code here...")
        self.code_highlighter = CodeHighlighter(self.code_editor.document(), "rust")
        code_layout.addWidget(self.code_editor)

        analyze_btn = QPushButton('Analyze Code')
        analyze_btn.clicked.connect(self.analyze_code)
        code_layout.addWidget(analyze_btn)

        self.code_analysis_display = QTextEdit()
        self.code_analysis_display.setReadOnly(True)
        code_layout.addWidget(self.code_analysis_display)

        # Knowledge Base Tab
        kb_tab = QWidget()
        kb_layout = QVBoxLayout(kb_tab)
        self.tabs.addTab(kb_tab, "Knowledge Base")

        self.kb_search_input = QTextEdit()
        self.kb_search_input.setFixedHeight(50)
        self.kb_search_input.setPlaceholderText("Search the knowledge base...")
        kb_layout.addWidget(self.kb_search_input)

        kb_search_btn = QPushButton('Search')
        kb_search_btn.clicked.connect(self.search_knowledge_base)
        kb_layout.addWidget(kb_search_btn)

        self.kb_display = QTextEdit()
        self.kb_display.setReadOnly(True)
        kb_layout.addWidget(self.kb_display)

    def eventFilter(self, source, event):
        if (source is self.chat_input and
            event.type() == QEvent.Type.KeyPress and
            event.key() == Qt.Key.Key_Return and
            event.modifiers() == Qt.KeyboardModifier.NoModifier):
            self.send_chat()
            return True
        return super().eventFilter(source, event)

    def send_chat(self):
        if not self.model_loaded:
            QMessageBox.warning(self, "Warning", "AI models are still loading. Please wait.")
            return

        message = self.chat_input.toPlainText().strip()
        if message:
            if len(gpt_j_tokenizer.encode(message)) > 150:
                message = ' '.join(message.split()[:150])

            self.chat_display.append(f"User: {message}")
            self.chat_input.clear()
            self.chat_display.append("AI: Processing...")
            QApplication.processEvents()

            worker = AsyncWorker(self.handle_ai_chat, message)
            worker.signals.result.connect(self.update_chat_display)
            worker.signals.error.connect(self.handle_chat_error)

            self.threadpool.start(worker)

    async def handle_ai_chat(self, message):
        try:
            self.chat_memory.append({"role": "user", "content": message})
            self.save_chat_memory()

            knowledge_context = await self.get_relevant_knowledge(message)

            context = (
                "You are an expert AI assistant specializing in Solana development and trading bots. "
                "Provide detailed, comprehensive responses with code examples where applicable. "
                f"Knowledge context: {knowledge_context}\n\n"
            )

            context += "\n".join([
                f"{msg['role'].capitalize()}: {msg['content']}" for msg in list(self.chat_memory)[-5:]
            ])

            context += f"\nUser: {message}\nAI:"

            if any(keyword in message.lower() for keyword in ["code", "script", "function", "class"]):
                sourcegraph_results = await self.web_crawler.crawl_sourcegraph(message)
                if sourcegraph_results:
                    context += f"\n\nRelevant code from Sourcegraph:\n" + "\n".join(sourcegraph_results[:3])

            response = await self.generate_enhanced_response(context)

            self.chat_memory.append({"role": "assistant", "content": response})
            self.save_chat_memory()

            await self.knowledge_db.add_knowledge(response, "ai_generated")

            return response
        except Exception as e:
            logger.error(f"Error in handle_ai_chat: {str(e)}")
            return f"An error occurred: {str(e)}"

    async def generate_enhanced_response(self, prompt):
        try:
            # Extract keywords using KeyBERT
            keywords = keybert_model.extract_keywords(prompt, keyphrase_ngram_range=(1, 2), stop_words='english', top_n=5)
            keyword_str = ", ".join([k for k, _ in keywords])

            # Use ALBERT for understanding
            albert_inputs = albert_tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
            albert_outputs = albert_model(**albert_inputs)
            albert_embedding = albert_outputs.last_hidden_state.mean(dim=1)

            # Use BART for initial response generation
            bart_inputs = bart_tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True).to(device)
            bart_outputs = bart_model.generate(**bart_inputs, max_length=150)
            bart_response = bart_tokenizer.decode(bart_outputs[0], skip_special_tokens=True)

            # Combine BART response with GPT-J
            combined_prompt = f"Keywords: {keyword_str}\nContext: {bart_response}\nUser query: {prompt}\nDetailed response:"
            gpt_j_response = await self.generate_with_gpt_j(combined_prompt)

            if len(gpt_j_response) < 200 or "I don't have enough information" in gpt_j_response:
                logger.info("GPT-J response insufficient. Using GROQ.")
                async with groq_limiter:
                    groq_response = await groq_client.chat.completions.create(
                        model="mixtral-8x7b-32768",
                        messages=[{"role": "user", "content": combined_prompt}],
                        max_tokens=2000
                    )
                response = groq_response.choices[0].message.content
            else:
                response = gpt_j_response

            # Enhance response with additional context and code examples if needed
            enhanced_response = self.enhance_response_with_context(response, prompt)

            return enhanced_response
        except Exception as e:
            logger.error(f"Error generating enhanced response: {e}")
            return "I'm sorry, I encountered an error while processing your request."

    async def generate_with_gpt_j(self, prompt):
        inputs = gpt_j_tokenizer.encode(prompt, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = gpt_j_model.generate(
                inputs,
                max_length=min(len(inputs[0]) + 500, self.MAX_TOKENS_LIMIT),
                num_return_sequences=1,
                no_repeat_ngram_size=2,
                temperature=0.7,
            )
        return gpt_j_tokenizer.decode(outputs[0], skip_special_tokens=True)

    def enhance_response_with_context(self, response, original_prompt):
        # Add more context, examples, or clarifications based on the original prompt and response
        if "code example" in original_prompt.lower() and "```" not in response:
            response += "\n\nHere's a code example to illustrate:\n```rust\n// Add relevant Rust code here\n```"
        
        if "explain" in original_prompt.lower() and len(response.split()) < 100:
            response += "\n\nTo elaborate further: [Add more detailed explanation]"
        
        if "compare" in original_prompt.lower() and "versus" not in response.lower():
            response += "\n\nTo compare more explicitly: [Add comparison details]"
        
        return response

    async def get_relevant_knowledge(self, query):
        results = await self.knowledge_db.search_knowledge(query)
        return "\n".join([result['content'] for result in results])

    def update_chat_display(self, message):
        self.chat_display.append(f"AI: {message}")
        self.chat_display.verticalScrollBar().setValue(
            self.chat_display.verticalScrollBar().maximum()
        )
        QApplication.processEvents()

    def handle_chat_error(self, error_message):
        self.chat_display.append(f"Error: {error_message}")
        QApplication.processEvents()

    def clear_chat(self):
        self.chat_display.clear()
        self.chat_memory.clear()
        self.current_token_count = 0
        self.chat_display.append("Chat cleared. Ready for new conversation.")
        self.save_chat_memory()
        QApplication.processEvents()

    def analyze_code(self):
        code = self.code_editor.toPlainText()
        if not code:
            QMessageBox.warning(self, "Warning", "Please enter some code to analyze.")
            return

        worker = AsyncWorker(self.perform_code_analysis, code)
        worker.signals.result.connect(self.update_code_analysis)
        worker.signals.error.connect(self.handle_analysis_error)

        self.threadpool.start(worker)

    async def perform_code_analysis(self, code):
        analysis = {
            "lines_of_code": len(code.split("\n")),
            "function_count": len(re.findall(r'\bfn\b', code)),
            "struct_count": len(re.findall(r'\bstruct\b', code)),
            "potential_issues": []
        }

        if "unwrap()" in code:
            analysis["potential_issues"].append("Use of unwrap() detected. Consider proper error handling.")

        if not re.search(r'#\[derive\(Debug\)\]', code):
            analysis["potential_issues"].append("Consider adding #[derive(Debug)] to structs for better debugging.")

        # Use ALBERT for code understanding
        albert_inputs = albert_tokenizer(code, return_tensors="pt", max_length=512, truncation=True).to(device)
        albert_outputs = albert_model(**albert_inputs)
        code_embedding = albert_outputs.last_hidden_state.mean(dim=1)

        # Use BART for initial code analysis
        bart_inputs = bart_tokenizer("Analyze this Solana Rust code: " + code, return_tensors="pt", max_length=1024, truncation=True).to(device)
        bart_outputs = bart_model.generate(**bart_inputs, max_length=200)
        bart_analysis = bart_tokenizer.decode(bart_outputs[0], skip_special_tokens=True)

        # Combine BART analysis with GPT-J for detailed insights
        combined_prompt = f"Code analysis: {bart_analysis}\n\nProvide a detailed analysis of this Solana Rust code, including best practices, potential improvements, and security considerations:\n\n{code}"
        gpt_j_analysis = await self.generate_with_gpt_j(combined_prompt)

        analysis["detailed_analysis"] = gpt_j_analysis

        return analysis

    def update_code_analysis(self, analysis):
        self.code_analysis_display.clear()
        self.code_analysis_display.append(f"Lines of code: {analysis['lines_of_code']}")
        self.code_analysis_display.append(f"Function count: {analysis['function_count']}")
        self.code_analysis_display.append(f"Struct count: {analysis['struct_count']}")

        if analysis['potential_issues']:
            self.code_analysis_display.append("\nPotential issues:")
            for issue in analysis['potential_issues']:
                self.code_analysis_display.append(f"- {issue}")

        self.code_analysis_display.append(f"\nDetailed Analysis:\n{analysis['detailed_analysis']}")

    def handle_analysis_error(self, error_message):
        self.code_analysis_display.append(f"Error during analysis: {error_message}")

    def search_knowledge_base(self):
        query = self.kb_search_input.toPlainText().strip()
        if not query:
            QMessageBox.warning(self, "Warning", "Please enter a search query.")
            return

        worker = AsyncWorker(self.perform_kb_search, query)
        worker.signals.result.connect(self.update_kb_display)
        worker.signals.error.connect(self.handle_kb_search_error)

        self.threadpool.start(worker)

    async def perform_kb_search(self, query):
        results = await self.knowledge_db.search_knowledge(query)
        
        # Enhance search results with ALBERT and BART
        enhanced_results = []
        for result in results:
            albert_inputs = albert_tokenizer(result['content'], return_tensors="pt", max_length=512, truncation=True).to(device)
            albert_outputs = albert_model(**albert_inputs)
            result_embedding = albert_outputs.last_hidden_state.mean(dim=1)

            bart_inputs = bart_tokenizer("Summarize this: " + result['content'], return_tensors="pt", max_length=1024, truncation=True).to(device)
            bart_outputs = bart_model.generate(**bart_inputs, max_length=100)
            summary = bart_tokenizer.decode(bart_outputs[0], skip_special_tokens=True)

            enhanced_results.append({
                **result,
                'summary': summary
            })

        return enhanced_results

    def update_kb_display(self, results):
        self.kb_display.clear()
        for result in results:
            self.kb_display.append(f"Category: {result['category']}")
            self.kb_display.append(f"Summary: {result['summary']}")
            self.kb_display.append(f"Content: {result['content']}")
            self.kb_display.append(f"Timestamp: {result['timestamp']}")
            self.kb_display.append("\n" + "-"*50 + "\n")

    def handle_kb_search_error(self, error_message):
        self.kb_display.append(f"Error during search: {error_message}")

    def graceful_shutdown(self):
        logger.info("Initiating graceful shutdown...")
        self.knowledge_update_timer.stop()
        self.threadpool.waitForDone()

        asyncio.run(self.web_crawler.close_session())
        asyncio.run(self.final_database_operations())

        self.save_chat_memory()
        logger.info("Graceful shutdown completed.")

    async def final_database_operations(self):
        try:
            await self.knowledge_db.add_knowledge("Shutdown occurred", "system_log")
            await self.knowledge_db.close_connection()
        except Exception as e:
            logger.error(f"Error during final database operations: {e}")

def main():
    app = QApplication(sys.argv)
    ex = App()
    ex.show()
    sys.exit(app.exec())

if __name__ == '__main__':
    main()   whats wrong with this script??