#!/usr/bin/env python3
"""
Gemini Image Generator App - PySide6 Version
Modern GUI application for generating images using Google's Gemini API
"""

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QDialog, QWidget,
    QVBoxLayout, QHBoxLayout, QFormLayout, QGridLayout,
    QLabel, QLineEdit, QPushButton, QCheckBox, QTextEdit,
    QComboBox, QScrollArea, QGroupBox, QFileDialog, QMessageBox
)
from PySide6.QtCore import Qt, QThread, Signal, QSize
from PySide6.QtGui import QPixmap, QFont, QIcon, QDesktopServices
from PySide6.QtCore import QUrl

import base64
import io
import json
import os
import sys
import tempfile
import platform
from pathlib import Path
from google import genai
from google.genai import types
import hashlib
import pymysql
from datetime import datetime


def hash_password(password: str) -> str:
    """使用 SHA256 哈希密码"""
    return hashlib.sha256(password.encode('utf-8')).hexdigest()


class DatabaseManager:
    """数据库连接管理类"""
    def __init__(self, db_config):
        self.config = db_config

    def authenticate(self, username, password):
        """
        验证用户凭证
        返回: (success: bool, message: str)
        """
        try:
            # 计算密码哈希
            password_hash = hash_password(password)

            # 连接数据库
            conn = pymysql.connect(
                host=self.config['host'],
                port=self.config.get('port', 3306),
                user=self.config['user'],
                password=self.config['password'],
                database=self.config['database'],
                connect_timeout=5
            )

            try:
                with conn.cursor() as cursor:
                    # 使用参数化查询防止 SQL 注入
                    sql = f"SELECT * FROM {self.config['table']} WHERE user_name=%s AND passwd=%s AND status='active'"
                    cursor.execute(sql, (username, password_hash))
                    result = cursor.fetchone()

                    if result:
                        return True, "认证成功"
                    else:
                        return False, "用户名或密码错误"
            finally:
                conn.close()

        except pymysql.OperationalError as e:
            return False, "无法连接到服务器,请检查网络连接"
        except Exception as e:
            return False, f"认证失败: {str(e)}"


class LoginDialog(QDialog):
    """Qt-based login dialog"""

    def __init__(self, db_config, last_user="", saved_password_hash=""):
        super().__init__()
        self.db_config = db_config
        self.last_user = last_user
        self.saved_password_hash = saved_password_hash

        self.success = False
        self.authenticated_user = ""
        self.password_changed = False
        self.current_password_hash = ""

        self.setup_ui()
        self.apply_styles()

    def setup_ui(self):
        """Build login dialog UI"""
        self.setWindowTitle("登录 - AI 图像生成器")
        self.setFixedSize(400, 400)

        # Main layout
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(40, 40, 40, 40)
        main_layout.setSpacing(20)

        # Title
        title_label = QLabel("登录")
        title_label.setObjectName("title")
        title_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(title_label)

        main_layout.addSpacing(10)

        # Username field
        username_label = QLabel("用户名")
        username_label.setObjectName("field_label")
        main_layout.addWidget(username_label)

        self.username_entry = QLineEdit()
        self.username_entry.setText(self.last_user)
        main_layout.addWidget(self.username_entry)

        main_layout.addSpacing(10)

        # Password field
        password_label = QLabel("密码")
        password_label.setObjectName("field_label")
        main_layout.addWidget(password_label)

        self.password_entry = QLineEdit()
        self.password_entry.setEchoMode(QLineEdit.Password)

        # Handle saved password placeholder
        if self.saved_password_hash:
            self.password_entry.setPlaceholderText("••••••••")
            self.password_entry.setStyleSheet("QLineEdit { color: #999999; }")

        self.password_entry.textChanged.connect(self.on_password_change)
        self.password_entry.returnPressed.connect(self.on_login)

        main_layout.addWidget(self.password_entry)

        # Checkboxes
        checkbox_layout = QHBoxLayout()
        self.remember_user_check = QCheckBox("记住用户名")
        self.remember_user_check.setChecked(bool(self.last_user))
        checkbox_layout.addWidget(self.remember_user_check)

        self.remember_password_check = QCheckBox("记住密码")
        self.remember_password_check.setChecked(bool(self.saved_password_hash))
        checkbox_layout.addWidget(self.remember_password_check)

        checkbox_layout.addStretch()
        main_layout.addLayout(checkbox_layout)

        # Login button
        self.login_button = QPushButton("登录")
        self.login_button.setObjectName("login_button")
        self.login_button.clicked.connect(self.on_login)
        main_layout.addWidget(self.login_button)

        # Error label
        self.error_label = QLabel("")
        self.error_label.setObjectName("error_label")
        self.error_label.setAlignment(Qt.AlignCenter)
        self.error_label.setWordWrap(True)
        main_layout.addWidget(self.error_label)

        main_layout.addStretch()

        self.setLayout(main_layout)

        # Set focus
        if self.last_user:
            self.password_entry.setFocus()
        else:
            self.username_entry.setFocus()

    def apply_styles(self):
        """Apply QSS stylesheet"""
        self.setStyleSheet("""
            QDialog {
                background-color: #ffffff;
            }
            QLabel#title {
                font-size: 20pt;
                font-weight: bold;
                color: #1d1d1f;
            }
            QLabel#field_label {
                font-size: 10pt;
                color: #666666;
            }
            QLineEdit {
                padding: 8px;
                border: 1px solid #e5e5e5;
                border-radius: 4px;
                background-color: #fafafa;
                font-size: 11pt;
                color: #000000;
            }
            QLineEdit:focus {
                border: 1px solid #007AFF;
            }
            QPushButton#login_button {
                background-color: #007AFF;
                color: white;
                font-size: 10pt;
                font-weight: bold;
                padding: 10px 20px;
                border: none;
                border-radius: 6px;
            }
            QPushButton#login_button:hover {
                background-color: #0051D5;
            }
            QPushButton#login_button:pressed {
                background-color: #003D99;
            }
            QPushButton#login_button:disabled {
                background-color: #cccccc;
            }
            QCheckBox {
                font-size: 9pt;
                color: #1d1d1f;
            }
            QLabel#error_label {
                color: #ff3b30;
                font-size: 9pt;
            }
        """)

    def on_password_change(self):
        """Handle password field changes"""
        if not self.password_changed and self.saved_password_hash:
            # First change - clear placeholder style
            self.password_entry.setStyleSheet("")
            self.password_changed = True

    def on_login(self):
        """Handle login button click"""
        username = self.username_entry.text().strip()
        password_input = self.password_entry.text()

        # Validate input
        if not username:
            self.show_error("请输入用户名")
            return

        # Check if password is provided or saved password exists
        if not password_input and not self.saved_password_hash:
            self.show_error("请输入密码")
            return

        # Disable button during authentication
        self.login_button.setEnabled(False)
        self.error_label.setText("正在验证...")
        self.error_label.setStyleSheet("QLabel { color: #666666; }")

        # Determine which password to use
        if not self.password_changed and self.saved_password_hash:
            password_hash = self.saved_password_hash
        else:
            password_hash = hash_password(password_input)

        # Authenticate
        try:
            conn = pymysql.connect(
                host=self.db_config['host'],
                port=self.db_config.get('port', 3306),
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database'],
                connect_timeout=5
            )

            try:
                with conn.cursor() as cursor:
                    sql = f"SELECT * FROM {self.db_config['table']} WHERE user_name=%s AND passwd=%s AND status='active'"
                    cursor.execute(sql, (username, password_hash))
                    result = cursor.fetchone()

                    if result:
                        self.success = True
                        self.authenticated_user = username
                        self.current_password_hash = password_hash
                        self.accept()  # Close dialog with success
                    else:
                        self.show_error("用户名或密码错误")
                        self.password_entry.clear()
                        self.password_changed = False
                        self.login_button.setEnabled(True)
            finally:
                conn.close()

        except pymysql.OperationalError:
            self.show_error("无法连接到服务器,请检查网络连接")
            self.login_button.setEnabled(True)
        except Exception as e:
            self.show_error(f"认证失败: {str(e)}")
            self.login_button.setEnabled(True)

    def show_error(self, message):
        """Display error message"""
        self.error_label.setText(message)
        self.error_label.setStyleSheet("QLabel { color: #ff3b30; }")

    def get_remember_user(self):
        """Get remember username checkbox state"""
        return self.remember_user_check.isChecked()

    def get_remember_password(self):
        """Get remember password checkbox state"""
        return self.remember_password_check.isChecked()

    def get_password_hash(self):
        """Get current password hash"""
        return getattr(self, 'current_password_hash', '')


class ImageGeneratorWindow(QMainWindow):
    """Qt-based main application window"""

    def __init__(self):
        super().__init__()
        self.api_key = ""
        self.uploaded_images = []  # List of file paths
        self.generated_image_data = None
        self.generated_image_bytes = None
        self.saved_prompts = []
        self.db_config = None
        self.last_user = ""
        self.saved_password_hash = ""

        self.load_config()
        self.setup_ui()
        self.apply_styles()

    def get_config_dir(self):
        """Get the appropriate directory for config files based on platform and mode"""
        if getattr(sys, 'frozen', False):
            system = platform.system()
            if system == 'Darwin':
                config_dir = Path.home() / 'Library' / 'Application Support' / 'ZB100ImageGenerator'
            elif system == 'Windows':
                config_dir = Path(os.getenv('APPDATA', Path.home())) / 'ZB100ImageGenerator'
            else:
                config_dir = Path.home() / '.config' / 'zb100imagegenerator'
        else:
            config_dir = Path('.')

        config_dir.mkdir(parents=True, exist_ok=True)
        return config_dir

    def get_config_path(self):
        """Get the full path to config.json"""
        return self.get_config_dir() / 'config.json'

    def load_config(self):
        """Load API key, saved prompts, and db config from config file"""
        config_path = self.get_config_path()

        if config_path.exists():
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    self.api_key = config.get("api_key", "")
                    self.saved_prompts = config.get("saved_prompts", [])
                    self.db_config = config.get("db_config")
                    self.last_user = config.get("last_user", "")
                    self.saved_password_hash = config.get("saved_password_hash", "")
            except Exception as e:
                print(f"Failed to load config from {config_path}: {e}")

        if not self.api_key and getattr(sys, 'frozen', False):
            try:
                bundle_dir = Path(sys._MEIPASS)
                bundled_config = bundle_dir / 'config.json'
                if bundled_config.exists():
                    with open(bundled_config, 'r', encoding='utf-8') as f:
                        config = json.load(f)
                        self.api_key = config.get("api_key", "")
                        self.db_config = config.get("db_config")
                    self.save_config()
            except Exception as e:
                print(f"Failed to load bundled config: {e}")

        if not self.api_key:
            QMessageBox.warning(self, "警告", f"未找到API密钥\n配置文件位置: {config_path}\n\n请在应用中输入API密钥或手动编辑配置文件")

    def save_config(self, last_user=None):
        """Save configuration to file"""
        config_path = self.get_config_path()

        try:
            config = {
                "api_key": self.api_key,
                "saved_prompts": self.saved_prompts
            }

            if self.db_config:
                config["db_config"] = self.db_config

            if last_user is not None:
                config["last_user"] = last_user
            elif hasattr(self, 'last_user'):
                config["last_user"] = self.last_user
            else:
                config["last_user"] = ""

            if hasattr(self, 'saved_password_hash'):
                config["saved_password_hash"] = self.saved_password_hash
            else:
                config["saved_password_hash"] = ""

            config_path.parent.mkdir(parents=True, exist_ok=True)

            with open(config_path, 'w', encoding='utf-8') as f:
                json.dump(config, f, indent=2, ensure_ascii=False)
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存配置失败: {str(e)}\n路径: {config_path}")

    def setup_ui(self):
        """Setup the user interface"""
        self.setWindowTitle("AI 图像生成器")
        self.setGeometry(100, 100, 1200, 850)
        self.setMinimumSize(1000, 700)

        # Central widget
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(15)

        # Reference images section
        ref_group = QGroupBox("参考图片")
        ref_layout = QVBoxLayout()

        # Upload button and count
        upload_header = QHBoxLayout()
        upload_btn = QPushButton("+ 添加图片")
        upload_btn.clicked.connect(self.upload_images)
        upload_header.addWidget(upload_btn)

        self.image_count_label = QLabel("已选择 0 张")
        upload_header.addWidget(self.image_count_label)
        upload_header.addStretch()
        ref_layout.addLayout(upload_header)

        # Image preview scroll area
        self.img_scroll = QScrollArea()
        self.img_scroll.setWidgetResizable(True)
        self.img_scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
        self.img_scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.img_scroll.setFixedHeight(140)

        self.img_container = QWidget()
        self.img_layout = QHBoxLayout()
        self.img_layout.addStretch()
        self.img_container.setLayout(self.img_layout)
        self.img_scroll.setWidget(self.img_container)

        ref_layout.addWidget(self.img_scroll)
        ref_group.setLayout(ref_layout)
        main_layout.addWidget(ref_group)

        # Content row: Prompt (left) + Settings (right)
        content_row = QHBoxLayout()

        # Prompt section
        prompt_group = QGroupBox("提示词")
        prompt_layout = QVBoxLayout()

        # Prompt toolbar
        prompt_toolbar = QHBoxLayout()
        self.save_prompt_btn = QPushButton("⭐ 收藏")
        self.save_prompt_btn.clicked.connect(self.toggle_favorite)
        prompt_toolbar.addWidget(self.save_prompt_btn)

        prompt_toolbar.addWidget(QLabel("快速选择:"))
        self.saved_prompts_combo = QComboBox()
        prompt_toolbar.addWidget(self.saved_prompts_combo)

        delete_prompt_btn = QPushButton("🗑️ 删除")
        delete_prompt_btn.clicked.connect(self.delete_saved_prompt)
        prompt_toolbar.addWidget(delete_prompt_btn)
        prompt_toolbar.addStretch()

        prompt_layout.addLayout(prompt_toolbar)

        # Prompt text area
        self.prompt_text = QTextEdit()
        self.prompt_text.setPlainText("一幅美丽的风景画，有山有湖，日落时分")
        self.prompt_text.textChanged.connect(self.check_favorite_status)
        prompt_layout.addWidget(self.prompt_text)

        # Update list first (before connecting signal to avoid triggering it during init)
        self.update_saved_prompts_list()

        prompt_group.setLayout(prompt_layout)
        content_row.addWidget(prompt_group, 2)

        # Settings section
        settings_group = QGroupBox("生成设置")
        settings_layout = QVBoxLayout()

        settings_layout.addWidget(QLabel("宽高比"))
        self.aspect_ratio = QComboBox()
        self.aspect_ratio.addItems(["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"])
        settings_layout.addWidget(self.aspect_ratio)

        settings_layout.addSpacing(10)

        settings_layout.addWidget(QLabel("图片尺寸"))
        self.image_size = QComboBox()
        self.image_size.addItems(["1K", "2K", "4K"])
        self.image_size.setCurrentIndex(1)  # Default to 2K
        settings_layout.addWidget(self.image_size)

        settings_layout.addStretch()
        settings_group.setLayout(settings_layout)
        content_row.addWidget(settings_group, 1)

        main_layout.addLayout(content_row)

        # Action buttons
        action_layout = QHBoxLayout()
        self.generate_btn = QPushButton("生成图片")
        self.generate_btn.clicked.connect(self.generate_image_async)
        action_layout.addWidget(self.generate_btn)

        self.download_btn = QPushButton("下载图片")
        self.download_btn.clicked.connect(self.download_image)
        self.download_btn.setEnabled(False)
        action_layout.addWidget(self.download_btn)

        self.status_label = QLabel("● 就绪")
        action_layout.addWidget(self.status_label)
        action_layout.addStretch()

        main_layout.addLayout(action_layout)

        # Preview section
        preview_group = QGroupBox("预览")
        preview_layout = QVBoxLayout()

        self.preview_label = QLabel("生成的图片将在这里显示\n双击用系统查看器打开")
        self.preview_label.setAlignment(Qt.AlignCenter)
        self.preview_label.setMinimumHeight(300)
        self.preview_label.setStyleSheet("QLabel { color: #999999; font-size: 10pt; }")
        self.preview_label.mouseDoubleClickEvent = self.open_fullsize_view

        preview_layout.addWidget(self.preview_label)
        preview_group.setLayout(preview_layout)
        main_layout.addWidget(preview_group, 1)

        central_widget.setLayout(main_layout)

        self.check_favorite_status()

    def apply_styles(self):
        """Apply QSS stylesheet"""
        """Apply QSS stylesheet - 最小化自定义样式"""
        self.setStyleSheet("""
                QMainWindow {
                    background-color: #ffffff;
                }
                QGroupBox {
                    font-weight: bold;
                    font-size: 10pt;
                    border: 1px solid #e5e5e5;
                    border-radius: 6px;
                    margin-top: 10px;
                    padding-top: 10px;
                }
                QGroupBox::title {
                    subcontrol-origin: margin;
                    left: 10px;
                    padding: 0 5px;
                }
                QPushButton {
                    padding: 6px 12px;
                    font-size: 9pt;
                    border: 1px solid #cccccc;
                    border-radius: 4px;
                }
                QPushButton:disabled {
                    color: #999999;
                }
                QComboBox {
                    padding: 5px;
                    min-width: 100px;
                }
                QTextEdit {
                    border: 1px solid #e5e5e5;
                    border-radius: 4px;
                    font-size: 10pt;
                }
                QScrollArea {
                    border: none;
                    background-color: #f6f6f6;
                }
                /* 只保留必要的边框和背景，不设置颜色 */
                QLineEdit, QComboBox, QTextEdit {
                    border: 1px solid #cccccc;
                    border-radius: 4px;
                }
                QLineEdit:focus, QComboBox:focus, QTextEdit:focus {
                    border: 1px solid #007AFF;  /* 保持焦点状态 */
                }
            """)

        # Connect signals after all widgets are created
        self.saved_prompts_combo.currentIndexChanged.connect(self.load_saved_prompt)

    def upload_images(self):
        """Upload reference images"""
        files, _ = QFileDialog.getOpenFileNames(
            self,
            "选择参考图片",
            "",
            "图片文件 (*.png *.jpg *.jpeg *.gif *.bmp);;所有文件 (*.*)"
        )

        if files:
            for file_path in files:
                try:
                    self.uploaded_images.append(file_path)
                except Exception as e:
                    QMessageBox.critical(self, "错误", f"无法加载图片: {file_path}\n{str(e)}")

            self.update_image_preview()
            self.image_count_label.setText(f"已选择 {len(self.uploaded_images)} 张")
            self.status_label.setText(f"● 已添加 {len(files)} 张参考图片")
            self.status_label.setStyleSheet("QLabel { color: #34C759; }")

    def update_image_preview(self):
        """Update image preview thumbnails"""
        # Clear existing previews
        while self.img_layout.count() > 1:  # Keep the stretch
            item = self.img_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()

        # Add thumbnails
        for idx, file_path in enumerate(self.uploaded_images):
            try:
                # Load and create thumbnail
                pixmap = QPixmap(file_path)
                pixmap = pixmap.scaled(100, 100, Qt.KeepAspectRatio, Qt.SmoothTransformation)

                # Container
                container = QWidget()
                container_layout = QVBoxLayout()
                container_layout.setContentsMargins(5, 5, 5, 5)

                # Image label
                img_label = QLabel()
                img_label.setPixmap(pixmap)
                img_label.setFixedSize(100, 100)
                img_label.setStyleSheet("QLabel { border: 1px solid #e5e5e5; }")
                container_layout.addWidget(img_label)

                # Info row
                info_layout = QHBoxLayout()
                index_label = QLabel(f"图 {idx + 1}")
                index_label.setStyleSheet("QLabel { font-size: 8pt; color: #666666; }")
                info_layout.addWidget(index_label)

                del_btn = QPushButton("✕")
                del_btn.setFixedSize(20, 20)
                del_btn.setStyleSheet("""
                    QPushButton {
                        background-color: #ff4444;
                        color: white;
                        font-weight: bold;
                        border: none;
                        border-radius: 3px;
                        padding: 0px;
                    }
                    QPushButton:hover {
                        background-color: #FF3B30;
                    }
                """)
                del_btn.clicked.connect(lambda checked, i=idx: self.delete_image(i))
                info_layout.addWidget(del_btn)

                container_layout.addLayout(info_layout)
                container.setLayout(container_layout)

                self.img_layout.insertWidget(self.img_layout.count() - 1, container)
            except Exception as e:
                print(f"Failed to create thumbnail for {file_path}: {e}")

    def delete_image(self, index):
        """Delete an image by index"""
        if 0 <= index < len(self.uploaded_images):
            self.uploaded_images.pop(index)
            self.update_image_preview()
            self.image_count_label.setText(f"已选择 {len(self.uploaded_images)} 张")
            self.status_label.setText("● 已删除图片")
            self.status_label.setStyleSheet("QLabel { color: #FF9500; }")

    def update_saved_prompts_list(self):
        """Update the saved prompts dropdown"""
        # Block signals to avoid triggering load_saved_prompt during update
        self.saved_prompts_combo.blockSignals(True)
        self.saved_prompts_combo.clear()
        if self.saved_prompts:
            display_prompts = [p[:50] + "..." if len(p) > 50 else p for p in self.saved_prompts]
            self.saved_prompts_combo.addItems(display_prompts)
        self.saved_prompts_combo.blockSignals(False)

    def check_favorite_status(self):
        """Check if current prompt is favorited"""
        prompt = self.prompt_text.toPlainText().strip()
        if prompt in self.saved_prompts:
            self.save_prompt_btn.setText("✓ 已收藏")
        else:
            self.save_prompt_btn.setText("⭐ 收藏")

    def toggle_favorite(self):
        """Toggle favorite status of current prompt"""
        prompt = self.prompt_text.toPlainText().strip()
        if not prompt:
            self.status_label.setText("● 提示词不能为空")
            self.status_label.setStyleSheet("QLabel { color: #FF3B30; }")
            return

        if prompt in self.saved_prompts:
            self.saved_prompts.remove(prompt)
            self.save_config()
            self.update_saved_prompts_list()
            self.status_label.setText("● 该提示词已取消收藏")
        else:
            self.saved_prompts.append(prompt)
            self.save_config()
            self.update_saved_prompts_list()
            self.status_label.setText("● 该提示词已收藏")

        self.status_label.setStyleSheet("QLabel { color: #34C759; }")
        self.check_favorite_status()

    def load_saved_prompt(self):
        """Load a saved prompt"""
        index = self.saved_prompts_combo.currentIndex()
        if 0 <= index < len(self.saved_prompts):
            self.prompt_text.setPlainText(self.saved_prompts[index])
            self.status_label.setText("● 已加载提示词")
            self.status_label.setStyleSheet("QLabel { color: #007AFF; }")

    def delete_saved_prompt(self):
        """Delete the currently selected saved prompt"""
        index = self.saved_prompts_combo.currentIndex()

        if index < 0 or index >= len(self.saved_prompts):
            self.status_label.setText("● 请先选择要删除的提示词")
            self.status_label.setStyleSheet("QLabel { color: #FF9500; }")
            return

        self.saved_prompts.pop(index)
        self.save_config()
        self.update_saved_prompts_list()
        self.status_label.setText("● 已删除提示词")
        self.status_label.setStyleSheet("QLabel { color: #34C759; }")

    def generate_image_async(self):
        """Start image generation in a separate thread"""
        # Create and start worker thread
        self.worker = ImageGenerationWorker(
            self.api_key,
            self.prompt_text.toPlainText().strip(),
            self.uploaded_images,
            self.aspect_ratio.currentText(),
            self.image_size.currentText()
        )
        self.worker.finished.connect(self.on_image_generated)
        self.worker.error.connect(self.on_generation_error)
        self.worker.progress.connect(self.update_status)

        self.generate_btn.setEnabled(False)
        self.download_btn.setEnabled(False)
        self.status_label.setText("● 正在生成图片...")
        self.status_label.setStyleSheet("QLabel { color: #FF9500; }")

        self.worker.start()

    def on_image_generated(self, image_bytes):
        """Handle successful image generation"""
        self.generated_image_bytes = image_bytes
        self.display_image()
        self.download_btn.setEnabled(True)
        self.generate_btn.setEnabled(True)
        self.status_label.setText("● 图片生成成功")
        self.status_label.setStyleSheet("QLabel { color: #34C759; }")

    def on_generation_error(self, error_msg):
        """Handle image generation error"""
        QMessageBox.critical(self, "错误", f"生成失败: {error_msg}")
        self.generate_btn.setEnabled(True)
        self.status_label.setText("● 生成失败")
        self.status_label.setStyleSheet("QLabel { color: #FF3B30; }")

    def update_status(self, message):
        """Update status label"""
        self.status_label.setText(f"● {message}")

    def display_image(self):
        """Display generated image in preview"""
        if not self.generated_image_bytes:
            return

        try:
            # Load image from bytes
            pixmap = QPixmap()
            pixmap.loadFromData(self.generated_image_bytes)

            # Scale to fit preview area
            available_width = self.preview_label.width() - 40
            available_height = self.preview_label.height() - 40

            scaled_pixmap = pixmap.scaled(
                available_width, available_height,
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            )

            self.preview_label.setPixmap(scaled_pixmap)
            self.preview_label.setStyleSheet("")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"图片显示失败: {str(e)}")

    def open_fullsize_view(self, event):
        """Open generated image with system default viewer"""
        if not self.generated_image_bytes:
            return

        try:
            # Create temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.png', mode='wb') as tmp_file:
                tmp_file.write(self.generated_image_bytes)
                tmp_path = tmp_file.name

            # Open with system default viewer
            url = QUrl.fromLocalFile(tmp_path)
            QDesktopServices.openUrl(url)

            self.status_label.setText("● 已用系统查看器打开")
            self.status_label.setStyleSheet("QLabel { color: #007AFF; }")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"无法打开系统图片查看器: {str(e)}")

    def download_image(self):
        """Download generated image"""
        if not self.generated_image_bytes:
            QMessageBox.critical(self, "错误", "没有可下载的图片！")
            return

        # Generate default filename
        default_filename = datetime.now().strftime("%Y%m%d%H%M%S.png")

        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "保存图片",
            default_filename,
            "PNG 文件 (*.png);;JPEG 文件 (*.jpg);;所有文件 (*.*)"
        )

        if file_path:
            try:
                with open(file_path, 'wb') as f:
                    f.write(self.generated_image_bytes)

                file_size = len(self.generated_image_bytes)
                QMessageBox.information(self, "成功", f"图片已保存到:\n{file_path}\n\n文件大小: {file_size:,} 字节")
                self.status_label.setText("● 图片已保存")
                self.status_label.setStyleSheet("QLabel { color: #34C759; }")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存失败: {str(e)}")


class ImageGenerationWorker(QThread):
    """Worker thread for image generation"""
    finished = Signal(bytes)
    error = Signal(str)
    progress = Signal(str)

    def __init__(self, api_key, prompt, images, aspect_ratio, image_size):
        super().__init__()
        self.api_key = api_key
        self.prompt = prompt
        self.images = images
        self.aspect_ratio = aspect_ratio
        self.image_size = image_size

    def run(self):
        """Execute image generation in background thread"""
        try:
            if not self.prompt:
                self.error.emit("请输入图片描述！")
                return

            if not self.api_key:
                self.error.emit("未找到API密钥，请在config.json中配置！")
                return

            self.progress.emit("正在连接 Gemini API...")

            client = genai.Client(api_key=self.api_key)

            # Build content parts
            content_parts = [self.prompt]

            # Add reference images
            for img_path in self.images:
                with open(img_path, 'rb') as f:
                    img_data = f.read()

                mime_type = "image/png"
                if img_path.lower().endswith(('.jpg', '.jpeg')):
                    mime_type = "image/jpeg"

                content_parts.append(
                    types.Part.from_bytes(
                        data=img_data,
                        mime_type=mime_type
                    )
                )

            self.progress.emit("正在生成图片...")

            # Generation config
            config = types.GenerateContentConfig(
                response_modalities=["IMAGE"],
                image_config=types.ImageConfig(
                    aspect_ratio=self.aspect_ratio,
                    image_size=self.image_size
                )
            )

            # Generate
            response = client.models.generate_content(
                model="gemini-3-pro-image-preview",
                contents=content_parts,
                config=config
            )

            # Extract image
            for part in response.parts:
                if hasattr(part, 'inline_data') and part.inline_data:
                    if isinstance(part.inline_data.data, bytes):
                        image_bytes = part.inline_data.data
                    else:
                        image_bytes = base64.b64decode(part.inline_data.data)

                    self.finished.emit(image_bytes)
                    return

            self.error.emit("响应中没有图片数据")

        except Exception as e:
            self.error.emit(str(e))


def main():
    """Main application entry point"""
    # Load config for database info
    config_dir = Path('.')
    if getattr(sys, 'frozen', False):
        system = platform.system()
        if system == 'Darwin':
            config_dir = Path.home() / 'Library' / 'Application Support' / 'ZB100ImageGenerator'
        elif system == 'Windows':
            config_dir = Path(os.getenv('APPDATA', Path.home())) / 'ZB100ImageGenerator'
        else:
            config_dir = Path.home() / '.config' / 'zb100imagegenerator'

    config_dir.mkdir(parents=True, exist_ok=True)
    config_path = config_dir / 'config.json'

    db_config = None
    last_user = ""
    saved_password_hash = ""

    if config_path.exists():
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
                db_config = config.get("db_config")
                last_user = config.get("last_user", "")
                saved_password_hash = config.get("saved_password_hash", "")
        except Exception as e:
            print(f"Failed to load config: {e}")

    # Create QApplication
    app = QApplication(sys.argv)

    # Check database config
    if not db_config:
        QMessageBox.critical(None, "配置错误",
                           f"未找到数据库配置\n配置文件: {config_path}\n\n"
                           "请确保 config.json 包含 db_config 字段")
        return

    # Show login dialog
    login_dialog = LoginDialog(db_config, last_user, saved_password_hash)

    if login_dialog.exec() == QDialog.Accepted:
        # Login successful
        authenticated_user = login_dialog.authenticated_user
        remember_user = login_dialog.get_remember_user()
        remember_password = login_dialog.get_remember_password()
        password_hash = login_dialog.get_password_hash()

        # Save/clear credentials
        if config_path.exists():
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)

                if remember_user:
                    config["last_user"] = authenticated_user
                else:
                    config["last_user"] = ""

                if remember_password:
                    config["saved_password_hash"] = password_hash
                else:
                    config["saved_password_hash"] = ""

                with open(config_path, 'w', encoding='utf-8') as f:
                    json.dump(config, f, indent=2, ensure_ascii=False)
            except Exception as e:
                print(f"Failed to save config: {e}")

        # Show main window
        main_window = ImageGeneratorWindow()
        main_window.show()

        sys.exit(app.exec())
    else:
        # Login cancelled or failed
        sys.exit(0)


if __name__ == "__main__":
    main()