import sys
import os
import time
import ctypes

if sys.platform == 'win32':
    kernel32 = ctypes.WinDLL('kernel32')
    user32 = ctypes.WinDLL('user32')
    hWnd = kernel32.GetConsoleWindow()
    if hWnd:
        user32.ShowWindow(hWnd, 0)
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget,
    QVBoxLayout, QPushButton, QLabel, QMessageBox, QDialog
)
from PyQt6.QtGui import QPixmap, QMovie, QIcon
from PyQt6.QtCore import QThread, pyqtSignal, QTimer, Qt, QSize
from pymobiledevice3.lockdown import create_using_usbmux
from pymobiledevice3.services.afc import AfcService
from pymobiledevice3.services.diagnostics import DiagnosticsService

SUPPORTED_DEVICES = {
    'iPhone4,1',
    'iPad2,1', 'iPad2,2', 'iPad2,3', 'iPad2,4',
    'iPad2,5', 'iPad2,6', 'iPad2,7',
    'iPad3,1', 'iPad3,2', 'iPad3,3',
    'iPod5,1'
}

SUPPORTED_VERSIONS = {'8.4.1', '9.3.5', '9.3.6'}

def resource_path(relative_path):
    try:
        base_path = sys._MEIPASS
    except Exception:
        base_path = os.path.abspath(".")

    return os.path.join(base_path, relative_path)

class WorkerThread(QThread):
    finished = pyqtSignal(str)
    error = pyqtSignal(str)
    status_update = pyqtSignal(str)

    def wait_for_device(self):
        self.status_update.emit("Waiting for device \nto return...")
        while True:
            try:
                lockdown = create_using_usbmux()
                if lockdown.get_value(key='ProductType'):
                    return lockdown
            except Exception:
                time.sleep(1)

    def run(self):
        try:
            lockdown = create_using_usbmux()
            
            activation_state = lockdown.get_value(key='ActivationState')
            if activation_state == 'Activated':
                self.finished.emit("Already activated")
                return

            product = lockdown.get_value(key='ProductType')
            if product not in SUPPORTED_DEVICES:
                self.finished.emit(f"Not A5")
                return

            version = lockdown.get_value(key='ProductVersion')
            if version not in SUPPORTED_VERSIONS:
                self.finished.emit(f"Not supported.")
                return

            self.status_update.emit("Activating...")
            
            payload_path = resource_path('payload')

            with AfcService(lockdown=lockdown) as afc:
                afc.set_file_contents(
                    'Downloads/downloads.28.sqlitedb',
                    open(payload_path, 'rb').read()
                )
            
            time.sleep(5) # may improve initial attempt success rate?

            DiagnosticsService(lockdown=lockdown).restart()
            time.sleep(10)
            
            lockdown = self.wait_for_device()

            diag = DiagnosticsService(lockdown=lockdown)
            
            gestalt = diag.mobilegestalt(keys=['ShouldHactivate'])
            should_hactivate = gestalt.get('ShouldHactivate')

            if should_hactivate is False:
                for i in range(5):
                    self.status_update.emit(f"Fixing activation\n(Attempt {i+1}/5)...")
                    
                    with AfcService(lockdown=lockdown) as afc:
                        afc.set_file_contents(
                            'Downloads/downloads.28.sqlitedb',
                            open(payload_path, 'rb').read()
                        )
                    
                    time.sleep(5+i*5)

                    diag.restart()
                    time.sleep(10) 
                    
                    lockdown = self.wait_for_device()
                    
                    diag = DiagnosticsService(lockdown=lockdown)
                    gestalt = diag.mobilegestalt(keys=['ShouldHactivate'])
                    should_hactivate = gestalt.get('ShouldHactivate')
                    
                    if should_hactivate is not False:
                        break
                
                if should_hactivate is False:
                    self.error.emit("Activation failed. Make sure your device is connected to Wi-Fi.")
                    return
                
            DiagnosticsService(lockdown=lockdown).restart()
            self.finished.emit("Done!")
            
        except Exception as e:
            self.error.emit(repr(e))


class CustomAlertBox(QDialog):
    def __init__(self, title, message, parent=None):
        super().__init__(parent)
        self.setWindowFlags(Qt.WindowType.FramelessWindowHint | Qt.WindowType.Dialog)
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        self.setFixedSize(400, 200)

        layout = QVBoxLayout()
        self.setLayout(layout)

        container = QWidget()
        container.setObjectName("AlertContainer")
        container.setStyleSheet("""
            QWidget#AlertContainer {
                background-color: rgba(30, 31, 33, 230);
                border-radius: 20px;
                border: 2px solid rgba(255, 255, 255, 50);
            }
            QLabel {
                color: white;
                font-family: 'Segoe UI Light', 'Helvetica Neue', 'Arial';
                font-weight: bold;
                font-size: 11pt;
            }
            QPushButton {
                background-color: rgba(49, 98, 222, 190);
                color: white;
                border: 1px solid rgba(49, 98, 222, 100);
                border-radius: 10px;
                font-family: 'Segoe UI Light', 'Helvetica Neue', 'Arial';
                font-size: 11pt;
                font-weight: bold;
                padding: 10px;
                min-width: 200px;
            }
            QPushButton:hover {
                background-color: rgba(0, 120, 215, 200);
                border: 1px solid rgba(255, 255, 255, 180);
            }
            QPushButton:pressed {
                background-color: rgba(0, 100, 180, 200);
            }
        """)
        
        container_layout = QVBoxLayout(container)
        container_layout.setContentsMargins(20, 20, 20, 20)
        
        # Title
        lbl_title = QLabel(title)
        lbl_title.setStyleSheet("font-size: 14pt; margin-bottom: 10px;")
        lbl_title.setAlignment(Qt.AlignmentFlag.AlignCenter)
        container_layout.addWidget(lbl_title)
        
        # Message
        lbl_message = QLabel(message)
        lbl_message.setWordWrap(True)
        lbl_message.setAlignment(Qt.AlignmentFlag.AlignCenter)
        container_layout.addWidget(lbl_message)
        
        container_layout.addStretch()
        
        # Button
        btn_ok = QPushButton("OK")
        btn_ok.setCursor(Qt.CursorShape.PointingHandCursor)
        btn_ok.clicked.connect(self.accept)
        container_layout.addWidget(btn_ok, alignment=Qt.AlignmentFlag.AlignCenter)
        
        layout.addWidget(container)

    def mousePressEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            self.drag_pos = event.globalPosition().toPoint() - self.frameGeometry().topLeft()
            event.accept()

    def mouseMoveEvent(self, event):
        if event.buttons() & Qt.MouseButton.LeftButton:
            self.move(event.globalPosition().toPoint() - self.drag_pos)
            event.accept()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Bookra1n A5 Afterlife")
        self.setWindowIcon(QIcon(resource_path('img/icon.ico')))
        self.setFixedSize(450, 550)
        
        # Remove window bar (make frameless)
        self.setWindowFlags(Qt.WindowType.FramelessWindowHint)
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)

        layout = QVBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        container = QWidget()
        container.setObjectName("Container")
        container.setLayout(layout)
        self.setCentralWidget(container)

        # Set background wallpaper
        wallpaper_path = resource_path('img/wall.jpg').replace('\\', '/')
        self.setStyleSheet(f"""
            #Container {{
                background-image: url({wallpaper_path});
                background-position: center;
                background-repeat: no-repeat;
                border-radius: 20px;
            }}
        """)

        # Logo
        self.logo_container = QWidget()
        self.logo_container.setObjectName("LogoContainer")
        self.logo_container.setStyleSheet("""
            #LogoContainer {
                background-color: rgba(0, 0, 0, 100);
                border-top-left-radius: 20px;
                border-top-right-radius: 20px;
                border-bottom-left-radius: 0px;
                border-bottom-right-radius: 0px;
            }
        """)
        self.logo_layout = QVBoxLayout(self.logo_container)
        self.logo_layout.setContentsMargins(10, 10, 10, 10)
        
        self.logo_label = QLabel()
        pixmap = QPixmap(resource_path('img/a5afterlife.png'))
        self.logo_label.setPixmap(pixmap.scaled(400, 200, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation))
        self.logo_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.logo_layout.addWidget(self.logo_label)
        layout.addWidget(self.logo_container)
        
        # Bottom Container (Device + Button)
        self.bottom_container = QWidget()
        self.bottom_layout = QVBoxLayout(self.bottom_container)
        self.bottom_layout.setContentsMargins(15, 0, 15, 10) # Bottom 30 to move button up
        
        # Device Image
        self.device_image_label = QLabel()
        self.pixmap_disconnected = QPixmap(resource_path('img/disconnected.png')).scaled(350, 350, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation)
        self.pixmap_connected = QPixmap(resource_path('img/connected.png')).scaled(350, 350, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation)
        
        self.device_image_label.setPixmap(self.pixmap_disconnected)
        self.device_image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.bottom_layout.addWidget(self.device_image_label)
        
        # Status Label (Overlaid involved device image)
        self.device_layout = QVBoxLayout(self.device_image_label)
        self.device_layout.setContentsMargins(0, 0, 0, 85)
        
        self.device_layout.addStretch()
        
        self.rocket_label = QLabel()
        self.rocket_movie = QMovie(resource_path('img/rocket.gif'))
        self.rocket_movie.setScaledSize(QSize(150, 150))
        self.rocket_label.setMovie(self.rocket_movie)
        self.rocket_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.rocket_label.setStyleSheet("padding-top: 10px;")
        self.rocket_label.hide()
        self.device_layout.addWidget(self.rocket_label)
        
        # self.device_layout.addStretch()
        
        self.status_label = QLabel("No device connected")
        self.status_label.setStyleSheet("color: white; font-weight: bold; font-family: 'Segoe UI Light', 'Helvetica Neue', 'Arial'; font-size: 10pt; background: transparent;")
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.device_layout.addWidget(self.status_label)
        
        self.btn_activate = QPushButton("Activate Device")
        self.btn_activate.setCursor(Qt.CursorShape.PointingHandCursor)
        self.btn_activate.setStyleSheet("""
            QPushButton {
                background-color: rgba(49, 98, 222, 190);
                color: white;
                border: 1px solid rgba(49, 98, 222, 100);
                border-radius: 10px;
                font-family: 'Segoe UI Light', 'Helvetica Neue', 'Arial';
                font-size: 11pt;
                font-weight: bold;
                padding: 10px;
            }
            QPushButton:hover {
                background-color: rgba(0, 120, 215, 200);
                border: 1px solid rgba(255, 255, 255, 180);
            }
            QPushButton:pressed {
                background-color: rgba(0, 100, 180, 200);
            }
            QPushButton:disabled {
                background-color: rgba(100, 100, 100, 100);
                color: rgba(200, 200, 200, 150);
            }
        """)
        self.bottom_layout.addWidget(self.btn_activate)
        
        layout.addWidget(self.bottom_container)

        self.btn_activate.clicked.connect(self.activate_device)

        # Close Button
        self.btn_close = QPushButton("X", self)
        self.btn_close.setGeometry(410, 10, 30, 30) # Top right, window width is 450
        self.btn_close.setCursor(Qt.CursorShape.PointingHandCursor)
        self.btn_close.setStyleSheet("""
            QPushButton {
                background-color: rgba(255, 255, 255, 30);
                color: white;
                border: 1px solid rgba(255, 255, 255, 50);
                border-radius: 15px;
                font-family: 'Segoe UI Light', 'Helvetica Neue', 'Arial';
                font-weight: bold;
                font-size: 12pt;
            }
            QPushButton:hover {
                background-color: rgba(232, 17, 35, 200);
                border: 1px solid rgba(255, 255, 255, 180);
            }
            QPushButton:pressed {
                background-color: rgba(180, 0, 0, 200);
            }
        """)
        self.btn_close.clicked.connect(QApplication.instance().quit)

        self.timer = QTimer(self)
        self.timer.timeout.connect(self.check_device)
        self.timer.start(1000)

    def mousePressEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            self.drag_pos = event.globalPosition().toPoint() - self.frameGeometry().topLeft()
            event.accept()

    def mouseMoveEvent(self, event):
        if event.buttons() & Qt.MouseButton.LeftButton:
            self.move(event.globalPosition().toPoint() - self.drag_pos)
            event.accept()
        

    def check_device(self):
        try:
            lockdown = create_using_usbmux()
            product = lockdown.get_value(key='ProductType')
            version = lockdown.get_value(key='ProductVersion')

            if product in SUPPORTED_DEVICES:
                if version in SUPPORTED_VERSIONS:
                    self.status_label.setText(f"Connected:\n{product} ({version})")
                    self.btn_activate.setText("Activate Device")
                    self.btn_activate.setEnabled(True)
                else:
                    self.status_label.setText(f"Unsupported iOS:\n{version}")
                    self.btn_activate.setText("t.me/Bookra1n")
                    self.btn_activate.setEnabled(False)
            else:
                self.status_label.setText(f"Unsupported Device:\n{product}")
                self.btn_activate.setText("t.me/Bookra1n")
                self.btn_activate.setEnabled(False)
            
            self.device_image_label.setPixmap(self.pixmap_connected)

        except Exception:
            self.status_label.setText("No device connected")
            self.btn_activate.setText("t.me/Bookra1n")
            self.btn_activate.setEnabled(False)
            self.device_image_label.setPixmap(self.pixmap_disconnected)

    def activate_device(self):
        self.show_custom_alert("Info", "The activation will begin now. Please MAKE SURE your device is connected to Wi-Fi.")
        
        self.timer.stop()
        self.btn_activate.setText("t.me/Bookra1n")
        self.btn_activate.setEnabled(False)
        
        self.rocket_label.show()
        self.rocket_movie.start()

        self.thread = WorkerThread()
        self.thread.status_update.connect(self.update_status)
        self.thread.finished.connect(self.on_success)
        self.thread.error.connect(self.on_error)
        self.thread.start()

    def update_status(self, message):
        self.status_label.setText(message)

    def on_success(self, message):
        self.rocket_movie.stop()
        self.rocket_label.hide()
        self.status_label.setText(message)
        self.show_custom_alert("Success!", message)
        self.btn_activate.setEnabled(True)
        self.timer.start(1000)

    def on_error(self, message):
        self.rocket_movie.stop()
        self.rocket_label.hide()
        self.show_custom_alert("Error", message)
        self.status_label.setText("Error.")
        self.btn_activate.setEnabled(True)
        self.timer.start(1000)

    def show_custom_alert(self, title, message):
        alert = CustomAlertBox(title, message, self)
        alert.exec()


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