from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, 
                            QComboBox, QPushButton, QLabel)
from PyQt6.QtCore import pyqtSignal
import requests
from typing import List, Dict
import os
from dotenv import load_dotenv

# Загружаем переменные окружения
load_dotenv()

class Model:
    def __init__(self, name: str, title: str = None):
        self.name = name
        self.title = title or name

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

    def setup_ui(self):
        self.setStyleSheet("""
            QWidget {
                background-color: transparent;
            }
            QLabel {
                color: #d4d4d4;
                font-size: 13px;
                font-weight: 500;
                margin-right: 8px;
            }
            QComboBox {
                background-color: #252526;
                color: #d4d4d4;
                border: 1px solid #3e3e3e;
                border-radius: 6px;
                padding: 6px 12px;
                min-width: 200px;
                font-size: 13px;
            }
            QComboBox:hover {
                border-color: #565656;
            }
            QComboBox:focus {
                border-color: #007acc;
            }
            QComboBox::drop-down {
                border: none;
                width: 20px;
            }
            QComboBox::down-arrow {
                image: url(data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iMTYiIGhlaWdodD0iMTYiIHZpZXdCb3g9IjAgMCAxNiAxNiIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPHBhdGggZD0iTTggMTFMNCBTTDEyIDUiIHN0cm9rZT0iI2Q0ZDRkNCIgc3Ryb2tlLXdpZHRoPSIyIiBzdHJva2UtbGluZWNhcD0icm91bmQiIHN0cm9rZS1saW5lam9pbj0icm91bmQiLz4KPC9zdmc+);
            }
            QComboBox QAbstractItemView {
                background-color: #252526;
                color: #d4d4d4;
                border: 1px solid #3e3e3e;
                selection-background-color: #094771;
                outline: none;
            }
            QPushButton {
                background-color: #2d2d2d;
                color: #d4d4d4;
                border: 1px solid #3e3e3e;
                border-radius: 6px;
                padding: 6px 12px;
                font-size: 13px;
            }
            QPushButton:hover {
                background-color: #3e3e3e;
                border-color: #565656;
            }
            QPushButton:pressed {
                background-color: #565656;
            }
        """)

        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(8)

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

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

        # Кнопка обновления
        refresh_button = QPushButton("🔄")
        refresh_button.setFixedSize(32, 32)
        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_data = response.json()
                self.models = [
                    Model(model.get('title', ''), model.get('model_name', ''))
                    for model in models_data
                ]
                
                # Получаем текущую модель
                options_response = requests.get(f"{self.api_url}/sdapi/v1/options")
                if options_response.status_code == 200:
                    current_model_name = options_response.json().get('sd_model_checkpoint')
                    self.current_model = next(
                        (model for model in self.models if model.name == current_model_name),
                        self.models[0] if self.models else None
                    )
                
                # Обновляем UI
                self.model_combo.clear()
                for model in self.models:
                    self.model_combo.addItem(model.title)
                
                # Устанавливаем текущую модель
                if self.current_model:
                    index = self.model_combo.findText(self.current_model.title)
                    if index >= 0:
                        self.model_combo.setCurrentIndex(index)
                
        except Exception as e:
            print(f"Ошибка при загрузке моделей: {str(e)}")
            self.models = []
            self.current_model = None

    def on_model_changed(self, model_title):
        """Обработка смены модели"""
        model = next(
            (m for m in self.models if m.title == model_title),
            None
        )
        if model:
            try:
                response = requests.post(
                    f"{self.api_url}/sdapi/v1/options",
                    json={"sd_model_checkpoint": model.name}
                )
                if response.status_code == 200:
                    self.current_model = model
                    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()

    def get_models(self):
        """Возвращает список доступных моделей"""
        return self.models

    def get_model_names(self):
        """Возвращает список имен моделей"""
        return [model.name for model in self.models]

    def get_current_model(self):
        """Возвращает текущую модель"""
        return self.current_model

    def switch_model(self, model_name: str) -> bool:
        """Переключение на другую модель"""
        try:
            response = requests.post(
                f"{self.api_url}/sdapi/v1/options",
                json={"sd_model_checkpoint": model_name}
            )
            if response.status_code == 200:
                self.current_model = next(
                    (model for model in self.models if model.name == model_name),
                    None
                )
                return True
            return False
        except Exception as e:
            print(f"Ошибка переключения модели: {str(e)}")
            return False