import sys
import os
import asyncio
import json
import subprocess
import ast
import re
import time
from datetime import datetime
from dotenv import load_dotenv
from PyQt6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QPushButton, QFileDialog,
    QTextEdit, QProgressBar, QMessageBox, QLabel, QTabWidget, QHBoxLayout, QLineEdit
)import os
import shutil
import tempfile   ONLY USEING GROQ API 8X7 MIXTRAL AND LOCAL ELEUTHER 2.7B WITH ENHANCED bitsandbytes 
import subprocess
import json
import time
from typing import List, Dict   import bitsandbytes as bnb     pip install transformers torch accelerate bitsandbytes

   model = bnb.optim.GlobalOptimManager.get_instance().quantize_model(model, quantization_type="bnb_8bit")from difflib import SequenceMatcher
   from transformers import AutoTokenizer, AutoModelForCausalLM
   import torch

   tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
   model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model.to(device)
from PyQt6.QtCore import Qt, QThread, pyqtSignal
import aiohttp
from groq import Groq
from bs4 import BeautifulSoup
import toml
from aiolimiter import AsyncLimiter
import torch
from transformers import GPTNeoForCausalLM, GPT2Tokenizer
import nltk
from nltk.tokenize import word_tokenize, sent_tokenize
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import spacy
from transformers import BertModel, BertTokenizer
import numpy as np
from functools import lru_cache

# Load environment variables from .env file
load_dotenv()

# Retrieve environment variables
GROQ_API_KEY = os.getenv('GROQ_API_KEY')
PYTHONPATH = os.getenv('PYTHONPATH')
GROQ_MODEL = os.getenv('GROQ_MODEL')
GPT_NEO_MODEL_NAME = os.getenv('GPT_NEO_MODEL', 'EleutherAI/gpt-neo-2.7B')
MODEL_PATH = os.getenv('MODEL_PATH')
KNOWLEDGE_PATH = os.getenv('KNOWLEDGE_PATH', 'solana_knowledge_base.json')
MEMORY_PATH = os.getenv('MEMORY_PATH', 'memory.json')

# Initialize the Groq client
groq_client = Groq(api_key=GROQ_API_KEY)
groq_model = GROQ_MODEL

# Initialize Rate Limiters
tokens_per_minute = 5000
calls_per_minute = 30
token_limiter = AsyncLimiter(tokens_per_minute, 60)
call_limiter = AsyncLimiter(calls_per_minute, 60)

# Download necessary NLTK data
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('wordnet')

# Load Spacy model
nlp = spacy.load('en_core_web_sm')

# Load configuration
if not os.path.exists(MEMORY_PATH):
    with open(MEMORY_PATH, 'w', encoding='utf-8') as f:
        json.dump([], f, indent=2)
    print("Memory file created.")
else:
    with open(MEMORY_PATH, 'r', encoding='utf-8') as f:
        memory = json.load(f)
if not os.path.exists('config.toml'):
    print("config.toml not found. Please create one with necessary configurations.")
    sys.exit(1)

with open('config.toml', 'r', encoding='utf-8') as f:
    config = toml.load(f)

# Load Solana knowledge base
def load_knowledge_base():
    if os.path.exists(KNOWLEDGE_PATH):
        with open(KNOWLEDGE_PATH, 'r', encoding='utf-8') as f:
            return json.load(f)
    else:
        return {}

knowledge_base = load_knowledge_base()

# Load or initialize memory for chat logs
def load_memory():
    if os.path.exists(MEMORY_PATH):
        with open(MEMORY_PATH, 'r', encoding='utf-8') as f:
            return json.load(f)
    else:
        return []

memory = load_memory()
def save_memory(memory):
    with open(MEMORY_PATH, 'w', encoding='utf-8') as f:
        json.dump(memory, f, indent=2)


memory = load_memory()

# Load local GPT Neo model
print("Loading local GPT Neo 2.7B model... This may take a while.")
tokenizer = GPT2Tokenizer.from_pretrained(GPT_NEO_MODEL_NAME)
model = GPTNeoForCausalLM.from_pretrained(GPT_NEO_MODEL_NAME)
model.eval()
if torch.cuda.is_available():
    model.to('cuda')
print("Local GPT Neo model loaded successfully.")

# Define the LegendaryNaturalLanguageProcessor class
class LegendaryNaturalLanguageProcessor:
    def __init__(self):
        self.stop_words = set(stopwords.words('english'))
        self.lemmatizer = WordNetLemmatizer()
        self.vectorizer = TfidfVectorizer()

        # Initialize BERT model
        self.bert_model = BertModel.from_pretrained('bert-base-uncased')
        self.bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

        # Initialize GPT-Neo model
        self.gpt_neo_model_path = os.getenv("MODEL_PATH")
        self.gpt_neo_model = GPTNeoForCausalLM.from_pretrained(self.gpt_neo_model_path)
        self.gpt_neo_tokenizer = GPT2Tokenizer.from_pretrained(self.gpt_neo_model_path)
        self.gpt_neo_tokenizer.pad_token = self.gpt_neo_tokenizer.eos_token
        self.gpt_neo_model.config.pad_token_id = self.gpt_neo_model.config.eos_token_id

        # Initialize Groq client
        self.groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
        self.groq_model = "mixtral_8x7"  # Assuming this is the model name

    def preprocess_text(self, text):
        doc = nlp(text.lower())
        return ' '.join([token.lemma_ for token in doc if token.is_alpha and token.lemma_ not in self.stop_words])

    def extract_keywords(self, text, num_keywords=5):
        preprocessed_text = self.preprocess_text(text)
        tfidf_matrix = self.vectorizer.fit_transform([preprocessed_text])
        feature_names = self.vectorizer.get_feature_names_out()
        sorted_items = sorted(zip(tfidf_matrix.tocsc().data, feature_names), reverse=True)
        return [item[1] for item in sorted_items[:num_keywords]]

    @lru_cache(maxsize=1000)
    def get_sentence_embedding(self, text):
        inputs = self.bert_tokenizer(text, return_tensors='pt', truncation=True, padding=True)
        outputs = self.bert_model(**inputs)
        return outputs.last_hidden_state.mean(dim=1).squeeze().detach().numpy()

    def calculate_similarity(self, text1, text2):
        embedding1 = self.get_sentence_embedding(text1)
        embedding2 = self.get_sentence_embedding(text2)
        return cosine_similarity([embedding1], [embedding2])[0][0]

    def extract_entities(self, text):
        doc = nlp(text)
        entities = {ent.label_: ent.text for ent in doc.ents}
        return entities

    def summarize_text(self, text, num_sentences=3):
        sentences = sent_tokenize(text)
        embeddings = np.array([self.get_sentence_embedding(sent) for sent in sentences])
        sentence_scores = np.sum(embeddings, axis=1)
        top_sentence_indices = sentence_scores.argsort()[-num_sentences:][::-1]
        summary = ' '.join([sentences[i] for i in sorted(top_sentence_indices)])
        return summary

    def analyze_sentiment(self, text):
        from textblob import TextBlob
        blob = TextBlob(text)
        sentiment_score = blob.sentiment.polarity
        if sentiment_score > 0.1:
            return "Positive"
        elif sentiment_score < -0.1:
            return "Negative"
        else:
            return "Neutral"

    def generate_text(self, prompt, max_length=100):
        input_ids = self.gpt_neo_tokenizer.encode(prompt, return_tensors='pt')
        output = self.gpt_neo_model.generate(input_ids, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2)
        return self.gpt_neo_tokenizer.decode(output[0], skip_special_tokens=True)

    def use_groq_api(self, data):
        # Example of using Groq API
        response = self.groq_client.query(self.groq_model, data)
        return response

# Initialize the NLP processor
nlp_processor = LegendaryNaturalLanguageProcessor()

# Web crawling function
async def web_crawl(url):
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, timeout=30) as response:
                if response.status == 200:
                    content = await response.text()
                    soup = BeautifulSoup(content, 'html.parser')
                    return soup.get_text()
    except Exception as e:
        return f"Error fetching {url}: {str(e)}"
    return None

# Function to get URLs for missing knowledge using the knowledge base
async def get_urls_for_missing_knowledge(topic):
    urls = knowledge_base.get(topic, [])
    if not urls:
        return []
    return urls

# Function to learn from URLs
async def learn_from_urls(topic):
    urls = await get_urls_for_missing_knowledge(topic)
    new_knowledge = ""
    for url in urls:
        content = await web_crawl(url)
        if content:
            new_knowledge += f"\n\nFrom {url}:\n{content}"
    return new_knowledge

# Function to update knowledge base
def update_knowledge_base(topic, new_knowledge):
    knowledge_base[topic] = knowledge_base.get(topic, "") + new_knowledge
    with open(KNOWLEDGE_PATH, 'w', encoding='utf-8') as f:
        json.dump(knowledge_base, f, indent=2)

# Functions to manage memory.json
def append_to_memory(entry):
    memory.append(entry)
    if len(memory) > 200:
        memory.pop(0)  # Remove the oldest entry to keep memory size at 200
    save_memory(memory)

# Function to analyze a file
def analyze_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()

    analysis = {
        'file_name': os.path.basename(file_path),
        'content': content,
        'imports': [],
        'functions': [],
        'structs': [],
    }

    if file_path.endswith('.py'):
        try:
            tree = ast.parse(content)
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for n in node.names:
                        analysis['imports'].append(n.name)
                elif isinstance(node, ast.ImportFrom):
                    for n in node.names:
                        analysis['imports'].append(f"{node.module}.{n.name}")
                elif isinstance(node, ast.FunctionDef):
                    analysis['functions'].append(node.name)
                elif isinstance(node, ast.ClassDef):
                    analysis['structs'].append(node.name)
        except SyntaxError:
            pass
    elif file_path.endswith('.rs'):
        for line in content.split('\n'):
            if line.strip().startswith('use '):
                analysis['imports'].append(line.strip())
            elif line.strip().startswith('fn '):
                match = re.match(r'fn\s+(\w+)', line.strip())
                if match:
                    analysis['functions'].append(match.group(1))
            elif line.strip().startswith('struct '):
                match = re.match(r'struct\s+(\w+)', line.strip())
                if match:
                    analysis['structs'].append(match.group(1))

    return analysis

# Worker class for managing background tasks
class Worker(QThread):
    progress = pyqtSignal(int)
    log = pyqtSignal(str)
    finished = pyqtSignal()

    def __init__(self, folder):
        super().__init__()
        self.folder = folder

    def run(self):
        files = [f for f in os.listdir(self.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.folder, file)
            analysis = analyze_file(file_path)
            self.log.emit(f"Analyzed {file}")
            self.progress.emit(int((i + 1) / total_files * 100))
            time.sleep(1)  # Simulate time-consuming analysis
        self.finished.emit()

# Function to make GROQ API requests with rate limit and fallback to local model
async def process_request_with_fallback(messages, max_retries=5):
    for attempt in range(max_retries):
        try:
            async with call_limiter, token_limiter:
                response = groq_client.chat.completions.create(
                    model=groq_model,
                    messages=messages,
                    max_tokens=2048  # Increased token capacity
                )
                if response and response.choices:
                    return response.choices[0].message.content
        except Exception as e:
            error_message = str(e)
            if 'rate_limit_exceeded' in error_message:
                print(f"Groq rate limit exceeded: {error_message}. Using local model.")
                return await generate_with_local_model(messages)
            else:
                print(f"Error during GROQ request: {error_message}. Retrying...")
                await asyncio.sleep(30)
    print("Max retries exceeded for GROQ API.")
    return await generate_with_local_model(messages)

# Function to generate response using local GPT Neo model
async def generate_with_local_model(messages):
    # Build prompt from conversation history
    prompt = ""
    for msg in messages:
        if msg["role"] == "system":
            prompt += f"{msg['content']}\n"
        elif msg["role"] == "user":
            prompt += f"User: {msg['content']}\n"
        elif msg["role"] == "assistant":
            prompt += f"AI: {msg['content']}\n"
    prompt += "AI:"

    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    max_total_tokens = 8000  # Overall maximum tokens (input and potential output)

    if input_ids.shape[1] > max_total_tokens:
        # Handle large inputs by splitting into sensible token chunks
        chunks = split_into_chunks(prompt, max_total_tokens - response_chunk_size)
    else:
        chunks = [prompt]

    responses = []
    for chunk in chunks:
        inputs = tokenizer(chunk, return_tensors="pt")
        if torch.cuda.is_available():
            inputs = {k: v.to('cuda') for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=min(inputs['input_ids'].shape[1] + response_chunk_size, model.config.n_ctx),
                do_sample=True,
                temperature=0.7
            )

        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        response = response[len(chunk):].strip()
        responses.append(response)

    return "\n".join(responses)

# Adjusted splitting function for larger sized text
def split_into_chunks(text, max_tokens):
    tokens = tokenizer.encode(text)
    chunks = []
    for i in range(0, len(tokens), max_tokens):
        chunk = tokenizer.decode(tokens[i:i + max_tokens], skip_special_tokens=True)
        chunks.append(chunk)
    return chunks

# Function to enhance code using AI tools
async def enhance_code(analysis, knowledge_base):
    prompt = (
        f"Enhance the following {'Rust' if analysis['file_name'].endswith('.rs') else 'Python'} code for a Solana Sniper Bot:\n\n"
        f"{analysis['content']}\n\n"
        f"File analysis:\n{json.dumps(analysis, indent=2)}\n\n"
        f"Solana knowledge base:\n{json.dumps(knowledge_base, indent=2)}\n\n"
        "Enhance the code by:\n"
        "1. Improving error handling\n"
        "2. Optimizing performance\n"
        "3. Ensuring best practices for Solana development\n"
        "4. Adding comments for clarity\n"
        "5. Implementing any missing functionality based on the file analysis\n"
        "6. Completing any unfinished functions or code blocks\n\n"
        "Important rules:\n"
        "- Do NOT remove or alter any existing functions\n"
        "- Do NOT weaken or remove any existing functionality\n"
        "- Only enhance or optimize code\n"
        "- Ensure all imports, functions, and structs are preserved\n\n"
        "Use the Solana knowledge base and official documentation as references."
    )
    messages = [
        {"role": "system", "content": "You are an expert Solana developer."},
        {"role": "user", "content": prompt}
    ]
    return await process_request_with_fallback(messages)

# Function to run and test the code
def run_and_test_code(file_path):
    try:
        if file_path.endswith('.py'):
            result = subprocess.run(['python', file_path], capture_output=True, text=True, timeout=30)
        elif file_path.endswith('.rs'):
            project_dir = os.path.dirname(file_path)
            cargo_toml_path = os.path.join(project_dir, 'Cargo.toml')
            if not os.path.exists(cargo_toml_path):
                return False, "", f"Cargo.toml not found in {project_dir}"
            result = subprocess.run(['cargo', 'test', '--manifest-path', cargo_toml_path], capture_output=True, text=True, timeout=180)
        return result.returncode == 0, result.stdout, result.stderr
    except subprocess.TimeoutExpired:
        return False, "", "Code execution timed out"
    except Exception as e:
        return False, "", str(e)

# Function to fix failing code using AI tools
async def fix_code(file_path, error_message, original_analysis):
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    prompt = (
        f"Given the following {'Rust' if file_path.endswith('.rs') else 'Python'} code with an error:\n\n"
        f"{content}\n\n"
        f"Error message:\n{error_message}\n\n"
        f"Original file analysis:\n{json.dumps(original_analysis, indent=2)}\n\n"
        "Please fix the code to resolve the error and ensure it runs correctly.\n"
        "Important rules:\n"
        "- Do NOT remove or alter any existing functions\n"
        "- Do NOT weaken or remove any existing functionality\n"
        "- Only fix, enhance or complete code\n"
        "- Ensure all imports, functions, and structs are preserved\n"
        "- If a function is incomplete, complete it based on its purpose\n\n"
    )
    messages = [
        {"role": "system", "content": "You are an expert Solana developer."},
        {"role": "user", "content": prompt}
    ]
    return await process_request_with_fallback(messages)

# Function to check if more knowledge is needed using AI tools
async def needs_more_knowledge(code):
    prompt = (
        f"Given the following code for a Solana Sniper Bot:\n\n{code}\n\n"
        "Do you need more information to further improve this code? Answer with 'Yes' or 'No'."
    )
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt}
    ]
    response = await process_request_with_fallback(messages)
    if response:
        response = response.strip().lower()
        if response.startswith('yes'):
            match = re.search(r'yes(?:,|\s+about\s+)?(.+)', response)
            if match:
                topic = match.group(1).strip()
            else:
                topic = "Solana development"
            return True, topic
    return False, None

# Function to handle AI Chat messages with enhanced chunking and knowledge integration
async def handle_ai_chat(message):
    message = ' '.join(message.split())  # Normalize whitespace

    # Save user message to memory
    append_to_memory({"role": "user", "content": message})

    # Prepare conversation history for context
    conversation_history = ""
    for entry in memory[-100:]:
        if entry["role"] == "user":
            conversation_history += f"User: {entry['content']}\n"
        elif entry["role"] == "assistant":
            conversation_history += f"AI: {entry['content']}\n"

    # Define larger maximum input and response token sizes
    input_chunk_size = 4096  # Maximum input chunk size based on model capacity
    response_chunk_size = 2048  # Maximum response chunk size

    # Efficiently handle large messages by splitting into chunks
    message_chunks = split_into_chunks(message, input_chunk_size)

    responses = []

    for chunk in message_chunks:
        # Process each chunk and pass the entire context
        messages = [
            {"role": "system", "content": "You are an expert Solana developer and AI assistant."},
            {"role": "user", "content": chunk}
        ]

        response = await process_request_with_fallback(messages)
        if response:
            responses.append(response)
            append_to_memory({"role": "assistant", "content": response})
        else:
            needs_knowledge, topic = await needs_more_knowledge(chunk)
            if needs_knowledge:
                new_knowledge = await learn_from_urls(topic)
                if new_knowledge:
                    update_knowledge_base(topic, new_knowledge)
                    response = await process_request_with_fallback(messages)
                    if response:
                        responses.append(response)
                        append_to_memory({"role": "assistant", "content": response})
                    else:
                        responses.append("AI: I'm still gathering information. Please try again later.")
                else:
                    responses.append("AI: I couldn't find additional information at the moment.")
            else:
                responses.append("AI: I'm sorry, I couldn't process your request at the moment.")

    final_response = "\n".join(responses)
    return final_response

# Function to generate a complete script based on user request and configuration
async def generate_script_from_request(request):
    config = load_config()

    prompt = (
        f"Generate a complete Rust script for a trading bot using the following configuration:\n\n"
        f"API RPC Path: {config.get('api_rpc_path', 'N/A')}\n"
        f"WebSocket Path: {config.get('websocket_path', 'N/A')}\n"
        f"Keypair: {config.get('keypair', 'N/A')}\n\n"
        f"User Request: {request}\n\n"
        "Ensure the script includes:\n"
        "- Proper API and WebSocket integration\n"
        "- Secure handling of keypairs\n"
        "- Efficient trading logic\n"
        "- Error handling and logging\n"
        "- Comments for clarity\n\n"
        "Provide the complete script without any explanations."
    )

    messages = [
        {"role": "system", "content": "You are an expert Rust developer specialized in trading bots."},
        {"role": "user", "content": prompt}
    ]

    return await process_request_with_fallback(messages)

# Function to load configuration from config.toml
def load_config():
    if os.path.exists('config.toml'):
        with open('config.toml', 'r', encoding='utf-8') as f:
            return toml.load(f)
    else:
        print("config.toml not found. Please create one with necessary configurations.")
        return {}

# PyQt6 GUI Application with AI Chat
class App(QWidget):
    def __init__(self):
        super().__init__()
        self.title = 'Solana Sniper Bot Analyzer and Enhancer'
        self.initUI()

    def initUI(self):
        self.setWindowTitle(self.title)
        self.setGeometry(100, 100, 900, 700)

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

        # Tab 1: Code Analyzer
        self.tab1 = QWidget()
        self.tabs.addTab(self.tab1, "Code Analyzer")
        self.tab1UI()

        # Tab 2: AI Chat
        self.tab2 = QWidget()
        self.tabs.addTab(self.tab2, "AI Chat")
        self.tab2UI()

        self.setLayout(layout)

    def tab1UI(self):
        layout = QVBoxLayout()

        self.title_label = QLabel("Solana Sniper Bot Code Analyzer and Enhancer")
        self.title_label.setStyleSheet("font-size: 18px; font-weight: bold;")
        self.title_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.title_label)

        self.select_folder_btn = QPushButton('Select Folder')
        self.select_folder_btn.clicked.connect(self.select_folder)
        layout.addWidget(self.select_folder_btn)

        self.selected_folder_label = QLabel('No folder selected.')
        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)
        layout.addWidget(self.start_btn)

        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        layout.addWidget(self.progress_bar)

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

        self.tab1.setLayout(layout)

    def tab2UI(self):
        layout = QVBoxLayout()

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

        h_layout = QHBoxLayout()
        self.chat_input = QLineEdit()
        self.chat_input.setPlaceholderText("Type your message here...")
        self.chat_input.returnPressed.connect(self.check_ctrl_enter)  # Connect to check for Ctrl+Enter
        self.send_btn = QPushButton("Send")
        self.send_btn.clicked.connect(self.send_chat)
        h_layout.addWidget(self.chat_input)
        h_layout.addWidget(self.send_btn)

        layout.addLayout(h_layout)
        self.tab2.setLayout(layout)

    def keyPressEvent(self, event):
        if event.key() == Qt.Key.Key_Return and (event.modifiers() & Qt.KeyboardModifier.ControlModifier):
            self.send_chat()

    def check_ctrl_enter(self):
        if QApplication.keyboardModifiers() == Qt.KeyboardModifier.ControlModifier:
            self.send_chat()

    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.worker = Worker(self.selected_folder)
        self.worker.progress.connect(self.update_progress)
        self.worker.log.connect(self.update_log)
        self.worker.finished.connect(self.analysis_finished)
        self.worker.start()
        self.update_log("Analysis started...")

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

    def update_log(self, message):
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        self.log_text.append(f"[{timestamp}] {message}")

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

    def send_chat(self):
        message = self.chat_input.text().strip()
        if message:
            self.chat_display.append(f"User: {message}")
            self.chat_input.clear()
            asyncio.run(self.process_chat(message))

    async def process_chat(self, message):
        if message.startswith("generate script"):
            response = await generate_script_from_request(message)
        else:
            response = await handle_ai_chat(message)

        if response:
            self.chat_display.append(f"AI: {response}")
        else:
            self.chat_display.append("AI: I'm still gathering information. Please try again later.")

# Entry point
if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = App()
    ex.show()
    sys.exit(app.exec())
 ENHANCE DONT REMOVE FUNCTIONS OR WEAKEN THE CODE BUT INSTEAD IMPROVE IT BY 1000% AND MAKE IT WORK PROPERLY ALL FUNCTIONS AND ADD SOME NEW SPECIAL THINGS YOU COME UP WITH PLEASE GO ALL OUT MAKE IT BEAUTIFUL