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

import torch
from transformers import GPTNeoForCausalLM, GPT2Tokenizer, BartForConditionalGeneration, BartTokenizer
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
    QPushButton, QTextEdit, QLabel, QProgressBar, QFileDialog, QMessageBox, QTabWidget, QSplitter, QLineEdit,
    QFrame
)
from PyQt6.QtCore import Qt, QRunnable, QThreadPool, QObject, pyqtSignal, pyqtSlot, QTimer, QEvent
from PyQt6.QtGui import QTextCursor, QColor
import aiosqlite
from bs4 import BeautifulSoup
import psutil
import ray
from sentence_transformers import SentenceTransformer
from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from accelerate import Accelerator
import aiofiles
from aiolimiter import AsyncLimiter
from sklearn.cluster import KMeans
import spacy
from rich.console import Console
from rich.logging import RichHandler
import bitsandbytes as bnb
from nltk import word_tokenize, sent_tokenize
from nltk.corpus import stopwords
from rake_nltk import Rake

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

# Load environment variables
GPT_NEO_MODEL = r'C:\solana_rust_bot\local_gpt_neo_model'
KNOWLEDGE_DB_PATH = 'solana_knowledge_base.db'
MEMORY_PATH = 'memory.json'
KB_JSON_PATH = r'C:\solana_rust_bot\solana_knowledge_base.json'

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

# Initialize Ray for distributed computing
ray.init(ignore_reinit_error=True)

# Initialize Accelerator
accelerator = Accelerator()

# Constants for Groq rate limiting and tokens
MAX_GROQ_CALLS_PER_MINUTE = 30
MAX_TOKENS_PER_MINUTE = 5000
MAX_RESPONSE_SIZE = 20000  # 20KB in bytes

# Initialize rate limiter
groq_limiter = AsyncLimiter(MAX_GROQ_CALLS_PER_MINUTE, 60)

# Load spaCy model
nlp = spacy.load("en_core_web_sm")

# Initialize RAKE for keyword extraction
rake = Rake()

@ray.remote(num_gpus=1 if torch.cuda.is_available() else 0)
def load_model():
    try:
        logger.info("Loading models...")
        gpt_model = GPTNeoForCausalLM.from_pretrained(GPT_NEO_MODEL)
        gpt_tokenizer = GPT2Tokenizer.from_pretrained(GPT_NEO_MODEL)

        bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
        bart_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")

        gpt_model, bart_model = accelerator.prepare(gpt_model, bart_model)
        
        # Quantize models
        gpt_model = quantize_model(gpt_model)
        bart_model = quantize_model(bart_model)

        logger.info("Models loaded successfully")
        return gpt_model, gpt_tokenizer, bart_model, bart_tokenizer
    except Exception as e:
        logger.error(f"Error loading models: {e}")
        return None

model_tokenizer_ref = load_model.remote()
sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')

async def load_memory():
    try:
        async with aiofiles.open(MEMORY_PATH, 'r') as f:
            return json.loads(await f.read())
    except FileNotFoundError:
        return []

async def save_memory(memory):
    async with aiofiles.open(MEMORY_PATH, 'w') as f:
        await f.write(json.dumps(memory, indent=4))

def quantize_model(model):
    for module in model.modules():
        if isinstance(module, torch.nn.Linear):
            module = bnb.nn.Linear8bitLt(module.in_features, module.out_features, bias=module.bias is not None)
    return model

def chunk_text(text, chunk_size=1000):
    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]

async def process_large_input(text):
    chunks = chunk_text(text)
    results = []
    total_tokens = 0
    
    for chunk in chunks:
        result = await process_chunk(chunk)
        results.append(result)
        
        # Count tokens and break if limit is reached
        total_tokens += len(gpt_tokenizer.encode(result))
        if total_tokens >= MAX_TOKENS_LIMIT:
            break
    
    return combine_results(results)

async def process_chunk(chunk):
    gpt_model, gpt_tokenizer, _, _ = await model_tokenizer_ref
    inputs = gpt_tokenizer.encode(chunk, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = gpt_model.generate(
            inputs,
            max_length=min(len(inputs[0]) + 500, 4000),
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            temperature=0.7,
        )
    
    result = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result.strip()

def combine_results(results):
    return "\n".join(results)

def cluster_results(embeddings, n_clusters=5):
    kmeans = KMeans(n_clusters=n_clusters)
    clusters = kmeans.fit_predict(embeddings)
    return clusters

def process_text(text):
    doc = nlp(text)
    entities = [(ent.text, ent.label_) for ent in doc.ents]
    return entities

def extract_keywords(text):
    rake.extract_keywords_from_text(text)
    return rake.get_ranked_phrases()[:5]  # Get top 5 keywords
class WorkerSignals(QObject):
    progress = pyqtSignal(int)
    log = pyqtSignal(str)
    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)
            self.signals.log.emit("Processing...")  # Emit log before starting
            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 SolanaKnowledgeDB:
    def __init__(self, db_path: str):
        self.db_path = db_path

    async def init_db(self):
        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()

    async def add_knowledge(self, content: str, category: str):
        content_hash = hashlib.sha256(content.encode()).hexdigest()
        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()

    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 = torch.cosine_similarity(torch.tensor(query_embedding), torch.tensor(content_embedding), dim=0)
                results.append((similarity.item(), {'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 load_from_json(self, json_path: str):
        try:
            async with aiofiles.open(json_path, 'r') as f:
                data = json.loads(await f.read())
            for item in data:
                await self.add_knowledge(item['content'], item.get('category', ''))
            logger.info("Knowledge base loaded from JSON successfully.")
        except Exception as e:
            logger.error(f"Failed to load knowledge base from JSON: {e}")
            raise  # Re-raise exception after logging

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_official_docs(self, url):
        await self.create_session()
        try:
            async with self.session.get(url, timeout=30) as response:
                if response.status == 200:
                    html = await response.text()
                    soup = BeautifulSoup(html, 'html.parser')
                    content = soup.find('main')
                    return content.get_text() if content else ""
                else:
                    return ""
        except asyncio.TimeoutError:
            return "Error: Request timed out"
        except Exception as e:
            return f"Error: {str(e)}"

class CodeAnalyzer:
    @staticmethod
    @ray.remote
    def analyze_code(code: str) -> Dict[str, Any]:
        analysis = {
            "num_functions": len(re.findall(r'\bdef\b', code)),
            "num_classes": len(re.findall(r'\bclass\b', code)),
            "num_imports": len(re.findall(r'\bimport\b', code)),
            "lines_of_code": len(code.splitlines()),
            "complexity": CodeAnalyzer.calculate_complexity(code),
            "potential_issues": CodeAnalyzer.identify_potential_issues(code),
        }
        return analysis

    @staticmethod
    def calculate_complexity(code: str) -> int:
        complexity = 0
        complexity += len(re.findall(r'\bif\b', code))
        complexity += len(re.findall(r'\bfor\b', code))
        complexity += len(re.findall(r'\bwhile\b', code))
        complexity += len(re.findall(r'\bexcept\b', code))
        return complexity

    @staticmethod
    def identify_potential_issues(code: str) -> List[str]:
        issues = []
        if 'print' in code:
            issues.append("Consider using logging instead of print statements.")
        if 'except:' in code:
            issues.append("Avoid bare except clauses.")
        if 'global ' in code:
            issues.append("Minimize use of global variables.")
        if 'import *' in code:
            issues.append("Avoid wildcard imports.")
        if 'assert' in code:
            issues.append("Use assertions judiciously, not for data validation in production.")
        return issues

class CodeEnhancer:
    @staticmethod
    @ray.remote
    def enhance_code(code: str, analysis: Dict[str, Any]) -> str:
        prompt = f"""Enhance the following Solana code:\n\n{code}\n\nCode analysis:\n{json.dumps(analysis, indent=2)}\n\nEnhance the code by providing a complete and detailed response without omissions."""
        return CodeEnhancer.generate_with_model.remote(prompt)

    @staticmethod
    @ray.remote
    def generate_with_model(prompt: str) -> str:
        gpt_model, gpt_tokenizer, _, _ = ray.get(model_tokenizer_ref)
        inputs = gpt_tokenizer.encode(prompt, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = gpt_model.generate(
                inputs,
                max_length=min(len(inputs[0]) + 500, 4000),  # Set a response limit
                num_return_sequences=1,
                no_repeat_ngram_size=2,
                temperature=0.7,
            )
        response = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response[len(prompt):].strip()

class App(QMainWindow):
    MAX_CHAT_MEMORY = 100  # Maximum number of chat messages to retain
    MAX_TOKENS_LIMIT = 200000  # Tokens threshold for cleaning up chat history

    def __init__(self):
        super().__init__()
        self.setWindowTitle('Legendary Solana Sniper Bot Analyzer and Enhancer (1.3B Model)')
        self.setGeometry(100, 100, 1200, 800)

        self.selected_folder = ''
        self.threadpool = QThreadPool()
        self.knowledge_db = SolanaKnowledgeDB(KNOWLEDGE_DB_PATH)
        self.web_crawler = WebCrawler()
        self.chat_memory = deque(maxlen=self.MAX_CHAT_MEMORY)  # Memory for chat context
        self.current_token_count = 0
        self.tokenizer = GPT2Tokenizer.from_pretrained(GPT_NEO_MODEL)

        self.init_ui()
        self.init_knowledge_db()

    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)

        worker_load_kb = AsyncWorker(self.knowledge_db.load_from_json, KB_JSON_PATH)
        worker_load_kb.signals.error.connect(self.handle_db_init_error)
        self.threadpool.start(worker_load_kb)

    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 init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

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

        # Analyzer tab
        analyzer_tab = QWidget()
        analyzer_layout = QVBoxLayout(analyzer_tab)
        self.tabs.addTab(analyzer_tab, "Analyzer")

        self.select_folder_btn = QPushButton('Select Folder')
        self.select_folder_btn.clicked.connect(self.select_folder)
        self.select_folder_btn.setStyleSheet("background-color: lightblue; color: black; font-weight: bold;")
        analyzer_layout.addWidget(self.select_folder_btn)

        self.selected_folder_label = QLabel('No folder selected')
        analyzer_layout.addWidget(self.selected_folder_label)

        self.start_btn = QPushButton('Start Analysis')
        self.start_btn.clicked.connect(self.start_analysis)
        self.start_btn.setEnabled(False)
        self.start_btn.setStyleSheet("background-color: limegreen; color: white; font-weight: bold;")
        analyzer_layout.addWidget(self.start_btn)

        self.progress_bar = QProgressBar()
        self.progress_bar.setStyleSheet("QProgressBar::chunk { background-color: limegreen; }")
        analyzer_layout.addWidget(self.progress_bar)

        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        analyzer_layout.addWidget(self.log_text)

        # Enhancer tab
        enhancer_tab = QWidget()
        enhancer_layout = QVBoxLayout(enhancer_tab)
        self.tabs.addTab(enhancer_tab, "Enhancer")

        splitter = QSplitter(Qt.Orientation.Horizontal)
        enhancer_layout.addWidget(splitter)

        self.original_code = QTextEdit()
        self.original_code.setPlaceholderText("Paste your original code here...")
        splitter.addWidget(self.original_code)

        self.enhanced_code = QTextEdit()
        self.enhanced_code.setReadOnly(True)
        self.enhanced_code.setPlaceholderText("Enhanced code will appear here...")
        splitter.addWidget(self.enhanced_code)

        enhance_btn = QPushButton('Enhance Code')
        enhance_btn.clicked.connect(self.enhance_code)
        enhance_btn.setStyleSheet("background-color: orange; color: white; font-weight: bold;")
        enhancer_layout.addWidget(enhance_btn)

        # 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_display_frame = QFrame()
        chat_display_frame.setLayout(QVBoxLayout())
        chat_display_frame.setStyleSheet("background-color: #f9f9f9; border: 1px solid #ccc; padding: 5px;")
        chat_display_frame.layout().addWidget(self.chat_display)
        chat_layout.addWidget(chat_display_frame)

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

        send_btn = QPushButton('Send')
        send_btn.clicked.connect(self.send_chat)
        send_btn.setStyleSheet("background-color: #4CAF50; color: white; font-weight: bold;")
        chat_layout.addWidget(send_btn)

        self.clear_chat_btn = QPushButton('Clear Chat')
        self.clear_chat_btn.clicked.connect(self.clear_chat)
        self.clear_chat_btn.setStyleSheet("background-color: #f44336; color: white; font-weight: bold;")
        chat_layout.addWidget(self.clear_chat_btn)

        # Web Crawler tab
        crawler_tab = QWidget()
        crawler_layout = QVBoxLayout(crawler_tab)
        self.tabs.addTab(crawler_tab, "Web Crawler")

        self.crawler_input = QLineEdit()
        self.crawler_input.setPlaceholderText("Enter search query or URL...")
        crawler_layout.addWidget(self.crawler_input)

        crawler_btn = QPushButton('Crawl')
        crawler_btn.clicked.connect(self.start_crawl)
        crawler_btn.setStyleSheet("background-color: #008CBA; color: white; font-weight: bold;")
        crawler_layout.addWidget(crawler_btn)

        self.crawler_results = QTextEdit()
        self.crawler_results.setReadOnly(True)
        crawler_layout.addWidget(self.crawler_results)

        # Export buttons
        export_layout = QHBoxLayout()
        layout.addLayout(export_layout)

        export_analysis_btn = QPushButton('Export Analysis')
        export_analysis_btn.clicked.connect(self.export_analysis)
        export_analysis_btn.setStyleSheet("background-color: #555555; color: white; font-weight: bold;")
        export_layout.addWidget(export_analysis_btn)

        export_chat_btn = QPushButton('Export Chat')
        export_chat_btn.clicked.connect(self.export_chat_history)
        export_chat_btn.setStyleSheet("background-color: #555555; color: white; font-weight: bold;")
        export_layout.addWidget(export_chat_btn)

        # Set up event filter for Ctrl+Enter in chat input
        self.chat_input.installEventFilter(self)

        # Timer for updating performance metrics
        self.performance_timer = QTimer(self)
        self.performance_timer.timeout.connect(self.update_performance_metrics)
        self.performance_timer.start(1000)  # Update every second

    def select_folder(self):
        folder = QFileDialog.getExistingDirectory(self, 'Select Folder')
        if folder:
            self.selected_folder = folder
            self.selected_folder_label.setText(f'Selected Folder: {folder}')
            self.start_btn.setEnabled(True)

    def start_analysis(self):
        self.start_btn.setEnabled(False)
        self.progress_bar.setValue(0)
        self.log_text.clear()

        worker = AsyncWorker(self.analyze_folder)
        worker.signals.progress.connect(self.update_progress)
        worker.signals.log.connect(self.update_log)
        worker.signals.finished.connect(self.analysis_finished)
        worker.signals.error.connect(self.handle_worker_error)

        self.threadpool.start(worker)
        self.update_log("Analysis started...")

    async def analyze_folder(self):
        files = [f for f in os.listdir(self.selected_folder) if f.endswith('.py') or f.endswith('.rs')]
        total_files = len(files)

        for i, file in enumerate(files):
            file_path = os.path.join(self.selected_folder, file)
            self.signals.log.emit(f"Analyzing {file}...")

            with open(file_path, 'r') as f:
                content = f.read()

            analysis = await CodeAnalyzer.analyze_code.remote(content)
            self.signals.log.emit(f"File: {file}")
            self.signals.log.emit(f"Analysis: {json.dumps(analysis, indent=2)}")

            enhanced_code = await CodeEnhancer.enhance_code.remote(content, analysis)
            self.signals.log.emit(f"Enhanced code for {file}")

            enhanced_file_path = os.path.join(self.selected_folder, f"enhanced_{file}")
            with open(enhanced_file_path, 'w') as f:
                f.write(enhanced_code)

            progress = int((i + 1) / total_files * 100)
            self.signals.progress.emit(progress)

    def enhance_code(self):
        original_code = self.original_code.toPlainText()
        if not original_code:
            QMessageBox.warning(self, "Warning", "Please enter some code to enhance.")
            return

        self.enhanced_code.clear()
        self.enhanced_code.setPlaceholderText("Enhancing code...")

        worker = AsyncWorker(self.perform_code_enhancement, original_code)
        worker.signals.result.connect(self.update_enhanced_code)
        worker.signals.error.connect(self.handle_enhancement_error)

        self.threadpool.start(worker)

    async def perform_code_enhancement(self, code):
        analysis = await CodeAnalyzer.analyze_code.remote(code)
        return await CodeEnhancer.enhance_code.remote(code, analysis)

    def update_enhanced_code(self, enhanced_code):
        self.enhanced_code.setPlainText(enhanced_code)

    def handle_enhancement_error(self, error_message):
        self.enhanced_code.setPlainText(f"Error during enhancement: {error_message}")

    def send_chat(self):
        message = self.chat_input.toPlainText().strip()
        if message:
            token_count = len(self.tokenizer.encode(message))
            self.current_token_count += token_count
            
            self.chat_display.append(f"User: {message}")
            self.chat_input.clear()
            self.chat_display.append("AI is thinking...")

            self.chat_memory.append({"role": "user", "content": message})

            if self.current_token_count > self.MAX_TOKENS_LIMIT:
                self.clean_chat_history()

            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:
            context_memory = "\n".join([f"{mem['role'].capitalize()}: {mem['content']}" for mem in self.chat_memory])
            context = f"You are an expert AI assistant. Please answer the following request:\n{context_memory}\nAssistant:"

            if any(keyword in message.lower() for keyword in ["code", "information", "example", "help", "find"]):
                logger.info("Automatically invoking web crawlers for additional context.")
                async with groq_limiter:
                    chunks = chunk_text(message)
                    additional_info = []
                    for chunk in chunks:
                        result = await self.web_crawler.crawl_sourcegraph(chunk)
                        additional_info.extend(result)
                
                if additional_info:
                    context += f"\nAdditional Information from Sourcegraph:\n{' '.join(additional_info)}"

            response = await process_large_input(context)
            self.chat_memory.append({"role": "assistant", "content": response})
            
            token_count = len(self.tokenizer.encode(response))
            self.current_token_count += token_count
            
            if self.current_token_count > self.MAX_TOKENS_LIMIT:
                self.clean_chat_history()
            
            return response
        except Exception as e:
            logger.error(f"Error in handle_ai_chat: {str(e)}")
            return f"An error occurred: {str(e)}"

    def update_chat_display(self, message):
        self.chat_display.append(f"AI: {message}")
        self.chat_display.moveCursor(QTextCursor.MoveOperation.End)

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

    def clean_chat_history(self):
        while self.chat_memory and self.current_token_count > self.MAX_TOKENS_LIMIT:
            oldest_message = self.chat_memory.popleft()
            self.current_token_count -= len(self.tokenizer.encode(oldest_message['content']))
        self.update_log("Old chat messages cleared to optimize performance.")
        self.update_chat_display("System: Chat history has been partially cleared to maintain performance.")

    def clear_chat(self):
        self.chat_display.clear()
        self.chat_memory.clear()
        self.current_token_count = 0
        self.update_log("Chat cleared.")
        self.update_chat_display("System: Chat has been cleared.")

    def start_crawl(self):
        query = self.crawler_input.text().strip()
        if not query:
            QMessageBox.warning(self, "Warning", "Please enter a search query or URL.")
            return

        self.crawler_results.clear()
        self.crawler_results.setPlaceholderText("Crawling...")

        worker = AsyncWorker(self.perform_crawl, query)
        worker.signals.result.connect(self.update_crawler_results)
        worker.signals.error.connect(self.handle_crawler_error)

        self.threadpool.start(worker)

    async def perform_crawl(self, query):
        results = []
        if query.startswith('http'):
            result = await self.web_crawler.crawl_official_docs(query)
            results.append(result)
        else:
            sourcegraph_results = await self.web_crawler.crawl_sourcegraph(query)
            results.extend(sourcegraph_results)

        for result in results:
            await self.knowledge_db.add_knowledge(result, 'web_crawl')

        return results

    def update_crawler_results(self, results):
        if isinstance(results, list):
            self.crawler_results.setPlainText("\n\n".join(results))
        else:
            self.crawler_results.setPlainText(results)

    def handle_crawler_error(self, error_message):
        self.crawler_results.setPlainText(f"Error during crawling: {error_message}")

    def export_analysis(self):
        file_path, _ = QFileDialog.getSaveFileName(self, "Save Analysis Results", "", "Text Files (*.txt);;All Files (*)")
        if file_path:
            with open(file_path, 'w') as f:
                f.write(self.log_text.toPlainText())
            QMessageBox.information(self, "Export Successful", f"Analysis results exported to {file_path}")

    def export_chat_history(self):
        file_path, _ = QFileDialog.getSaveFileName(self, "Save Chat History", "", "Text Files (*.txt);;All Files (*)")
        if file_path:
            with open(file_path, 'w') as f:
                f.write(self.chat_display.toPlainText())
            QMessageBox.information(self, "Export Successful", f"Chat history exported to {file_path}")

    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.ControlModifier):
            self.send_chat()
            return True
        return super().eventFilter(source, event)

    def update_performance_metrics(self):
        cpu_percent = psutil.cpu_percent()
        memory_percent = psutil.virtual_memory().percent
        gpu_utilization = self.get_gpu_utilization()
        
        metrics = f"CPU: {cpu_percent}% | RAM: {memory_percent}% | GPU: {gpu_utilization}%"
        self.statusBar().showMessage(metrics)

    def get_gpu_utilization(self):
        try:
            output = subprocess.check_output(['nvidia-smi', '--query-gpu=utilization.gpu', '--format=csv,noheader,nounits'])
            utilization = output.decode('utf-8').strip()
            return f"{utilization}%"
        except Exception as e:
            logger.error(f"Error getting GPU utilization: {e}")
            return "N/A"

    def update_progress(self, value):
        self.progress_bar.setValue(value)

    def update_log(self, message):
        if "Error" in message:
            self.log_text.setTextColor(QColor("red"))
        else:
            self.log_text.setTextColor(QColor("black"))
        self.log_text.append(message)
        self.log_text.moveCursor(QTextCursor.MoveOperation.End)

    def analysis_finished(self):
        self.update_log("Analysis completed.")
        QMessageBox.information(self, "Completed", "Analysis and enhancement completed.")
        self.start_btn.setEnabled(True)

    def handle_worker_error(self, error_message):
        self.update_log(f"Error: {error_message}")
        QMessageBox.critical(self, "Error", error_message)
        self.start_btn.setEnabled(True)

    def closeEvent(self, event):
        ray.shutdown()
        event.accept()

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

class CodeEnhanceRequest(BaseModel):
    code: str

@app.post("/enhance_code")
async def enhance_code_api(request: CodeEnhanceRequest, background_tasks: BackgroundTasks):
    analysis = await CodeAnalyzer.analyze_code.remote(request.code)
    enhanced_code = await CodeEnhancer.enhance_code.remote(request.code, analysis)
    background_tasks.add_task(SolanaKnowledgeDB(KNOWLEDGE_DB_PATH).add_knowledge, enhanced_code, "enhanced_code")
    return {"enhanced_code": enhanced_code}

# Main PyQt application setup
def main():
    qapp = QApplication(sys.argv)
    app = App()
    app.show()
    sys.exit(qapp.exec())

if __name__ == '__main__':
    main()
     please write my entire script every single word of it dont remove code just fix any issues to it runs very well and high performing with fixing the tokens for the local model please all of the code dont remove anything just adjust minor issues