import sys
import subprocess
import ctypes
from PyQt6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QPushButton, QComboBox,
    QLabel, QLineEdit, QFileDialog
)
from PyQt6.QtCore import Qt

def get_adb_mounts():
    """Return a list of device mounts with SD card and internal storage at top."""
    try:
        # Check if any devices are connected
        devices_result = subprocess.run(
            ["adb", "devices"],
            capture_output=True, text=True, check=True
        )
        lines = devices_result.stdout.strip().splitlines()[1:]  # skip header
        connected = [line for line in lines if "device" in line]
        if not connected:
            return ["No device connected or unauthorized"]

        # Get mounts from the first connected device
        mount_result = subprocess.run(
            ["adb", "shell", "mount"],
            capture_output=True, text=True, check=True
        )
        
        mounts = []
        sd_card_mount = None
        internal_storage = None
        other_mounts = []
        
        for line in mount_result.stdout.splitlines():
            parts = line.split()
            if len(parts) >= 3:
                device, mount_point = parts[0], parts[2]
                mount_entry = f"{device} -> {mount_point}"
                
                # Check for SD card patterns (xxxx-xxxx format in device or common SD mount points)
                if any(pattern in mount_point.lower() for pattern in ['/mnt/media_rw/', '/storage/', '/sdcard', '/external_sd']) \
                   or '-' in device and len(device.split('-')[0]) == 4:
                    sd_card_mount = mount_entry
                # Check for internal storage
                elif '/storage/emulated/0' in mount_point or '/data/media' in mount_point:
                    internal_storage = mount_entry
                else:
                    other_mounts.append(mount_entry)
        
        # Add SD card at top if found (try to get specific format)
        if sd_card_mount:
            # Try to extract the xxxx-xxxx pattern for better display
            device_part = sd_card_mount.split('->')[0].strip()
            mount_part = sd_card_mount.split('->')[1].strip()
            if '-' in device_part and len(device_part) >= 9:
                mounts.append(f"{device_part} -> {mount_part} (micro SD card)")
            else:
                mounts.append(f"{sd_card_mount} (micro SD card)")
        else:
            # Try alternative method to find SD card
            try:
                sd_result = subprocess.run(
                    ["adb", "shell", "ls", "/storage/"],
                    capture_output=True, text=True, check=True
                )
                for item in sd_result.stdout.split():
                    if item and item not in ['emulated', 'self']:
                        mounts.append(f"SD Card -> /storage/{item}")
                        break
            except:
                pass
        
        # Add internal storage second
        if internal_storage:
            mounts.append(f"{internal_storage} (Internal Storage)")
        else:
            # Try to find internal storage via other means
            mounts.append("/data/media/0 -> /storage/emulated/0 (Internal Storage)")
        
        # Add all other mounts
        mounts.extend(other_mounts)
        
        if not mounts:
            mounts.append("No mounts found")
        return mounts
    except subprocess.CalledProcessError as e:
        return [f"ADB error: {e.stderr.strip() or e}"]
    except FileNotFoundError:
        return ["ADB not found. Make sure it is installed and in PATH."]

def get_volume_id(drive_letter):
    """Return the volume serial number for a given drive letter."""
    volume_name_buf = ctypes.create_unicode_buffer(1024)
    file_system_name_buf = ctypes.create_unicode_buffer(1024)
    serial_number = ctypes.c_uint()
    max_component_length = ctypes.c_uint()
    file_system_flags = ctypes.c_uint()
    ret = ctypes.windll.kernel32.GetVolumeInformationW(
        ctypes.c_wchar_p(drive_letter + "\\"),
        volume_name_buf,
        ctypes.sizeof(volume_name_buf),
        ctypes.byref(serial_number),
        ctypes.byref(max_component_length),
        ctypes.byref(file_system_flags),
        file_system_name_buf,
        ctypes.sizeof(file_system_name_buf)
    )
    if ret:
        return f"{serial_number.value:X}"
    return "Unavailable"

class AdbStorageApp(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("ADB Storage & Drive Info")
        self.setGeometry(300, 300, 450, 220)
        layout = QVBoxLayout()

        # ADB Mounts Picker
        self.mount_label = QLabel("Device Storage Mounts:")
        self.mount_combo = QComboBox()
        self.refresh_btn = QPushButton("Refresh Mounts")
        self.refresh_btn.clicked.connect(self.load_mounts)

        # Drive input
        self.drive_label = QLabel("Drive:")
        self.drive_input = QLineEdit("F:")
        self.browse_btn = QPushButton("Browse")
        self.browse_btn.clicked.connect(self.browse_drive)

        # Volume ID display
        self.volume_label = QLabel("Volume ID:")
        self.volume_display = QLabel("")

        # Connect drive input change
        self.drive_input.textChanged.connect(self.update_volume_id)

        layout.addWidget(self.mount_label)
        layout.addWidget(self.mount_combo)
        layout.addWidget(self.refresh_btn)
        layout.addWidget(self.drive_label)
        layout.addWidget(self.drive_input)
        layout.addWidget(self.browse_btn)
        layout.addWidget(self.volume_label)
        layout.addWidget(self.volume_display)

        self.setLayout(layout)
        self.load_mounts()
        self.update_volume_id()

    def load_mounts(self):
        self.mount_combo.clear()
        mounts = get_adb_mounts()
        self.mount_combo.addItems(mounts)

    def browse_drive(self):
        folder = QFileDialog.getExistingDirectory(self, "Select Drive")
        if folder:
            drive = folder.split(":")[0] + ":"
            self.drive_input.setText(drive)

    def update_volume_id(self):
        drive = self.drive_input.text().strip()
        if drive and len(drive) == 2 and drive[1] == ":":
            vid = get_volume_id(drive)
            self.volume_display.setText(vid)
        else:
            self.volume_display.setText("Invalid Drive")

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