import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QTableWidget, QTableWidgetItem, QPushButton, QRadioButton, QMessageBox
from subprocess import Popen, PIPE

class DiskpartGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Diskpart GUI - 부팅 디스크 설정")
        self.setGeometry(100, 100, 600, 400)

        # 메인 위젯과 레이아웃
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)

        # 디스크 목록 테이블
        self.disk_table = QTableWidget(0, 3)  # 행은 동적, 열은 3개(Disk, Size, Status)
        self.disk_table.setHorizontalHeaderLabels(["Disk", "Size", "Status"])
        layout.addWidget(self.disk_table)

        # 디스크 새로고침 버튼
        refresh_btn = QPushButton("디스크 목록 새로고침")
        refresh_btn.clicked.connect(self.refresh_disks)
        layout.addWidget(refresh_btn)

        # 부팅 방식 선택
        self.mbr_radio = QRadioButton("BIOS (MBR)")
        self.gpt_radio = QRadioButton("UEFI (GPT)")
        self.mbr_radio.setChecked(True)  # 기본값
        layout.addWidget(self.mbr_radio)
        layout.addWidget(self.gpt_radio)

        # 실행 버튼
        apply_btn = QPushButton("선택한 디스크 준비")
        apply_btn.clicked.connect(self.prepare_disk)
        layout.addWidget(apply_btn)

    def run_diskpart(self, commands):
        process = Popen(['diskpart'], stdin=PIPE, stdout=PIPE, stderr=PIPE, text=True)
        output, error = process.communicate(input='\n'.join(commands))
        return output

    def refresh_disks(self):
        output = self.run_diskpart(['list disk'])
        lines = output.splitlines()
        self.disk_table.setRowCount(0)  # 테이블 초기화
        for line in lines:
            if "Disk" in line and "Status" in line:
                continue  # 헤더 건너뛰기
            if "Disk" in line:
                parts = line.split()
                if len(parts) >= 5:
                    disk_num, size, status = parts[1], parts[4] + " " + parts[5], parts[3]
                    row = self.disk_table.rowCount()
                    self.disk_table.insertRow(row)
                    self.disk_table.setItem(row, 0, QTableWidgetItem(f"Disk {disk_num}"))
                    self.disk_table.setItem(row, 1, QTableWidgetItem(size))
                    self.disk_table.setItem(row, 2, QTableWidgetItem(status))

    def prepare_disk(self):
        selected = self.disk_table.currentRow()
        if selected == -1:
            QMessageBox.warning(self, "오류", "디스크를 선택하세요!")
            return
        disk_num = self.disk_table.item(selected, 0).text().split()[1]
        partition_type = "mbr" if self.mbr_radio.isChecked() else "gpt"

        commands = [
            f"select disk {disk_num}",
            "clean",
            f"convert {partition_type}",
            "create partition primary",
            "format fs=ntfs quick",
            "active" if partition_type == "mbr" else "",
        ]
        result = self.run_diskpart(commands)
        QMessageBox.information(self, "완료", f"Disk {disk_num}이 {partition_type}로 준비되었습니다.")

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

•