from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, 
                            QComboBox, QPushButton, QLabel)
from PyQt6.QtCore import pyqtSignal
import requests

class ModelManager(QWidget):
    model_changed = pyqtSignal(str)
    
    def __init__(self, api_url, parent=None):
        super().__init__(parent)
        self.api_url = api_url
        self.setup_ui()
        self.load_models()

    def setup_ui(self):
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)

        # Метка
        label = QLabel("Модель:")
        layout.addWidget(label)

        # Выпадающий список моделей
        self.model_combo = QComboBox()
        self.model_combo.currentTextChanged.connect(self.on_model_changed)
        layout.addWidget(self.model_combo)

        # Кнопка обновления
        refresh_button = QPushButton("🔄")
        refresh_button.setToolTip("Обновить список моделей")
        refresh_button.clicked.connect(self.load_models)
        layout.addWidget(refresh_button)

        layout.addStretch()

    def load_models(self):
        """Загрузить список доступных моделей"""
        try:
            response = requests.get(f"{self.api_url}/sdapi/v1/sd-models")
            if response.status_code == 200:
                models = response.json()
                current_model = self.model_combo.currentText()
                
                self.model_combo.clear()
                for model in models:
                    self.model_combo.addItem(model.get('title', ''))
                
                # Восстанавливаем выбранную модель, если она всё ещё доступна
                index = self.model_combo.findText(current_model)
                if index >= 0:
                    self.model_combo.setCurrentIndex(index)
        except Exception as e:
            print(f"Ошибка при загрузке моделей: {str(e)}")

    def on_model_changed(self, model_name):
        """Обработка смены модели"""
        try:
            response = requests.post(
                f"{self.api_url}/sdapi/v1/options",
                json={"sd_model_checkpoint": model_name}
            )
            if response.status_code == 200:
                self.model_changed.emit(model_name)
        except Exception as e:
            print(f"Ошибка при смене модели: {str(e)}")

    def update_api_url(self, new_url):
        """Обновить URL API"""
        self.api_url = new_url
        self.load_models()