import sys
import os
import random
import math
import time
import secrets

from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                               QHBoxLayout, QLabel, QLineEdit, QPushButton, 
                               QTextEdit, QSpinBox, QComboBox, QGroupBox, 
                               QRadioButton, QButtonGroup, QFileDialog, QFrame, QMessageBox)
from PySide6.QtCore import QThread, Signal, Qt

# ==============================================================================
# Z-Image-Turbo GUI Wrapper
#
# This GUI is a wrapper for the Z-Image-MPS project.
# Original Core Logic Copyright (c) 2025 Ivan Fioravanti (MIT License)
# Source: https://github.com/ivanfioravanti/z-image-mps
# ==============================================================================

try:
    import torch
    from diffusers import ZImagePipeline # Pipeline specified in the source code
except ImportError:
    print("Error: torch or diffusers library is not installed.")
    sys.exit(1)

# ==========================================
# 1. Background Worker (Model Resident & Generation)
# ==========================================
class ModelWorker(QThread):
    status_signal = Signal(str)    # Status message
    result_signal = Signal(str)    # Result message
    error_signal = Signal(str)     # Error message

    def __init__(self):
        super().__init__()
        self.pipe = None
        self.is_loaded = False
        self.request_queue = None

    def pick_device(self):
        """Porting pick_device logic from the provided source code"""
        if torch.backends.mps.is_available():
            return "mps", torch.bfloat16
        if torch.cuda.is_available():
            return "cuda", torch.bfloat16
        return "cpu", torch.float32

    def run(self):
        """Loads the model upon app start and waits."""
        self.status_signal.emit("Loading model... (Tongyi-MAI/Z-Image-Turbo)")
        
        try:
            # 1. Select Device and Dtype (Source Code Logic)
            device, dtype = self.pick_device()
            self.status_signal.emit(f"Device set: {device} ({dtype})")

            # 2. Load Pipeline (Source Code Logic)
            # Set load_kwargs
            load_kwargs = {"low_cpu_mem_usage": False}
            load_kwargs["torch_dtype"] = dtype

            # Load Model
            self.pipe = ZImagePipeline.from_pretrained(
                "Tongyi-MAI/Z-Image-Turbo", 
                **load_kwargs
            )

            # 3. VAE Setup (Critical logic from source: Prevent NaNs)
            if dtype != torch.float32 and hasattr(self.pipe, "vae"):
                self.pipe.vae.to(dtype=torch.float32)
                self.pipe.vae.config.force_upcast = True
                self.status_signal.emit("VAE setup complete (Float32 Upcast)")

            # 4. Attention Backend (Default SDPA applied)
            try:
                self.pipe.transformer.set_attention_backend("sdpa")
            except Exception:
                pass # Ignore if not supported

            # 5. Move to Device
            self.pipe.to(device)

            self.is_loaded = True
            self.status_signal.emit("✅ Model loaded! Ready to generate.")

        except Exception as e:
            self.error_signal.emit(f"Model loading failed: {str(e)}")
            return

        # === Waiting Loop ===
        while True:
            if self.request_queue is None:
                self.msleep(50) # Wait
                continue
            
            # Process Request
            params = self.request_queue
            self.request_queue = None
            self.generate_image(params)

    def generate_image(self, params):
        if not self.pipe:
            return

        try:
            self.status_signal.emit(f"Generating... [Seed: {params['seed']}]")
            
            # Create Generator (Source logic: MPS uses CPU Generator)
            device = self.pipe.device
            seed = int(params['seed'])
            
            generator_device = "cpu" if device.type == "mps" else device
            generator = torch.Generator(device=generator_device).manual_seed(seed)

            # Run Inference
            with torch.inference_mode():
                result = self.pipe(
                    prompt=params['prompt'],
                    negative_prompt=params['negative_prompt'],
                    height=int(params['height']),
                    width=int(params['width']),
                    num_inference_steps=int(params['steps']),
                    guidance_scale=float(params['guidance']), # Usually 0.0
                    generator=generator
                )

            image = result.images[0]
            
            # Save
            output_dir = params['output_dir']
            if not output_dir:
                output_dir = os.path.join(os.getcwd(), "output")
            os.makedirs(output_dir, exist_ok=True)
            
            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") if 'datetime' in globals() else str(int(time.time()))
            filename = f"z-image-{timestamp}-{seed}.png"
            save_path = os.path.join(output_dir, filename)
            
            image.save(save_path)
            
            self.result_signal.emit(f"Saved: {save_path}")
            self.status_signal.emit("Ready...")

        except Exception as e:
            self.error_signal.emit(f"Generation Error: {str(e)}")

    def request_generation(self, params):
        if not self.is_loaded:
            self.error_signal.emit("Model is loading. Please wait.")
            return
        self.request_queue = params


# ==========================================
# 2. Main GUI Class (DINKI Spec Applied)
# ==========================================
class ZImageGUIFast(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Z-Image-Turbo GUI (ANG!)")
        self.resize(720, 820)

        # Initial Resolution
        self.calc_width = 896
        self.calc_height = 1152

        # Start Worker Thread
        self.worker = ModelWorker()
        self.worker.status_signal.connect(self.update_log)
        self.worker.result_signal.connect(self.generation_finished)
        self.worker.error_signal.connect(self.show_error)
        self.worker.start()

        self.init_ui()

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        # 1. Prompts
        prompt_group = QGroupBox("Prompts")
        prompt_layout = QVBoxLayout()
        
        prompt_layout.addWidget(QLabel("Positive Prompt:"))
        self.positive_prompt = QTextEdit()
        # Apply default prompt from source code
        DEFAULT_PROMPT = "On the desk, there is an iMac along with a keyboard, a mouse, and an iPhone."
        self.positive_prompt.setText(DEFAULT_PROMPT)
        self.positive_prompt.setFixedHeight(65)
        prompt_layout.addWidget(self.positive_prompt)

        prompt_layout.addWidget(QLabel("Negative Prompt:"))
        self.negative_prompt = QTextEdit()
        self.negative_prompt.setFixedHeight(65)
        prompt_layout.addWidget(self.negative_prompt)
        prompt_group.setLayout(prompt_layout)
        main_layout.addWidget(prompt_group)

        # 2. Settings (Steps, Seed, Guidance)
        settings_group = QGroupBox("Settings")
        settings_layout = QVBoxLayout()
        
        # Row 1: Steps & Seed
        row1 = QHBoxLayout()
        
        row1.addWidget(QLabel("Steps:"))
        self.step_spin = QSpinBox()
        self.step_spin.setRange(1, 100)
        self.step_spin.setValue(9) # Default from source
        self.step_spin.setFixedWidth(70)
        row1.addWidget(self.step_spin)
        
        row1.addSpacing(20)
        
        # Seed Control
        self.radio_random = QRadioButton("Random")
        self.radio_fixed = QRadioButton("Fixed")
        self.radio_random.setChecked(True)
        btn_group = QButtonGroup(self)
        btn_group.addButton(self.radio_random)
        btn_group.addButton(self.radio_fixed)
        
        row1.addWidget(self.radio_random)
        row1.addWidget(self.radio_fixed)
        
        self.seed_input = QLineEdit()
        self.seed_input.setPlaceholderText("Seed")
        self.seed_input.setEnabled(False)
        self.seed_input.setFixedWidth(120)
        row1.addWidget(self.seed_input)
        
        self.radio_random.toggled.connect(lambda: self.seed_input.setEnabled(self.radio_fixed.isChecked()))
        
        settings_layout.addLayout(row1)

        # DINKI Specs
        h_line = QFrame()
        h_line.setFrameShape(QFrame.HLine)
        settings_layout.addWidget(h_line)

        row2 = QHBoxLayout()
        
        # MP
        self.mp_combo = QComboBox()
        self.mp_combo.addItems(["1MP", "2MP", "3MP", "4MP"])
        self.mp_combo.currentIndexChanged.connect(self.calculate_dinki_resolution)
        row2.addWidget(QLabel("MP:"))
        row2.addWidget(self.mp_combo)
        
        # Ratio
        self.ratio_combo = QComboBox()
        self.ratio_list = [
            "Photo 3:4", "Photo 3.5:5", "Photo 4:6", 
            "Photo 5:7", "Photo 6:8", "Photo 8:10", 
            "Photo 10:13", "Photo 10:15", "Photo 11:14",
            "35mm Academy 1.37:1", "35mm Flat 1.85:1", 
            "35mm Scope (Anamorphic) 2.39:1", "70mm Todd-AO 2.20:1", 
            "IMAX 70mm 1.43:1", "Super 35 1.85:1", 
            "Super 35 2.39:1", "Super 16 1.66:1", "Super 16 1.78:1",
        ]
        self.ratio_combo.addItems(self.ratio_list)
        self.ratio_combo.currentIndexChanged.connect(self.calculate_dinki_resolution)
        row2.addWidget(QLabel("Ratio:"))
        row2.addWidget(self.ratio_combo)

        # Orientation
        self.orient_combo = QComboBox()
        self.orient_combo.addItems(["Portrait", "Landscape"])
        self.orient_combo.currentIndexChanged.connect(self.calculate_dinki_resolution)
        row2.addWidget(QLabel("Orient:"))
        row2.addWidget(self.orient_combo)
        
        settings_layout.addLayout(row2)

        self.res_label = QLabel("896 x 1152")
        self.res_label.setStyleSheet("font-weight: bold; color: #4CAF50;")
        self.res_label.setAlignment(Qt.AlignCenter)
        settings_layout.addWidget(self.res_label)
        
        settings_group.setLayout(settings_layout)
        main_layout.addWidget(settings_group)

        # 3. Output
        out_layout = QHBoxLayout()
        self.output_path = QLineEdit()
        self.output_path.setPlaceholderText("Output Folder (Default: ./output)")
        btn_browse = QPushButton("Select Folder")
        btn_browse.clicked.connect(self.browse_folder)
        out_layout.addWidget(self.output_path)
        out_layout.addWidget(btn_browse)
        main_layout.addLayout(out_layout)

        # 4. Action & Log
        self.run_btn = QPushButton("Generate")
        self.run_btn.setMinimumHeight(60)
        self.run_btn.setStyleSheet("font-size: 18px; font-weight: bold; background-color: #2196F3; color: white;")
        self.run_btn.clicked.connect(self.start_generation)
        main_layout.addWidget(self.run_btn)

        self.log_view = QTextEdit()
        self.log_view.setReadOnly(True)
        self.log_view.setStyleSheet("background-color: #1e1e1e; color: #00ff00; font-family: Consolas; font-size: 12px;")
        main_layout.addWidget(self.log_view)

        # Init Calc
        self.calculate_dinki_resolution()

    def calculate_dinki_resolution(self):
        megapixels = self.mp_combo.currentText()
        aspect_ratio_str = self.ratio_combo.currentText()
        orientation = self.orient_combo.currentText()

        mp_multiplier = int(megapixels.replace("MP", ""))
        target_area = 1024 * 1024 * mp_multiplier

        ratio_part_str = aspect_ratio_str.split(" ")[-1]
        ratio_parts = ratio_part_str.split(":")
        w_ratio = float(ratio_parts[0])
        h_ratio = float(ratio_parts[1])
        target_ratio = w_ratio / h_ratio

        height_val = math.sqrt(target_area / target_ratio)
        width_val = height_val * target_ratio

        # [Important] 16-multiple adjustment (Prevent model errors)
        width = round(width_val / 16) * 16
        height = round(height_val / 16) * 16

        is_portrait = "Portrait" in orientation
        if is_portrait:
            if width > height: width, height = height, width
        else:
            if width < height: width, height = height, width

        self.calc_width = width
        self.calc_height = height
        self.res_label.setText(f"{width} x {height}")

    def browse_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "Select Output Folder")
        if folder: self.output_path.setText(folder)

    def start_generation(self):
        # Seed Processing
        final_seed = "0"
        if self.radio_random.isChecked():
            # Use secrets module for strong random number generation
            random_num = secrets.randbits(63) 
            final_seed = str(random_num)
            self.seed_input.setText(final_seed)
        else:
            final_seed = self.seed_input.text().strip() or "0"

        params = {
            "prompt": self.positive_prompt.toPlainText(),
            "negative_prompt": self.negative_prompt.toPlainText(),
            "steps": self.step_spin.value(),
            "width": self.calc_width,
            "height": self.calc_height,
            "seed": final_seed,
            "guidance": 0.0, # Turbo models usually recommend 0.0 (Based on source)
            "output_dir": self.output_path.text()
        }
        
        self.run_btn.setEnabled(False)
        self.worker.request_generation(params)

    def update_log(self, msg):
        self.log_view.append(msg)
        self.log_view.verticalScrollBar().setValue(self.log_view.verticalScrollBar().maximum())

    def generation_finished(self, msg):
        self.update_log(msg)
        self.run_btn.setEnabled(True)

    def show_error(self, msg):
        self.log_view.append(f"[ERROR] {msg}")
        self.run_btn.setEnabled(True)
        QMessageBox.critical(self, "Error", msg)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ZImageGUIFast()
    window.show()
    sys.exit(app.exec())