#!/usr/bin/env python3
"""
MTK Flash Tool GUI for macOS
A graphical interface for flashing MediaTek devices,
similar to SP Flash Tool but native on macOS.
Built with PySide6 (Qt6).
"""

import sys
import os
import time
import threading

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QGridLayout, QPushButton, QLabel, QLineEdit, QFileDialog,
    QTreeWidget, QTreeWidgetItem, QTextEdit, QProgressBar,
    QGroupBox, QRadioButton, QCheckBox, QMessageBox, QSplitter,
    QHeaderView, QButtonGroup, QSizePolicy, QSpacerItem,
)
from PySide6.QtCore import Qt, Signal, QObject, QTimer
from PySide6.QtGui import QFont, QColor, QTextCharFormat, QBrush, QPalette, QIcon

from scatter_parser import parse_scatter, auto_load_roms, ScatterFile, PartitionEntry, EmmcRegion
from serial_connection import SerialConnection, SerialError
from brom import BromProtocol, BromError
from da_parser import find_da_for_chip, load_da_regions
from da_protocol import (
    DAInterface, DAError,
    HW_STORAGE_EMMC, HW_STORAGE_UFS,
    EMMC_BOOT1, EMMC_BOOT2, EMMC_USER,
)


# ============================================================================
# Colors - SP Flash Tool style (light/classic Windows look)
# ============================================================================
C_BG = "#f0f0f0"
C_BG2 = "#e8e8e8"
C_PANEL = "#ffffff"
C_INPUT = "#ffffff"
C_BORDER = "#b0b0b0"
C_TEXT = "#000000"
C_DIM = "#808080"
C_BRIGHT = "#000000"
C_ACCENT = "#0066cc"
C_GREEN = "#008000"
C_RED = "#cc0000"
C_YELLOW = "#cc8800"
C_ORANGE = "#ff6600"
C_CYAN = "#0088cc"


def _fmt_size(size: int) -> str:
    if size >= 1024 ** 3:
        return f"{size / 1024**3:.2f} GB"
    if size >= 1024 ** 2:
        return f"{size / 1024**2:.1f} MB"
    if size >= 1024:
        return f"{size / 1024:.1f} KB"
    return f"{size} B"


def _fmt_hex(v: int) -> str:
    return f"0x{v:X}"


def _region_to_part_id(region: EmmcRegion) -> int:
    return {
        EmmcRegion.EMMC_BOOT1: EMMC_BOOT1,
        EmmcRegion.EMMC_BOOT2: EMMC_BOOT2,
    }.get(region, EMMC_USER)


# ============================================================================
# Signals bridge (thread → UI)
# ============================================================================
class Signals(QObject):
    log = Signal(str, str)           # (level, message)
    status = Signal(str)             # status text
    speed = Signal(str)              # speed text
    progress = Signal(float)         # 0-100
    partition_status = Signal(str, str, str)  # (name, status_text, color)
    finished = Signal(bool, str)     # (success, message)


# ============================================================================
# Main Window
# ============================================================================
class MTKFlashWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Smart Phone Flash Tool for macOS")
        self.setMinimumSize(900, 650)
        self.resize(960, 700)

        self.scatter: ScatterFile = None
        self.scatter_path = ""
        self.da_path = ""
        self.is_flashing = False
        self.signals = Signals()

        self._connect_signals()
        self._apply_theme()
        self._build_ui()

    # ---------------------------------------------------------------- theme
    def _apply_theme(self):
        self.setStyleSheet(f"""
            QMainWindow, QWidget {{
                background-color: {C_BG};
                color: {C_TEXT};
                font-family: "Segoe UI", "Helvetica Neue", "Arial", sans-serif;
                font-size: 12px;
            }}
            QGroupBox {{
                background-color: {C_PANEL};
                border: 1px solid {C_BORDER};
                border-radius: 3px;
                margin-top: 12px;
                padding: 12px 8px 8px 8px;
                font-weight: bold;
                font-size: 12px;
                color: {C_TEXT};
            }}
            QGroupBox::title {{
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 4px;
            }}
            QLineEdit {{
                background-color: {C_INPUT};
                border: 1px solid {C_BORDER};
                border-radius: 2px;
                padding: 4px 6px;
                color: {C_TEXT};
                font-size: 11px;
            }}
            QLineEdit:focus {{
                border-color: {C_ACCENT};
            }}
            QPushButton {{
                background-color: #e1e1e1;
                border: 1px solid {C_BORDER};
                border-radius: 2px;
                padding: 5px 14px;
                color: {C_TEXT};
                font-size: 11px;
                min-width: 70px;
            }}
            QPushButton:hover {{
                background-color: #d0d0d0;
                border-color: {C_ACCENT};
            }}
            QPushButton:pressed {{
                background-color: #c0c0c0;
            }}
            QPushButton:disabled {{
                background-color: {C_BG};
                color: {C_DIM};
            }}
            QPushButton#flashBtn {{
                background-color: #4CAF50;
                color: white;
                font-size: 14px;
                font-weight: bold;
                padding: 10px;
                border: 1px solid #388E3C;
                border-radius: 3px;
            }}
            QPushButton#flashBtn:hover {{
                background-color: #45a049;
            }}
            QPushButton#flashBtn:disabled {{
                background-color: #cccccc;
                color: #888888;
                border-color: #aaaaaa;
            }}
            QPushButton#stopBtn {{
                background-color: #f44336;
                color: white;
                font-size: 14px;
                font-weight: bold;
                padding: 10px;
                border: 1px solid #d32f2f;
                border-radius: 3px;
            }}
            QPushButton#stopBtn:hover {{
                background-color: #e53935;
            }}
            QTreeWidget {{
                background-color: white;
                border: 1px solid {C_BORDER};
                border-radius: 1px;
                font-size: 11px;
                outline: none;
                gridline-color: #d0d0d0;
            }}
            QTreeWidget::item {{
                padding: 2px 0;
                border-bottom: 1px solid #e8e8e8;
            }}
            QTreeWidget::item:selected {{
                background-color: #cce5ff;
                color: black;
            }}
            QHeaderView::section {{
                background-color: #e8e8e8;
                color: {C_TEXT};
                border: 1px solid #d0d0d0;
                border-top: none;
                padding: 4px 6px;
                font-weight: bold;
                font-size: 11px;
            }}
            QTextEdit {{
                background-color: white;
                border: 1px solid {C_BORDER};
                border-radius: 1px;
                font-family: "Consolas", "SF Mono", "Menlo", monospace;
                font-size: 11px;
                color: {C_TEXT};
                padding: 4px;
            }}
            QProgressBar {{
                background-color: white;
                border: 1px solid {C_BORDER};
                border-radius: 1px;
                text-align: center;
                color: {C_TEXT};
                font-size: 11px;
                min-height: 18px;
            }}
            QProgressBar::chunk {{
                background-color: #3399ff;
            }}
            QRadioButton, QCheckBox {{
                color: {C_TEXT};
                spacing: 5px;
                font-size: 12px;
            }}
            QComboBox {{
                background-color: white;
                border: 1px solid {C_BORDER};
                border-radius: 2px;
                padding: 3px 6px;
            }}
            QSplitter::handle {{
                background-color: {C_BORDER};
                width: 2px;
            }}
            QTabWidget::pane {{
                border: 1px solid {C_BORDER};
                background: white;
            }}
            QTabBar::tab {{
                background: #e8e8e8;
                border: 1px solid {C_BORDER};
                padding: 6px 16px;
                margin-right: 2px;
            }}
            QTabBar::tab:selected {{
                background: white;
                border-bottom: none;
            }}
        """)

    # ---------------------------------------------------------------- signals
    def _connect_signals(self):
        self.signals.log.connect(self._on_log)
        self.signals.status.connect(self._on_status)
        self.signals.speed.connect(self._on_speed)
        self.signals.progress.connect(self._on_progress)
        self.signals.partition_status.connect(self._on_partition_status)
        self.signals.finished.connect(self._on_finished)

    # ---------------------------------------------------------------- build
    def _build_ui(self):
        central = QWidget()
        self.setCentralWidget(central)
        root_layout = QVBoxLayout(central)
        root_layout.setContentsMargins(16, 10, 16, 12)
        root_layout.setSpacing(8)

        # No title bar - SP Flash Tool uses window title only
        root_layout.setContentsMargins(6, 6, 6, 6)
        root_layout.setSpacing(4)

        # Splitter: left (files + table) | right (actions + log)
        splitter = QSplitter(Qt.Horizontal)
        splitter.setHandleWidth(3)

        left = QWidget()
        left_layout = QVBoxLayout(left)
        left_layout.setContentsMargins(0, 0, 0, 0)
        left_layout.setSpacing(6)

        right = QWidget()
        right_layout = QVBoxLayout(right)
        right_layout.setContentsMargins(0, 0, 0, 0)
        right_layout.setSpacing(6)

        # --- Left: DA + Scatter (SP Flash Tool layout)
        files_grp = QGroupBox("Download")
        files_lay = QGridLayout(files_grp)
        files_lay.setSpacing(4)
        files_lay.setColumnStretch(1, 1)

        files_lay.addWidget(QLabel("Download Agent:"), 0, 0)
        self.da_edit = QLineEdit()
        self.da_edit.setReadOnly(True)
        files_lay.addWidget(self.da_edit, 0, 1)
        btn_da = QPushButton("DA Agent")
        btn_da.clicked.connect(self._browse_da)
        files_lay.addWidget(btn_da, 0, 2)

        files_lay.addWidget(QLabel("Scatter-loading File:"), 1, 0)
        self.scatter_edit = QLineEdit()
        files_lay.addWidget(self.scatter_edit, 1, 1)
        btn_scatter = QPushButton("Scatter-loading")
        btn_scatter.clicked.connect(self._browse_scatter)
        files_lay.addWidget(btn_scatter, 1, 2)

        self.platform_lbl = QLabel("")
        self.platform_lbl.setStyleSheet(f"color: {C_ACCENT}; font-size: 11px; font-weight: bold;")
        files_lay.addWidget(self.platform_lbl, 2, 0, 1, 3)

        left_layout.addWidget(files_grp)

        # --- Left: Partition table
        parts_grp = QGroupBox("Partitions")
        parts_lay = QVBoxLayout(parts_grp)

        self.select_all_cb = QCheckBox("Select All")
        self.select_all_cb.setChecked(True)
        self.select_all_cb.stateChanged.connect(self._toggle_all)
        parts_lay.addWidget(self.select_all_cb)

        self.tree = QTreeWidget()
        self.tree.setHeaderLabels(["", "Name", "Begin Address", "End Address", "Size", "Location"])
        self.tree.setColumnWidth(0, 28)
        self.tree.setColumnWidth(1, 120)
        self.tree.setColumnWidth(2, 130)
        self.tree.setColumnWidth(3, 130)
        self.tree.setColumnWidth(4, 80)
        self.tree.setColumnWidth(5, 240)
        self.tree.itemDoubleClicked.connect(self._on_partition_double_click)
        self.tree.setRootIsDecorated(False)
        self.tree.setAlternatingRowColors(True)
        self.tree.setStyleSheet(self.tree.styleSheet() + """
            QTreeWidget {
                alternate-background-color: #f5f5ff;
            }
        """)
        parts_lay.addWidget(self.tree)

        left_layout.addWidget(parts_grp, 1)

        # --- Right: Actions
        actions_grp = QGroupBox("Actions")
        actions_lay = QVBoxLayout(actions_grp)
        actions_lay.setSpacing(8)

        # Mode selector (like SP Flash Tool combo)
        from PySide6.QtWidgets import QComboBox
        mode_row = QHBoxLayout()
        mode_row.addWidget(QLabel("Mode:"))
        self.mode_combo = QComboBox()
        self.mode_combo.addItems(["Download Only", "Format All + Download"])
        self.mode_combo.setCurrentIndex(0)
        mode_row.addWidget(self.mode_combo, 1)
        actions_lay.addLayout(mode_row)

        # Dummy button group for compatibility
        self.mode_group = QButtonGroup(self)
        self.mode_download = QRadioButton()
        self.mode_download.setChecked(True)
        self.mode_download.hide()
        self.mode_group.addButton(self.mode_download, 0)
        self.mode_combo.currentIndexChanged.connect(
            lambda i: None)  # Mode 0=download, 1=format+download

        actions_lay.addSpacing(6)

        # Big Download button (SP Flash Tool style)
        self.flash_btn = QPushButton("Download")
        self.flash_btn.setObjectName("flashBtn")
        self.flash_btn.setMinimumHeight(44)
        self.flash_btn.clicked.connect(self._start_flash)
        actions_lay.addWidget(self.flash_btn)

        self.stop_btn = QPushButton("Stop")
        self.stop_btn.setObjectName("stopBtn")
        self.stop_btn.setMinimumHeight(36)
        self.stop_btn.setEnabled(False)
        self.stop_btn.clicked.connect(self._stop_flash)
        actions_lay.addWidget(self.stop_btn)

        actions_lay.addStretch()

        right_layout.addWidget(actions_grp)

        # --- Right: Log
        log_grp = QGroupBox("Log")
        log_lay = QVBoxLayout(log_grp)

        log_header = QHBoxLayout()
        log_header.addStretch()
        clear_btn = QPushButton("Clear")
        clear_btn.clicked.connect(lambda: self.log_text.clear())
        log_header.addWidget(clear_btn)
        log_lay.addLayout(log_header)

        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        log_lay.addWidget(self.log_text)

        right_layout.addWidget(log_grp, 1)

        # Add to splitter
        splitter.addWidget(left)
        splitter.addWidget(right)
        splitter.setSizes([600, 400])

        root_layout.addWidget(splitter, 1)

        # --- Bottom: progress
        bottom = QHBoxLayout()
        bottom.setSpacing(10)

        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        self.progress_bar.setFormat("%p%")
        bottom.addWidget(self.progress_bar, 1)

        self.status_lbl = QLabel("Ready")
        self.status_lbl.setStyleSheet("font-size: 11px; color: #006600;")
        self.status_lbl.setMinimumWidth(160)
        bottom.addWidget(self.status_lbl)

        self.speed_lbl = QLabel("")
        self.speed_lbl.setStyleSheet("font-size: 11px; color: #666666;")
        self.speed_lbl.setMinimumWidth(120)
        self.speed_lbl.setAlignment(Qt.AlignRight)
        bottom.addWidget(self.speed_lbl)

        root_layout.addLayout(bottom)

        # Initial log
        self._log_append("info", "MTK Flash Tool for macOS")

        # Auto-load DA_PL.bin from tool directory
        da_bundled = os.path.join(os.path.dirname(os.path.abspath(__file__)), "DA_PL.bin")
        if os.path.isfile(da_bundled):
            self.da_edit.setText(da_bundled)
            self.da_path = da_bundled
            self._log_append("success", f"DA loaded: {da_bundled}")
        else:
            self._log_append("info", "Load a scatter file and DA to begin")

        # Auto-load last scatter file
        self._settings_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".last_scatter")
        try:
            with open(self._settings_file, "r") as f:
                last_scatter = f.read().strip()
            if os.path.isfile(last_scatter):
                self.scatter_edit.setText(last_scatter)
                self._load_scatter(last_scatter)
                self._log_append("info", f"Loaded last scatter: {os.path.basename(last_scatter)}")
        except Exception:
            pass

    # ---------------------------------------------------------------- browse
    def _browse_scatter(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "Select Scatter File", "",
            "Scatter Files (*.txt *.scatter);;All Files (*)")
        if path:
            self.scatter_edit.setText(path)
            self._load_scatter(path)

    def _browse_da(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "Select DA File", "",
            "DA Files (*.bin *.img);;All Files (*)")
        if path:
            self.da_edit.setText(path)
            self.da_path = path
            self._log_append("info", f"DA: {os.path.basename(path)}")

    def _load_scatter(self, path: str):
        try:
            self.scatter = parse_scatter(path)
            self.scatter_path = path
            auto_load_roms(self.scatter)

            # Remember last scatter path
            try:
                with open(self._settings_file, "w") as f:
                    f.write(path)
            except Exception:
                pass

            self.platform_lbl.setText(
                f"Platform: {self.scatter.header.platform}   |   "
                f"Storage: {self.scatter.header.storage}   |   "
                f"Partitions: {len(self.scatter.partitions)}")

            self._populate_tree()

            # Auto-detect DA
            scatter_dir = os.path.dirname(path)
            for f in os.listdir(scatter_dir):
                if f.lower().startswith("da_") and f.lower().endswith(".bin"):
                    da = os.path.join(scatter_dir, f)
                    self.da_edit.setText(da)
                    self.da_path = da
                    self._log_append("info", f"Auto-detected DA: {f}")
                    break

            en = len(self.scatter.get_enabled_partitions())
            self._log_append("success",
                             f"Loaded: {self.scatter.header.platform} — "
                             f"{en}/{len(self.scatter.partitions)} partitions ready")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to load scatter:\n{e}")
            self._log_append("error", str(e))

    def _populate_tree(self):
        self.tree.clear()
        if not self.scatter:
            return
        for p in self.scatter.partitions:
            has_file = bool(p.file_path and os.path.isfile(p.file_path))
            fname = os.path.basename(p.file_name) if p.file_name else "—"
            fsize_str = ""
            if has_file:
                fsize_str = f"  ({_fmt_size(os.path.getsize(p.file_path))})"

            # SP Flash Tool style: Name, Begin, End, Size, Location
            end_addr = p.start_addr + p.size if p.size else 0
            location = p.file_path if has_file else "(double-click to load)"
            file_size = os.path.getsize(p.file_path) if has_file else p.size
            size_str = _fmt_size(file_size) if file_size else "—"

            item = QTreeWidgetItem([
                "", p.name.upper(),
                f"0x{p.start_addr:016X}",
                f"0x{end_addr:016X}",
                size_str,
                location,
            ])
            item.setData(0, Qt.UserRole, p.name)  # store partition name
            item.setCheckState(0, Qt.Checked if has_file else Qt.Unchecked)

            if has_file:
                for col in range(1, 6):
                    item.setForeground(col, QBrush(QColor("#000080")))
            else:
                for col in range(1, 6):
                    item.setForeground(col, QBrush(QColor(C_DIM)))

            # Make row selectable even without file (for double-click to load)
            item.setFlags(item.flags() | Qt.ItemIsSelectable)

            self.tree.addTopLevelItem(item)

    def _on_partition_double_click(self, item, column):
        """Double-click partition row to load ROM file from custom location."""
        if not self.scatter:
            return
        part_name = item.data(0, Qt.UserRole)
        part = self.scatter.get_partition_by_name(part_name)
        if not part:
            return

        # Open file dialog
        default_dir = os.path.dirname(self.scatter_path) if self.scatter_path else ""
        path, _ = QFileDialog.getOpenFileName(
            self, f"Load ROM for {part_name}",
            default_dir,
            "Image Files (*.img *.bin *.ext4);;All Files (*)")

        if path and os.path.isfile(path):
            part.file_path = path
            part.enabled = True
            self._log_append("success", f"Loaded {part_name}: {os.path.basename(path)}")
            self._populate_tree()

    def _toggle_all(self, state):
        checked = Qt.Checked if state else Qt.Unchecked
        for i in range(self.tree.topLevelItemCount()):
            item = self.tree.topLevelItem(i)
            if item.flags() & Qt.ItemIsUserCheckable:
                item.setCheckState(0, checked)

    # ---------------------------------------------------------------- log
    def _log_append(self, level: str, msg: str):
        colors = {"info": "#000000", "success": "#006600", "error": "#cc0000",
                  "warning": "#cc6600"}
        color = colors.get(level, "#000000")
        ts = time.strftime("%H:%M:%S")
        self.log_text.append(
            f'<span style="color:#0066cc">[{ts}]</span> '
            f'<span style="color:{color}">{msg}</span>')
        sb = self.log_text.verticalScrollBar()
        sb.setValue(sb.maximum())

    def _on_log(self, level, msg):
        self._log_append(level, msg)

    def _on_status(self, text):
        self.status_lbl.setText(text)

    def _on_speed(self, text):
        self.speed_lbl.setText(text)

    def _on_progress(self, val):
        self.progress_bar.setValue(int(val))

    def _on_partition_status(self, name, status_text, color):
        for i in range(self.tree.topLevelItemCount()):
            item = self.tree.topLevelItem(i)
            if item.text(1) == name:
                item.setText(5, status_text)
                item.setForeground(5, QBrush(QColor(color)))
                break

    def _on_finished(self, success, message):
        self.is_flashing = False
        self.flash_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        if success:
            QMessageBox.information(self, "Success", message)
        else:
            QMessageBox.critical(self, "Error", message)

    # ---------------------------------------------------------------- flash
    def _get_selected_partitions(self) -> list:
        selected = []
        if not self.scatter:
            return selected
        for i in range(self.tree.topLevelItemCount()):
            item = self.tree.topLevelItem(i)
            if item.checkState(0) == Qt.Checked:
                name = item.text(1)
                p = self.scatter.get_partition_by_name(name)
                if p and p.file_path:
                    selected.append(p)
        return selected

    def _start_flash(self):
        if self.is_flashing:
            return

        if not self.scatter:
            QMessageBox.warning(self, "Warning", "Load a scatter file first.")
            return

        da_path = self.da_edit.text().strip()
        if not da_path or not os.path.isfile(da_path):
            QMessageBox.warning(self, "Warning", "Select a valid DA file.")
            return
        self.da_path = da_path

        selected = self._get_selected_partitions()
        if not selected:
            QMessageBox.warning(self, "Warning", "No partitions selected.")
            return

        total = sum(os.path.getsize(p.file_path) for p in selected)
        names = "\n".join(f"  • {p.name}" for p in selected)
        ret = QMessageBox.question(
            self, "Confirm Flash",
            f"Flash {len(selected)} partitions ({_fmt_size(total)})?\n\n"
            f"{names}\n\n"
            f"This will OVERWRITE data on the device!",
            QMessageBox.Yes | QMessageBox.No)

        if ret != QMessageBox.Yes:
            return

        self.is_flashing = True
        self.flash_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.progress_bar.setValue(0)

        thread = threading.Thread(target=self._flash_thread,
                                   args=(selected,), daemon=True)
        thread.start()

    def _flash_thread_OLD_UNUSED(self, partitions_UNUSED):
        pass

    def _flash_thread(self, partitions):
        sig = self.signals
        mtk = None
        caffeinate_proc = None

        try:
            import subprocess
            from mtk_device import MTKDevice, MTKError
            from da_parser import find_da_for_chip, load_da_regions

            def log_cb(level, msg):
                sig.log.emit(level, msg)

            # Prevent Mac from sleeping during flash
            try:
                caffeinate_proc = subprocess.Popen(
                    ["caffeinate", "-d", "-i", "-s"],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL)
                sig.log.emit("info", "Sleep prevention enabled (caffeinate)")
            except Exception as e:
                sig.log.emit("warning", f"Sleep prevention: {e}")

            sig.status.emit("Waiting for device...")

            attempt = 0
            t0 = time.time()

            while time.time() - t0 < 120:
                if not self.is_flashing:
                    raise Exception("Cancelled")

                attempt += 1
                mtk = MTKDevice(log_cb=log_cb)

                try:
                    sig.status.emit(f"Connecting ({attempt})...")
                    mtk.connect(timeout_sec=10)
                    mtk.brom_handshake()
                    mtk.get_hw_code()
                    mtk.get_hw_sw_ver()
                    try:
                        mtk.get_target_config()
                    except Exception:
                        pass

                    hw = mtk.hw_code
                    da_entry = find_da_for_chip(self.da_path, hw)
                    if not da_entry:
                        da_entry = find_da_for_chip(self.da_path,
                            {0x0717:0x6761, 0x0507:0x6761}.get(hw, hw))
                    if not da_entry:
                        raise MTKError(f"No DA for {hw:#06x}")

                    regions = load_da_regions(self.da_path, da_entry)

                    sig.status.emit("Uploading DA...")
                    mtk.full_init(self.da_path, regions)
                    sig.log.emit("success", "DA ready")
                    break

                except Exception as e:
                    sig.log.emit("warning", f"Attempt {attempt}: {e}")
                    mtk.disconnect()
                    time.sleep(0.5)

            if not mtk or not mtk.dev:
                raise Exception("Connection failed after 120s")

            # Flash partitions
            total = sum(os.path.getsize(p.file_path) for p in partitions)
            sent = 0
            t1 = time.time()
            skip_fmt = (self.mode_group.checkedId() == 0)

            for i, part in enumerate(partitions):
                if not self.is_flashing:
                    break
                sig.status.emit(f"{part.name} ({i+1}/{len(partitions)})")
                sig.partition_status.emit(part.name, "⚡", C_YELLOW)

                with open(part.file_path, "rb") as f:
                    rom = f.read()

                # Map scatter region to eMMC partition type
                from scatter_parser import EmmcRegion
                region_map = {
                    EmmcRegion.EMMC_BOOT1: 1,
                    EmmcRegion.EMMC_BOOT2: 2,
                    EmmcRegion.EMMC_USER: 8,
                }
                pt = region_map.get(part.region, 8)
                # EMMC_BOOT1 means write to both BOOT1 and BOOT2
                write_to_boot2 = (part.region == EmmcRegion.EMMC_BOOT1 and
                                  "preloader" in part.name.lower())
                sig.log.emit("info", f"  region={part.region.name} parttype={pt}"
                             f"{' (+BOOT2)' if write_to_boot2 else ''}")

                # FORMAT disabled - leaves stale responses that break WRITE
                # TODO: fix FORMAT drain to re-enable
                # if not skip_fmt:
                #     try:
                #         mtk.format_partition(...)
                #     except: pass

                tp = time.time()
                def pcb(s, t, _s=sent, _t=total):
                    sig.progress.emit((_s+s)*100/_t if _t else 0)
                mtk.write_partition(part.start_addr, rom, progress_cb=pcb,
                                    parttype=pt)

                # NOTE: BOOT2 duplicate write disabled - may corrupt preloader
                # eMMC automatically mirrors BOOT1→BOOT2 on most devices

                el = time.time() - tp
                sig.log.emit("success",
                    f"  ✓ {_fmt_size(len(rom))} "
                    f"({len(rom)/el/1048576:.1f} MB/s)" if el > 0.01 else
                    f"  ✓ {_fmt_size(len(rom))}")
                sig.partition_status.emit(part.name, "✓", C_GREEN)
                sent += len(rom)

            sig.progress.emit(100)
            sig.status.emit("Flash complete!")
            sig.finished.emit(True, f"Done! {_fmt_size(sent)}")

        except Exception as e:
            sig.log.emit("error", f"ERROR: {e}")
            sig.finished.emit(False, str(e))
        finally:
            if mtk:
                try:
                    mtk.disconnect()
                except Exception:
                    pass
            # Release sleep prevention
            if caffeinate_proc:
                try:
                    caffeinate_proc.terminate()
                    caffeinate_proc.wait(timeout=2)
                except Exception:
                    pass

    def _flash_thread_REPLACED(self, partitions):
        sig = self.signals
        usb = None
        da = None

        try:
            sig.status.emit("Waiting for device...")
            sig.log.emit("info", "Keep device connected - catching bootloop window...")

            # Bootloop-aware connection loop using Serial (CDC ACM)
            # Like SP Flash Tool's FlashTool_Connect_BROM_Ex with loop_until_found=true
            info = None
            brom = None
            conn = None
            start_time = time.time()
            attempt = 0

            while time.time() - start_time < 120:
                if not self.is_flashing:
                    raise Exception("Cancelled")

                attempt += 1
                conn = SerialConnection()

                try:
                    if not conn.find_and_connect(timeout_sec=10):
                        continue

                    sig.status.emit(f"BROM handshake (attempt {attempt})...")
                    sig.log.emit("info", f"Port: {conn.port_name}, handshake attempt {attempt}...")

                    brom = BromProtocol(conn)
                    info = brom.init_device(max_handshake_retries=1)

                    sig.log.emit("success", f"Device: {info.chip_name} ({info.hw_code:#06x})")
                    break

                except (BromError, SerialError, Exception) as e:
                    sig.log.emit("warning", f"Attempt {attempt}: {e}")
                    sig.log.emit("info", "Waiting for device reboot cycle...")
                    try:
                        conn.disconnect()
                    except Exception:
                        pass
                    conn = None
                    brom = None
                    time.sleep(0.5)
                    continue

            if not info or not brom or not conn:
                raise Exception("Could not connect after 120s. Check USB cable and device.")

            # DA upload - parse DA blob and send correct chip entry
            sig.status.emit("Uploading DA...")
            sig.log.emit("info", f"Uploading DA: {os.path.basename(self.da_path)}")

            # Find DA entry for this chip
            # Try exact hw_code first, then base chip family (0x6761 for MT6761D)
            hw = info.hw_code
            da_entry = find_da_for_chip(self.da_path, hw)
            if not da_entry:
                # Try common base codes: 0x0717 → 0x6761, etc.
                base_codes = {
                    0x0717: 0x6761, 0x0507: 0x6761,
                    0x0690: 0x6762, 0x0707: 0x6765,
                    0x0725: 0x6768, 0x0788: 0x6771,
                }
                base = base_codes.get(hw, hw)
                da_entry = find_da_for_chip(self.da_path, base)

            if not da_entry:
                raise Exception(f"No DA entry for hw_code {hw:#06x} in {self.da_path}")

            sig.log.emit("info", f"DA entry: hw={da_entry.hw_code:#06x}, "
                         f"{da_entry.entry_region_count} regions")

            # Load DA regions
            regions = load_da_regions(self.da_path, da_entry)
            if len(regions) < 2:
                raise Exception("DA entry has less than 2 regions")

            # Stage 1: Send region[1] via BROM and jump to it
            # region[0] is a tiny stub/descriptor, region[1] is the actual DA1
            da1_addr, da1_data, da1_sig = regions[1]
            sig.log.emit("info", f"  DA1: {len(da1_data)} bytes → {da1_addr:#010x}")
            brom.send_da(da1_addr, da1_data, da1_sig)
            brom.jump_da(da1_addr)
            sig.log.emit("success", "DA1 running")

            # === THE COMBO THAT WORKED ===
            # 1. Reconnect serial at 921600
            # 2. Open pyusb ALONGSIDE serial
            # 3. Detach kernel driver Interface 0 only
            # 4. CDC control transfers via pyusb
            # 5. XFlash data via serial
            import struct as _struct
            import usb.core
            import usb.util

            conn.disconnect()
            time.sleep(0.5)

            # Step 1: Reconnect serial at 921600
            new_conn = None
            for _ in range(15):
                try:
                    c = SerialConnection()
                    if c.find_and_connect(timeout_sec=1, baud=921600):
                        new_conn = c
                        break
                except Exception:
                    pass
                time.sleep(0.2)
            if not new_conn:
                raise Exception("Serial reconnect failed after DA1")
            conn = new_conn
            sig.log.emit("info", f"Serial: {conn.port_name} @ 921600")

            # Step 2: Open pyusb alongside (for CDC control transfers)
            dev = usb.core.find(idVendor=0x0E8D)
            if dev:
                # Detach kernel driver Interface 0 ONLY
                try:
                    if dev.is_kernel_driver_active(0):
                        dev.detach_kernel_driver(0)
                        sig.log.emit("info", "Detached intf 0 kernel driver")
                except Exception:
                    pass

                # CDC control transfers
                try:
                    dev.ctrl_transfer(0x21, 0x20, 0, 0,
                                      _struct.pack("<IBBB", 921600, 0, 0, 8))
                    dev.ctrl_transfer(0x21, 0x22, 0x0003, 0)
                    sig.log.emit("info", "CDC: LINE_CODING + DTR/RTS set")
                except Exception as e:
                    sig.log.emit("warning", f"CDC ctrl: {e}")
            else:
                sig.log.emit("warning", "pyusb: device not found")

            # XFlash through serial
            XFLASH_MAGIC = 0xFEEEEEEF

            def xwrite(data):
                conn.write(data)

            def xread_raw(size=512, timeout=1000):
                try:
                    return conn.read(size, timeout_ms=timeout)
                except Exception:
                    return b""

            def xsend(cmd_val):
                payload = _struct.pack("<I", cmd_val)
                header = _struct.pack("<III", XFLASH_MAGIC, 1, len(payload))
                xwrite(header + payload)

            def xsend_data(data_bytes):
                header = _struct.pack("<III", XFLASH_MAGIC, 1, len(data_bytes))
                xwrite(header + data_bytes)

            sig.log.emit("info", "XFlash init...")

            def xsend_read(name, cmd, params=None):
                """Send XFlash cmd, read response, send params."""
                try:
                    xsend(cmd)
                    r = xread_raw(64, 1000)
                    if r:
                        sig.log.emit("info", f"  {name} → [{' '.join(f'{b:02X}' for b in r[:20])}]")
                    else:
                        sig.log.emit("info", f"  {name}: no resp")
                    if params:
                        xsend_data(params)
                        r2 = xread_raw(64, 1000)
                        if r2:
                            sig.log.emit("info", f"  {name}+ → [{' '.join(f'{b:02X}' for b in r2[:20])}]")
                    return r
                except Exception as e:
                    sig.log.emit("warning", f"  {name}: {e}")
                    return None

            # Rapid sequence: SYNC → ENV → HW_INIT
            xsend_read("SYNC", 0x434E5953)
            xsend_read("ENV", 0x010100, _struct.pack("<IIIII", 0, 1, 1, 0, 0))
            xsend_read("HW", 0x010101, _struct.pack("<I", 0))

            # DA2 upload via BOOT_TO
            if len(regions) >= 3:
                da2_addr, da2_data, da2_sig = regions[2]
                sig.log.emit("info", f"DA2: {len(da2_data)} bytes → {da2_addr:#x}")
                try:
                    xsend(0x010008)  # BOOT_TO
                    r = xread_raw(64, 2000)
                    if r:
                        sig.log.emit("info", f"  BOOT_TO → [{' '.join(f'{b:02X}' for b in r[:16])}]")

                    xsend_data(_struct.pack("<QQ", da2_addr, len(da2_data)))

                    # Send DA2 in chunks through serial
                    chunk = 4096
                    for off in range(0, len(da2_data), chunk):
                        conn.write(da2_data[off:off+chunk])
                    sig.log.emit("info", "  DA2 sent")

                    time.sleep(1)
                    r = xread_raw(64, 5000)
                    if r:
                        sig.log.emit("info", f"  DA2 done → [{' '.join(f'{b:02X}' for b in r[:16])}]")
                except Exception as e:
                    sig.log.emit("warning", f"DA2: {e}")

            sig.log.emit("success", "DA init complete")

            # DA interface on serial
            da = DAInterface(conn)
            da.is_xflash = True
            da.xflash = __import__('da_protocol').XFlashProtocol(conn)
            da.report = __import__('da_protocol').DAReport()
            sig.log.emit("success", "DA ready")

            storage_type = HW_STORAGE_EMMC
            if self.scatter and self.scatter.header.storage.upper() == "UFS":
                storage_type = HW_STORAGE_UFS

            mode_id = self.mode_group.checkedId()
            skip_format = (mode_id == 0)

            total_size = sum(os.path.getsize(p.file_path) for p in partitions)
            total_sent = 0
            flash_start = time.time()

            for i, part in enumerate(partitions):
                if not self.is_flashing:
                    sig.log.emit("warning", "Cancelled by user")
                    break

                sig.status.emit(f"Flashing {part.name} ({i+1}/{len(partitions)})")
                sig.partition_status.emit(part.name, "⚡ Flashing...", C_YELLOW)
                sig.log.emit("info", f"[{i+1}/{len(partitions)}] {part.name}")

                with open(part.file_path, "rb") as f:
                    rom_data = f.read()

                part_id = _region_to_part_id(part.region)
                file_size = len(rom_data)

                if not skip_format:
                    sig.log.emit("info", f"  Formatting...")
                    try:
                        fmt_size = max(file_size, part.size) if part.size else file_size
                        da.format_partition(part.start_addr, fmt_size,
                                            part_id=part_id, storage_type=storage_type)
                    except DAError as e:
                        sig.log.emit("warning", f"  Format: {e}")

                sig.log.emit("info", f"  Writing {_fmt_size(file_size)}...")
                part_start = time.time()

                def progress_cb(sent, total, _ts=total_sent, _tt=total_size, _ps=part_start):
                    pct = ((_ts + sent) * 100) / _tt if _tt else 0
                    sig.progress.emit(pct)
                    el = time.time() - _ps
                    spd = sent / el if el > 0 else 0
                    sig.speed.emit(f"{spd / (1024*1024):.1f} MB/s")

                da.write_partition(part.start_addr, rom_data,
                                   part_id=part_id, storage_type=storage_type,
                                   progress_cb=progress_cb)

                elapsed = time.time() - part_start
                spd = file_size / elapsed if elapsed else 0
                sig.log.emit("success",
                             f"  ✓ {_fmt_size(file_size)} in {elapsed:.1f}s "
                             f"({spd/(1024*1024):.1f} MB/s)")
                sig.partition_status.emit(part.name, "✓ Done", C_GREEN)
                total_sent += file_size

            # Done
            total_time = time.time() - flash_start
            avg = total_sent / total_time if total_time else 0

            sig.progress.emit(100)
            sig.status.emit("Flash complete!")
            sig.speed.emit(f"Avg: {avg/(1024*1024):.1f} MB/s")

            if self.reboot_cb.isChecked():
                sig.log.emit("info", "Rebooting device...")
                da.shutdown(reboot=True)
            else:
                da.shutdown(reboot=False)

            sig.finished.emit(True,
                              f"Flash complete!\n\n"
                              f"{len(partitions)} partitions\n"
                              f"{_fmt_size(total_sent)} in {total_time:.1f}s\n"
                              f"Avg speed: {avg/(1024*1024):.1f} MB/s")

        except Exception as e:
            sig.log.emit("error", f"ERROR: {e}")
            sig.status.emit(f"Error")
            sig.finished.emit(False, str(e))
            try:
                if da:
                    da.shutdown(reboot=False)
            except Exception:
                pass
        finally:
            if conn:
                try:
                    conn.disconnect()
                except Exception:
                    pass

    def _stop_flash(self):
        self.is_flashing = False
        self._log_append("warning", "Stopping...")
        self.status_lbl.setText("Stopping...")

    # ---------------------------------------------------------------- detect
    def _detect_device(self):
        def _run():
            sig = self.signals
            try:
                sig.log.emit("info", "Detecting device (60s, auto-retry for bootloop)...")
                sig.status.emit("Waiting for device...")

                start = time.time()
                attempt = 0
                while time.time() - start < 60:
                    attempt += 1
                    conn = SerialConnection()
                    try:
                        if not conn.find_and_connect(timeout_sec=5):
                            continue
                        brom = BromProtocol(conn)
                        info = brom.init_device(max_handshake_retries=1)
                        sig.log.emit("success", f"Found: {info.chip_name}")
                        sig.log.emit("info", str(info))
                        sig.status.emit(f"Device: {info.chip_name}")
                        conn.disconnect()
                        return
                    except Exception as e:
                        sig.log.emit("warning", f"Attempt {attempt}: {e}")
                        try:
                            conn.disconnect()
                        except Exception:
                            pass
                        time.sleep(0.3)

                sig.log.emit("error", "No device found after 60s")
                sig.status.emit("No device")
            except Exception as e:
                sig.log.emit("error", f"Detection failed: {e}")
                sig.status.emit("Detection failed")

        threading.Thread(target=_run, daemon=True).start()


# ============================================================================
# Main
# ============================================================================
def main():
    app = QApplication(sys.argv)
    app.setApplicationName("MTK Flash Tool")
    window = MTKFlashWindow()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()
