import sys
import json
import os
import random
import urllib.request
import urllib.parse
import websocket
import requests
import uuid
import io
import logging
import re
import copy
from PIL import Image
from PIL.PngImagePlugin import PngInfo

from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QTextEdit, QSpinBox, QComboBox, 
                             QPushButton, QLabel, QFileDialog, QInputDialog, QMessageBox,
                             QCheckBox, QProgressBar, QFrame, QGroupBox, QSizePolicy)
from PyQt5.QtGui import QPixmap, QImage, QFont, QIcon
from PyQt5.QtCore import Qt, QThread, pyqtSignal

# ================= ЛОГИРОВАНИЕ =================
logger = logging.getLogger("ComfyUIRemote")
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s | %(levelname)-8s | %(message)s')
file_handler = logging.FileHandler("app_detailed_log.txt", encoding='utf-8')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)

# ================= КОНФИГУРАЦИЯ =================
CONFIG_FILE = "settings.json"
DEFAULT_CONFIG = {
    "comfyui": {"server_address": "192.168.1.118:8188"},
    "lmstudio": {
        "base_url": "http://192.168.1.118:1234/v1", 
        "api_key": "lm-studio",
        "models": ["luna-ai-llama2"]
    },
    "cloud_ai": {
        "base_url": "https://neuroapi.host/v1",
        "api_key": "ВАШ_КЛЮЧ_ЗДЕСЬ",
        "models": ["gemini-2.5-flash", "deepseek-v3.2", "grok-4-fast-reasoning"]
    },
    "ollama": {
        "base_url": "http://127.0.0.1:11434/v1"
    }
}

# ================= WORKFLOWS =================
# 1. Базовый Text-to-Image (стандартный изначальный воркфлоу)
WORKFLOW_T2I = {
  "39": {"inputs": {"filename_prefix": "ComfyUI", "images": ["35:8", 0]}, "class_type": "SaveImage"},
  "35:30": {"inputs": {"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}, "class_type": "CLIPLoader"},
  "35:29": {"inputs": {"vae_name": "ae.safetensors"}, "class_type": "VAELoader"},
  "35:28": {"inputs": {"unet_name": "SwarmUI_Z-Image-Turbo-FP8Mix.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"},
  "35:27": {"inputs": {"text": "", "clip": ["35:30", 0]}, "class_type": "CLIPTextEncode"},
  "35:13": {"inputs": {"width": 1024, "height": 1024, "batch_size": 1}, "class_type": "EmptySD3LatentImage"},
  "35:3": {"inputs": {"seed": 0, "steps": 4, "cfg": 1, "sampler_name": "res_multistep", "scheduler": "simple", "denoise": 1, "model": ["35:11", 0], "positive": ["35:27", 0], "negative": ["35:33", 0], "latent_image": ["35:13", 0]}, "class_type": "KSampler"},
  "35:11": {"inputs": {"shift": 3, "model": ["35:28", 0]}, "class_type": "ModelSamplingAuraFlow"},
  "35:33": {"inputs": {"conditioning": ["35:27", 0]}, "class_type": "ConditioningZeroOut"},
  "35:8": {"inputs": {"samples": ["35:3", 0], "vae": ["35:29", 0]}, "class_type": "VAEDecode"}
}

# 2. Image Editing (Qwen-Image-Edit)
WORKFLOW_I2I = {
  "60": {"inputs": {"filename_prefix": "ComfyUI", "images": ["102:8", 0]}, "class_type": "SaveImage"},
  "78": {"inputs": {"image": "image_qwen_image_edit_input_image.png"}, "class_type": "LoadImage"},
  "93": {"inputs": {"upscale_method": "lanczos", "megapixels": 1.5, "resolution_steps": 1, "image": ["78", 0]}, "class_type": "ImageScaleToTotalPixels"},
  "102:39": {"inputs": {"vae_name": "qwen_image_vae.safetensors"}, "class_type": "VAELoader"},
  "102:77": {"inputs": {"prompt": "", "clip": ["102:38", 0], "vae": ["102:39", 0], "image": ["78", 0]}, "class_type": "TextEncodeQwenImageEdit"},
  "102:75": {"inputs": {"strength": 1, "model": ["102:66", 0]}, "class_type": "CFGNorm"},
  "102:66": {"inputs": {"shift": 3, "model": ["102:108", 0]}, "class_type": "ModelSamplingAuraFlow"},
  "102:8": {"inputs": {"samples": ["102:3", 0], "vae": ["102:39", 0]}, "class_type": "VAEDecode"},
  "102:38": {"inputs": {"clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors", "type": "qwen_image", "device": "default"}, "class_type": "CLIPLoader"},
  "102:76": {"inputs": {"prompt": "", "clip": ["102:38", 0], "vae": ["102:39", 0], "image": ["78", 0]}, "class_type": "TextEncodeQwenImageEdit"},
  "102:88": {"inputs": {"pixels": ["78", 0], "vae": ["102:39", 0]}, "class_type": "VAEEncode"},
  "102:89": {"inputs": {"lora_name": "Qwen-Image-Edit-Lightning-4steps-V1.0-bf16.safetensors", "strength_model": 1, "model": ["102:37", 0]}, "class_type": "LoraLoaderModelOnly"},
  "102:37": {"inputs": {"unet_name": "qwen_image_edit_fp8_e4m3fn.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"},
  "102:105": {"inputs": {"value": 1}, "class_type": "PrimitiveFloat"},
  "102:106": {"inputs": {"value": 20}, "class_type": "PrimitiveInt"},
  "102:103": {"inputs": {"value": 4}, "class_type": "PrimitiveInt"},
  "102:107": {"inputs": {"value": 2.5}, "class_type": "PrimitiveFloat"},
  "102:111": {"inputs": {"value": False}, "class_type": "PrimitiveBoolean"},
  "102:3": {"inputs": {"seed": 0, "steps": ["102:110", 0], "cfg": ["102:109", 0], "sampler_name": "euler", "scheduler": "simple", "denoise": 1, "model": ["102:75", 0], "positive": ["102:76", 0], "negative": ["102:77", 0], "latent_image": ["102:88", 0]}, "class_type": "KSampler"},
  "102:109": {"inputs": {"switch": ["102:111", 0], "on_false": ["102:107", 0], "on_true": ["102:105", 0]}, "class_type": "ComfySwitchNode"},
  "102:110": {"inputs": {"switch": ["102:111", 0], "on_false": ["102:106", 0], "on_true": ["102:103", 0]}, "class_type": "ComfySwitchNode"},
  "102:108": {"inputs": {"switch": ["102:111", 0], "on_false": ["102:37", 0], "on_true": ["102:89", 0]}, "class_type": "ComfySwitchNode"}
}

# ================= СТИЛИ (QSS) =================
DARK_STYLE = """
QMainWindow { background-color: #1e1e2e; }
QGroupBox { color: #cdd6f4; font-weight: bold; border: 1px solid #45475a; margin-top: 15px; border-radius: 8px; padding-top: 20px; }
QLabel { color: #bac2de; font-size: 13px; }
QTextEdit { background-color: #313244; color: #cdd6f4; border: 1px solid #45475a; border-radius: 5px; padding: 10px; selection-background-color: #585b70; }
QSpinBox, QComboBox { background-color: #313244; color: #cdd6f4; border: 1px solid #45475a; border-radius: 4px; padding: 5px; min-height: 30px; }
QPushButton { background-color: #45475a; color: #cdd6f4; border-radius: 6px; padding: 10px; font-weight: bold; }
QPushButton:hover { background-color: #585b70; }
QPushButton#start_btn { background-color: #a6e3a1; color: #11111b; font-size: 14px; }
QPushButton#start_btn:hover { background-color: #94e2d5; }
QPushButton#save_btn { background-color: #f38ba8; color: #11111b; }
QProgressBar { border: 1px solid #45475a; border-radius: 5px; text-align: center; color: white; height: 15px; }
QProgressBar::chunk { background-color: #89b4fa; }
QCheckBox { color: #bac2de; font-size: 13px; spacing: 8px; }
"""

# ================= WORKERS =================

class ComfyUIWorker(QThread):
    finished_signal = pyqtSignal(QPixmap, bytes)
    error_signal = pyqtSignal(str)
    progress_signal = pyqtSignal(int, int)

    def __init__(self, workflow, server_address, upload_image_path=None, image_node_id=None):
        super().__init__()
        self.workflow = workflow
        self.server_address = server_address
        self.client_id = str(uuid.uuid4())
        self.upload_image_path = upload_image_path
        self.image_node_id = image_node_id

    def run(self):
        ws = None
        try:
            # 1. Если задан путь к картинке (I2I Workflow) -> загружаем ее на сервер
            if self.upload_image_path and self.image_node_id:
                try:
                    with open(self.upload_image_path, 'rb') as f:
                        upload_url = f"http://{self.server_address}/upload/image"
                        files = {'image': f}
                        req = requests.post(upload_url, files=files)
                        if req.status_code == 200:
                            # Получаем реальное имя файла на сервере
                            filename = req.json()['name']
                            self.workflow[self.image_node_id]["inputs"]["image"] = filename
                        else:
                            self.error_signal.emit("Failed to upload image to ComfyUI.")
                            return
                except Exception as e:
                    self.error_signal.emit(f"Image upload error: {str(e)}")
                    return

            # 2. Подключаемся к WebSocket и отправляем workflow
            ws_url = f"ws://{self.server_address}/ws?clientId={self.client_id}"
            ws = websocket.WebSocket()
            ws.connect(ws_url)
            prompt_url = f"http://{self.server_address}/prompt"
            data = json.dumps({"prompt": self.workflow, "client_id": self.client_id}).encode('utf-8')
            req = urllib.request.Request(prompt_url, data=data)
            
            with urllib.request.urlopen(req) as response:
                prompt_id = json.loads(response.read())['prompt_id']
                
            # Ждем выполнения
            while True:
                out = ws.recv()
                if isinstance(out, str):
                    message = json.loads(out)
                    if message['type'] == 'progress':
                        self.progress_signal.emit(message['data']['value'], message['data']['max'])
                    if message['type'] == 'executing' and message['data']['node'] is None:
                        if message['data']['prompt_id'] == prompt_id: break
                        
            # Получаем результат
            with urllib.request.urlopen(f"http://{self.server_address}/history/{prompt_id}") as response:
                history = json.loads(response.read())[prompt_id]
                
            for node_id in history['outputs']:
                node_output = history['outputs'][node_id]
                if 'images' in node_output:
                    img_data = node_output['images'][0]
                    view_url = f"http://{self.server_address}/view?filename={img_data['filename']}&subfolder={img_data['subfolder']}&type={img_data['type']}"
                    with urllib.request.urlopen(view_url) as img_res:
                        img_bytes = img_res.read()
                        self.finished_signal.emit(QPixmap.fromImage(QImage.fromData(img_bytes)), img_bytes)
                        return
                        
            self.error_signal.emit("Generation finished, but no image found in outputs.")
        except Exception as e: 
            self.error_signal.emit(str(e))
        finally: 
            if ws: ws.close()

class LLMWorker(QThread):
    finished_signal = pyqtSignal(str)
    error_signal = pyqtSignal(str)
    
    def __init__(self, base_url, api_key, model, prompt, provider="default"):
        super().__init__()
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key
        self.model = model
        self.prompt = prompt
        self.provider = provider

    def run(self):
        try:
            headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
            
            if self.provider == "ollama":
                system_instruction = (
                    "You are a professional image prompt engineer. "
                    "Task: Translate the user's input into ENGLISH (if it's not already) and expand it into a highly detailed visual prompt. "
                    "CRITICAL RULES: "
                    "1. MUST BE IN ENGLISH ONLY. Do NOT output Russian words. "
                    "2. NO REASONING, NO THOUGHTS. Do not explain your process. "
                    "3. OUTPUT EXACTLY THE RAW PROMPT. No introductions like 'Here is...' or 'Sure'. "
                )
                temperature = 0.6
            else:
                system_instruction = "Improve prompt for AI image generation. Expanded English description only."
                temperature = 0.7

            payload = {
                "model": self.model,
                "messages": [
                    {"role": "system", "content": system_instruction},
                    {"role": "user", "content": self.prompt}
                ],
                "temperature": temperature
            }
            
            response = requests.post(f"{self.base_url}/chat/completions", headers=headers, json=payload, timeout=180)
            
            if response.status_code == 200:
                result_text = response.json()['choices'][0]['message']['content'].strip()
                result_text = re.sub(r'<think>.*?</think>', '', result_text, flags=re.DOTALL).strip()
                if self.provider == "ollama":
                    if result_text.startswith('"') and result_text.endswith('"'):
                        result_text = result_text[1:-1].strip()
                self.finished_signal.emit(result_text)
            else: 
                self.error_signal.emit(f"Error {response.status_code}: {response.text}")
                
        except requests.exceptions.Timeout:
            self.error_signal.emit("Тайм-аут: Сервер слишком долго загружал модель или генерировал ответ. Попробуйте еще раз.")
        except Exception as e: 
            self.error_signal.emit(str(e))

# ================= ГЛАВНОЕ ОКНО =================

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("ComfyUI Studio v3")
        self.resize(1280, 850)
        self.setStyleSheet(DARK_STYLE)
        self.load_config()
        self.current_bytes = None
        self.selected_image_path = None
        self.init_ui()

    def load_config(self):
        if not os.path.exists(CONFIG_FILE):
            with open(CONFIG_FILE, 'w', encoding='utf-8') as f: json.dump(DEFAULT_CONFIG, f, indent=4)
            self.config = DEFAULT_CONFIG
        else:
            try:
                with open(CONFIG_FILE, 'r', encoding='utf-8') as f: self.config = json.load(f)
            except: self.config = DEFAULT_CONFIG

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(15, 15, 15, 15)

        left_panel = QFrame()
        left_panel.setMinimumWidth(550) 
        left_layout = QVBoxLayout(left_panel)
        left_layout.setContentsMargins(0, 0, 0, 0)

        # Выбор Workflow
        wf_group = QGroupBox("Workflow Selector")
        wfg_layout = QVBoxLayout()
        self.workflow_cb = QComboBox()
        self.workflow_cb.addItems([
            "Text-to-Image (Default Turbo)", 
            "Image Editing (Qwen I2I)"
        ])
        self.workflow_cb.currentIndexChanged.connect(self.on_workflow_changed)
        wfg_layout.addWidget(self.workflow_cb)
        wf_group.setLayout(wfg_layout)
        left_layout.addWidget(wf_group)

        # Prompt Input
        prompt_group = QGroupBox("Prompt Input")
        pg_layout = QVBoxLayout()
        self.prompt_edit = QTextEdit()
        self.prompt_edit.setPlaceholderText("Describe your image idea here...")
        self.prompt_edit.setMinimumHeight(150) 
        pg_layout.addWidget(self.prompt_edit)
        
        self.neg_check = QCheckBox("Use Negative Prompt")
        self.neg_check.toggled.connect(lambda c: self.neg_edit.setVisible(c))
        pg_layout.addWidget(self.neg_check)
        self.neg_edit = QTextEdit()
        self.neg_edit.setPlaceholderText("Describe what to exclude...")
        self.neg_edit.setMinimumHeight(80)
        self.neg_edit.setVisible(False)
        pg_layout.addWidget(self.neg_edit)
        prompt_group.setLayout(pg_layout)
        left_layout.addWidget(prompt_group, 3) 

        # Блок Изображения (Только для Image Editing)
        self.image_picker_group = QGroupBox("Input Image Settings")
        ip_layout = QVBoxLayout()
        btn_img_layout = QHBoxLayout()
        self.btn_select_image = QPushButton("Select Local Image")
        self.btn_select_image.clicked.connect(self.select_input_image)
        self.lbl_input_image = QLabel("No image selected")
        self.lbl_input_image.setWordWrap(True)
        btn_img_layout.addWidget(self.btn_select_image)
        btn_img_layout.addWidget(self.lbl_input_image)
        self.fast_edit_cb = QCheckBox("Enable 4-Steps LoRA (Fast Edit)")
        self.fast_edit_cb.setChecked(True)
        ip_layout.addLayout(btn_img_layout)
        ip_layout.addWidget(self.fast_edit_cb)
        self.image_picker_group.setLayout(ip_layout)
        self.image_picker_group.setVisible(False) # Скрыто по умолчанию
        left_layout.addWidget(self.image_picker_group)

        # Settings Group
        settings_group = QGroupBox("Core Parameters")
        sg_layout = QVBoxLayout()
        
        row1 = QHBoxLayout()
        self.step_s = QSpinBox(); self.step_s.setRange(1, 100); self.step_s.setValue(4)
        row1.addWidget(QLabel("Steps:")); row1.addWidget(self.step_s)
        self.seed_s = QSpinBox(); self.seed_s.setRange(0, 2147483647); self.seed_s.setValue(random.randint(0, 1000000))
        row1.addWidget(QLabel("Seed:")); row1.addWidget(self.seed_s)
        sg_layout.addLayout(row1)

        row2 = QHBoxLayout()
        self.size_cb = QComboBox()
        self.size_cb.addItems(["1024x1024", "832x1216", "1216x832", "1920x1080"])
        row2.addWidget(QLabel("Resolution:")); row2.addWidget(self.size_cb)
        self.font_cb = QComboBox()
        self.font_cb.addItems(["10", "12", "14", "16", "18", "20", "24", "26"]); self.font_cb.setCurrentText("14")
        self.font_cb.currentIndexChanged.connect(self.update_font)
        row2.addWidget(QLabel("Text Size:")); row2.addWidget(self.font_cb)
        sg_layout.addLayout(row2)

        self.vae_cb = QComboBox()
        self.vae_cb.addItem("ae.safetensors")
        sg_layout.addWidget(QLabel("VAE Checkpoint (T2I):"))
        sg_layout.addWidget(self.vae_cb)
        settings_group.setLayout(sg_layout)
        left_layout.addWidget(settings_group, 1)

        # AI Tools Group
        ai_group = QGroupBox("AI Prompt Assistance")
        ag_layout = QHBoxLayout()
        
        self.btn_lms = QPushButton("LM Studio")
        self.btn_lms.clicked.connect(self.use_lms)
        self.btn_ollama = QPushButton("Ollama")
        self.btn_ollama.setStyleSheet("color: #fab387;")
        self.btn_ollama.clicked.connect(self.use_ollama)
        self.btn_ai = QPushButton("NeuroAPI")
        self.btn_ai.setStyleSheet("color: #ca9ee6;")
        self.btn_ai.clicked.connect(self.use_ai)
        
        ag_layout.addWidget(self.btn_lms)
        ag_layout.addWidget(self.btn_ollama)
        ag_layout.addWidget(self.btn_ai)
        ai_group.setLayout(ag_layout)
        left_layout.addWidget(ai_group)

        main_layout.addWidget(left_panel, 2) 

        # ПРАВАЯ ПАНЕЛЬ (ПРОСМОТР)
        right_panel = QFrame()
        right_layout = QVBoxLayout(right_panel)
        right_layout.setContentsMargins(0, 0, 0, 0)

        self.img_lbl = QLabel("Image Viewer")
        self.img_lbl.setAlignment(Qt.AlignCenter)
        self.img_lbl.setStyleSheet("background-color: #11111b; border-radius: 10px; border: 2px solid #313244;")
        self.img_lbl.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        right_layout.addWidget(self.img_lbl)

        self.pbar = QProgressBar()
        right_layout.addWidget(self.pbar)

        controls = QHBoxLayout()
        self.btn_start = QPushButton("▶  START GENERATION")
        self.btn_start.setObjectName("start_btn")
        self.btn_start.setFixedHeight(55)
        self.btn_start.clicked.connect(self.run_comfy)
        
        btn_upd = QPushButton("🔄")
        btn_upd.setFixedWidth(60); btn_upd.setFixedHeight(55)
        btn_upd.clicked.connect(self.refresh)
        
        btn_sav = QPushButton("💾 SAVE IMAGE")
        btn_sav.setObjectName("save_btn")
        btn_sav.setFixedHeight(55)
        btn_sav.clicked.connect(self.save_img)
        
        controls.addWidget(self.btn_start, 3)
        controls.addWidget(btn_upd, 1)
        controls.addWidget(btn_sav, 2)
        right_layout.addLayout(controls)

        main_layout.addWidget(right_panel, 3) 
        
        self.update_font()
        self.fetch_vaes_from_comfy()

    def update_font(self):
        f = QFont(); f.setPointSize(int(self.font_cb.currentText()))
        self.prompt_edit.setFont(f); self.neg_edit.setFont(f)

    def fetch_vaes_from_comfy(self):
        server = self.config['comfyui']['server_address']
        try:
            resp = requests.get(f"http://{server}/object_info/VAELoader", timeout=2)
            if resp.status_code == 200:
                vaes = resp.json()['VAELoader']['input']['required']['vae_name'][0]
                self.vae_cb.clear(); self.vae_cb.addItems(vaes)
        except: pass

    def on_workflow_changed(self):
        idx = self.workflow_cb.currentIndex()
        if idx == 1: # Qwen Image Edit (I2I)
            self.image_picker_group.setVisible(True)
            self.step_s.setEnabled(False) # Регулируется через Fast Mode
            self.size_cb.setEnabled(False) # Регулируется через исходное изображение
            self.vae_cb.setEnabled(False) # Воркфлоу требует специфичный qwen_image_vae
        else: # Text-to-Image Default
            self.image_picker_group.setVisible(False)
            self.step_s.setEnabled(True)
            self.size_cb.setEnabled(True)
            self.vae_cb.setEnabled(True)

    def select_input_image(self):
        path, _ = QFileDialog.getOpenFileName(self, "Select Input Image", "", "Images (*.png *.jpg *.jpeg *.webp)")
        if path:
            self.selected_image_path = path
            self.lbl_input_image.setText(os.path.basename(path))

    def run_comfy(self):
        txt = self.prompt_edit.toPlainText().strip()
        if not txt: 
            QMessageBox.warning(self, "Warning", "Prompt cannot be empty!")
            return
            
        self.btn_start.setEnabled(False)
        self.btn_start.setText("G E N E R A T I N G . . .")
        
        idx = self.workflow_cb.currentIndex()
        upload_path = None
        image_node = None
        
        if idx == 0:
            # --- Text to Image ---
            wf = copy.deepcopy(WORKFLOW_T2I)
            wf["35:27"]["inputs"]["text"] = txt
            wf["35:3"]["inputs"]["steps"] = self.step_s.value()
            wf["35:3"]["inputs"]["seed"] = self.seed_s.value()
            w, h = map(int, self.size_cb.currentText().split('x'))
            wf["35:13"]["inputs"]["width"] = w
            wf["35:13"]["inputs"]["height"] = h
            if self.vae_cb.currentText():
                wf["35:29"]["inputs"]["vae_name"] = self.vae_cb.currentText()
                
        elif idx == 1:
            # --- Image to Image / Qwen Edit ---
            if not self.selected_image_path:
                QMessageBox.warning(self, "Warning", "Please select an input image first!")
                self.btn_start.setEnabled(True)
                self.btn_start.setText("▶  START GENERATION")
                return
                
            wf = copy.deepcopy(WORKFLOW_I2I)
            wf["102:76"]["inputs"]["prompt"] = txt
            neg_txt = self.neg_edit.toPlainText().strip() if self.neg_check.isChecked() else ""
            wf["102:77"]["inputs"]["prompt"] = neg_txt
            wf["102:3"]["inputs"]["seed"] = self.seed_s.value()
            
            # Включаем или выключаем 4-steps LoRA (True = 4 Steps/CFG 1, False = 20 Steps/CFG 2.5)
            wf["102:111"]["inputs"]["value"] = self.fast_edit_cb.isChecked()
            
            upload_path = self.selected_image_path
            image_node = "78" # Узел LoadImage в Qwen-воркфлоу

        # Передаем параметры в Worker
        self.worker = ComfyUIWorker(wf, self.config['comfyui']['server_address'], upload_path, image_node)
        self.worker.finished_signal.connect(self.on_comfy_finished)
        self.worker.error_signal.connect(self.on_comfy_error)
        self.worker.progress_signal.connect(lambda v, m: self.pbar.setValue(int(v/m*100) if m > 0 else 0))
        self.worker.start()

    def on_comfy_finished(self, pix, raw):
        self.current_bytes = raw
        self.img_lbl.setPixmap(pix.scaled(self.img_lbl.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
        self.btn_start.setEnabled(True)
        self.btn_start.setText("▶  START GENERATION")

    def on_comfy_error(self, err):
        QMessageBox.critical(self, "Generation Error", err)
        self.btn_start.setEnabled(True)
        self.btn_start.setText("▶  START GENERATION")

    def refresh(self):
        self.seed_s.setValue(random.randint(0, 2147483647))
        self.run_comfy()

    def save_img(self):
        if not self.current_bytes: return
        path, _ = QFileDialog.getSaveFileName(self, "Save Image", "generated_art.png", "*.png")
        if path:
            img = Image.open(io.BytesIO(self.current_bytes))
            img.save(path)

    def use_lms(self):
        models = self.config['lmstudio'].get('models', [])
        m, ok = QInputDialog.getItem(self, "LM Studio", "Select Model:", models, 0, False)
        if ok: self._run_llm(self.config['lmstudio']['base_url'], self.config['lmstudio']['api_key'], m, "lmstudio")

    def use_ollama(self):
        url = self.config.get('ollama', {}).get('base_url', 'http://127.0.0.1:11434/v1')
        try:
            resp = requests.get(f"{url}/models", timeout=3)
            if resp.status_code == 200:
                models = [m['id'] for m in resp.json().get('data', [])]
            else:
                raise Exception(f"HTTP Error {resp.status_code}")
        except Exception as e:
            QMessageBox.warning(self, "Ollama Error", f"Failed to fetch models from Ollama.\nEnsure Ollama is running.\nError: {e}")
            return
            
        if not models:
            QMessageBox.warning(self, "Ollama Info", "No models found. Please pull a model first.")
            return

        m, ok = QInputDialog.getItem(self, "Ollama", "Select Model:", models, 0, False)
        if ok: self._run_llm(url, "ollama", m, "ollama")

    def use_ai(self):
        models = self.config['cloud_ai'].get('models', [])
        m, ok = QInputDialog.getItem(self, "NeuroAPI", "Select Model:", models, 0, False)
        if ok: self._run_llm(self.config['cloud_ai']['base_url'], self.config['cloud_ai']['api_key'], m, "cloud_ai")

    def _run_llm(self, url, key, model, provider="default"):
        p = self.prompt_edit.toPlainText().strip()
        if not p: return
        self.prompt_edit.setReadOnly(True)
        self.llm_worker = LLMWorker(url, key, model, p, provider)
        self.llm_worker.finished_signal.connect(self._on_llm_ok)
        self.llm_worker.error_signal.connect(self._on_llm_err)
        self.llm_worker.start()

    def _on_llm_ok(self, t):
        self.prompt_edit.setPlainText(t); self.prompt_edit.setReadOnly(False)

    def _on_llm_err(self, e):
        QMessageBox.warning(self, "AI Error", e); self.prompt_edit.setReadOnly(False)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setFont(QFont("Segoe UI", 10))
    window = MainWindow(); window.show()
    sys.exit(app.exec_())