import os
import json
from PySide2.QtWidgets import (
    QWidget,
    QApplication,
    QVBoxLayout,
    QHBoxLayout,
    QPushButton,
    QTreeWidget,
    QTreeWidgetItem,
    QFileDialog,
    QCheckBox,
    QLabel,
    QMessageBox,
    QDialog,
    QScrollArea,
)
from PySide2.QtCore import Qt
import Metashape


class BatchImageLoader(QDialog):  # 继承自QDialog
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setModal(True)  # 设置为模态对话框

        self.setWindowTitle("Batch Image Loader")
        self.setMinimumSize(600, 800)

        # Main layout
        self.layout = QVBoxLayout()

        # Folder selection
        self.folder_layout = QHBoxLayout()
        self.folder_label = QLabel("Root Folder:")
        self.folder_path = QLabel("No folder selected")
        self.select_folder_btn = QPushButton("Browse...")
        self.select_folder_btn.clicked.connect(self.select_folder)

        self.folder_layout.addWidget(self.folder_label)
        self.folder_layout.addWidget(self.folder_path)
        self.folder_layout.addWidget(self.select_folder_btn)

        # Camera group checkbox
        self.camera_group_cb = QCheckBox("Use second-level folders as camera groups")
        self.camera_group_cb.stateChanged.connect(self.update_preview)

        # Preview tree
        self.preview_layout = QHBoxLayout()
        self.preview_label = QLabel("Import Preview:")

        self.btn_expand_all = QPushButton("Expand All")
        self.btn_expand_all.clicked.connect(self.expand_all_items)
        self.btn_collapse_all = QPushButton("Collapse All")
        self.btn_collapse_all.clicked.connect(self.collapse_all_items)

        self.btn_select_all = QPushButton("Select All")
        self.btn_select_all.clicked.connect(self.select_all_items)
        self.btn_clear_all = QPushButton("Clear")
        self.btn_clear_all.clicked.connect(self.clear_all_items)

        self.preview_layout.addWidget(self.preview_label)
        self.preview_layout.addStretch()
        self.preview_layout.addWidget(self.btn_expand_all)
        self.preview_layout.addWidget(self.btn_collapse_all)
        self.preview_layout.addWidget(self.btn_select_all)
        self.preview_layout.addWidget(self.btn_clear_all)

        self.tree_widget = QTreeWidget()
        self.tree_widget.setHeaderLabel("Workspace Structure")

        # Help button
        self.help_btn = QPushButton("Help")
        self.help_btn.clicked.connect(self.show_help)

        # Import button
        self.import_btn = QPushButton("Import Images")
        self.import_btn.clicked.connect(self.import_images)
        self.import_btn.setEnabled(False)

        # Add widgets to layout
        self.layout.addLayout(self.folder_layout)
        self.layout.addWidget(self.camera_group_cb)

        self.layout.addLayout(self.preview_layout)
        self.layout.addWidget(self.tree_widget)
        self.layout.addWidget(self.help_btn)
        self.layout.addWidget(self.import_btn)

        self.setLayout(self.layout)

        # Variables
        self.root_path = ""
        self.img_ext = (".jpg", ".jpeg", ".png", ".tif", ".tiff")

        from .. import system_info

        self.config_manager = system_info.config_manager

    def select_folder(self):
        last_path = self.config_manager.load("last_batch_import_folder")
        folder = QFileDialog.getExistingDirectory(
            self, "Select Root Image Folder", dir=last_path
        )

        if folder:
            self.config_manager.save("last_batch_import_folder", folder)

            self.root_path = folder
            self.folder_path.setText(folder)
            self.update_preview()
            self.import_btn.setEnabled(True)

    def get_folder_structure(self, root_path, use_camera_groups):
        """Shared function to scan folder structure and return organized data"""
        structure = {
            "chunks": [],
        }

        # Build regular structure
        #  增加文件名的排序
        for item in sorted(os.listdir(root_path)):
            item_path = os.path.join(root_path, item)
            if os.path.isdir(item_path):
                chunk_data = {"name": item, "path": item_path, "groups": []}
                if use_camera_groups:
                    #  增加文件名的排序
                    for sub_item in sorted(os.listdir(item_path)):
                        sub_item_path = os.path.join(item_path, sub_item)
                        if os.path.isdir(sub_item_path):
                            group_data = {
                                "name": sub_item,
                                "path": sub_item_path,
                                "images": [
                                    f
                                    for f in os.listdir(sub_item_path)
                                    if f.lower().endswith(self.img_ext)
                                ],
                            }
                            chunk_data["groups"].append(group_data)
                else:
                    chunk_data["images"] = [
                        f
                        for f in os.listdir(item_path)
                        if f.lower().endswith(self.img_ext)
                    ]
                structure["chunks"].append(chunk_data)

        return structure

    def update_preview(self):
        """Update the tree widget preview with checkboxes for folder selection"""

        if not self.root_path:
            return

        self.tree_widget.clear()
        use_camera_groups = self.camera_group_cb.isChecked()
        structure = self.get_folder_structure(self.root_path, use_camera_groups)

        self.ignored_folders = set()  # 存储用户选择忽略的文件夹

        self.chunk_labels = [chunk.label for chunk in Metashape.app.document.chunks]

        try:
            for chunk in structure["chunks"]:
                # 检查chunk是否有效（包含图片或有效group）
                if not use_camera_groups:
                    # 情况一：不使用相机组。
                    # 直接检查当前 chunk 是否包含任何照片。
                    is_valid_chunk = bool(chunk["images"])
                else:
                    # 情况二：使用相机组。
                    # 遍历 chunk 内的所有相机组（groups），检查是否有任何一个组（group）包含照片。
                    # any() 函数会在找到第一个包含照片的组后立即返回 True，效率较高。
                    is_valid_chunk = any(group["images"] for group in chunk["groups"])

                # Add chunk item with checkbox
                chunk_item = QTreeWidgetItem(self.tree_widget)

                chunk_suffix = (
                    "Chunk" if is_valid_chunk else "invalid chunk without images"
                )

                # Check if chunk exists in project
                is_exist_chunk = chunk["name"] in self.chunk_labels
                if is_exist_chunk:
                    if is_valid_chunk:
                        chunk_suffix += ", exists"

                chunk_item.setText(0, f"📁 {chunk['name']} ({chunk_suffix})")
                chunk_item.setData(0, Qt.UserRole, chunk["path"])  # 存储完整路径

                if is_exist_chunk:
                    # >>> 后续还需要同步self.ignored_folders 来确保忽略
                    self.ignored_folders.add(chunk["path"])

                chunk_item.setCheckState(
                    0,
                    Qt.Checked
                    if is_valid_chunk and not is_exist_chunk
                    else Qt.Unchecked,
                )

                if not is_valid_chunk:
                    chunk_item.setFlags(
                        chunk_item.flags() & ~Qt.ItemIsEnabled
                    )  # 禁用无效chunk
                    continue

                if use_camera_groups:
                    # Add camera groups with checkboxes
                    for group in chunk["groups"]:
                        is_valid_group = bool(group["images"])

                        group_item = QTreeWidgetItem(chunk_item)
                        group_suffix = (
                            "Camera Group"
                            if is_valid_group
                            else "invalid camera group without images"
                        )

                        group_item.setText(0, f"📷 {group['name']} ({group_suffix})")
                        group_item.setData(0, Qt.UserRole, group["path"])
                        group_item.setData(0, Qt.UserRole + 1, is_valid_group)

                        if not is_valid_group:
                            group_item.setCheckState(0, Qt.Unchecked)
                            group_item.setFlags(
                                group_item.flags() & ~Qt.ItemIsEnabled
                            )  # 禁用无效group
                            continue

                        if is_exist_chunk:
                            # Parent is unchecked (due to existence), so child should be disabled
                            group_item.setCheckState(0, Qt.Unchecked)
                            group_item.setFlags(group_item.flags() & ~Qt.ItemIsEnabled)
                        else:
                            group_item.setCheckState(0, Qt.Checked)

                        # Add sample images (no checkboxes for images)
                        for img in group["images"][:3]:  # Show first 3 as sample
                            img_item = QTreeWidgetItem(group_item)
                            img_item.setText(0, f"🖼 {img}")
                        if len(group["images"]) > 3:
                            more_item = QTreeWidgetItem(group_item)
                            more_item.setText(
                                0, f"... and {len(group['images']) - 3} more"
                            )
                else:
                    # Add images directly under chunk (no checkboxes for images)
                    images = chunk.get("images", [])
                    for img in images[:5]:
                        img_item = QTreeWidgetItem(chunk_item)
                        img_item.setText(0, f"🖼 {img}")
                    if len(images) > 5:
                        more_item = QTreeWidgetItem(chunk_item)
                        more_item.setText(0, f"... and {len(images) - 5} more")

            # 连接itemChanged信号以跟踪复选框状态变化<sup>1</sup>
            self.tree_widget.itemChanged.connect(self.on_item_changed)
            self.tree_widget.expandAll()

        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error scanning folder: {str(e)}")

    def on_item_changed(self, item, column):
        """Handle checkbox state changes"""
        state = item.checkState(0)

        # 同步子项状态
        # 同步子项状态
        if item.parent() is None:  # 如果是chunk顶级项
            for i in range(item.childCount()):
                child = item.child(i)

                # Check if this is a structural child (group) or just an image leaf
                # Structural items have path in UserRole
                if not child.data(0, Qt.UserRole):
                    continue

                is_valid = child.data(0, Qt.UserRole + 1)

                if state == Qt.Unchecked:
                    # Parent disabled -> Disable and uncheck all children
                    child.setFlags(child.flags() & ~Qt.ItemIsEnabled)
                    child.setCheckState(0, Qt.Unchecked)
                elif state == Qt.Checked:
                    # Parent enabled -> Enable valid children and check them
                    if is_valid:
                        child.setFlags(child.flags() | Qt.ItemIsEnabled)
                        child.setCheckState(0, Qt.Checked)
                    else:
                        # Invalid children remain disabled
                        child.setFlags(child.flags() & ~Qt.ItemIsEnabled)
                        child.setCheckState(0, Qt.Unchecked)

        if state == Qt.Unchecked:
            self.ignored_folders.add(
                item.data(0, Qt.UserRole)
            )  # 添加到忽略列表<sup>1</sup>
        else:
            self.ignored_folders.discard(
                item.data(0, Qt.UserRole)
            )  # 从忽略列表移除<sup>1</sup>

    def get_ignored_folders(self):
        """Return the list of folders marked for ignoring"""
        return list(self.ignored_folders)

    def select_all_items(self):
        """Select all enabled items in the tree"""
        root = self.tree_widget.invisibleRootItem()
        for i in range(root.childCount()):
            item = root.child(i)
            if item.flags() & Qt.ItemIsEnabled:
                item.setCheckState(0, Qt.Checked)

    def clear_all_items(self):
        """Unselect all items in the tree"""
        root = self.tree_widget.invisibleRootItem()
        for i in range(root.childCount()):
            item = root.child(i)
            # Even disabled items can be unchecked safely
            item.setCheckState(0, Qt.Unchecked)

    def expand_all_items(self):
        """Expand all items in the tree"""
        self.tree_widget.expandAll()

    def collapse_all_items(self):
        """Collapse all items in the tree"""
        self.tree_widget.collapseAll()

    def import_images(self):
        """Perform the actual import"""
        if not self.root_path:
            return

        doc = Metashape.app.document

        use_camera_groups = self.camera_group_cb.isChecked()
        structure = self.get_folder_structure(self.root_path, use_camera_groups)

        existing_labels = set(c.label for c in doc.chunks)
        apply_to_all_choice = None  # None, 'create', 'skip'

        try:
            for chunk in structure["chunks"]:
                # 跳过未选中的chunk
                chunk_path = os.path.join(self.root_path, chunk["name"])
                if chunk_path in self.ignored_folders:
                    continue

                # Check for duplicates and ask user if needed
                if chunk["name"] in existing_labels:
                    valid_choice = "create"  # default

                    if apply_to_all_choice:
                        valid_choice = apply_to_all_choice
                    else:
                        # Show Warning Dialog
                        msg = QMessageBox(self)
                        msg.setWindowTitle("Duplicate Chunk Detected")
                        msg.setText(
                            f"The chunk '{chunk['name']}' already exists in the project."
                        )
                        msg.setInformativeText(
                            f"Do you want to create a new chunk for '{chunk['name']}' or skip it?"
                        )

                        btn_create = msg.addButton("Create New", QMessageBox.AcceptRole)
                        btn_skip = msg.addButton("Skip", QMessageBox.RejectRole)

                        cb_apply = QCheckBox("Apply to all future conflicts")
                        msg.setCheckBox(cb_apply)

                        msg.exec_()

                        if msg.clickedButton() == btn_skip:
                            valid_choice = "skip"
                        else:
                            valid_choice = "create"

                        if cb_apply.isChecked():
                            apply_to_all_choice = valid_choice

                    if valid_choice == "skip":
                        continue

                # Create new chunk
                ms_chunk = doc.addChunk()
                ms_chunk.label = chunk["name"]

                if use_camera_groups and chunk["groups"]:
                    # Add camera groups
                    for group in chunk["groups"]:
                        # 跳过未选中的group
                        group_path = os.path.join(chunk_path, group["name"])
                        if group_path in self.ignored_folders:
                            continue

                        if group["images"]:
                            # Create camera group
                            camera_group = ms_chunk.addCameraGroup()
                            camera_group.label = group["name"]

                            # Add images to group <sup>1</sup>
                            image_paths = [
                                os.path.join(group["path"], img)
                                for img in group["images"]
                            ]
                            ms_chunk.addPhotos(
                                image_paths,
                                group=camera_group.key,
                                load_xmp_accuracy=True,  # support loading RTK info
                                load_rpc_txt=True,
                            )

                else:
                    # Add images directly to chunk
                    images = chunk.get("images", [])
                    if images:
                        image_paths = [
                            os.path.join(chunk["path"], img) for img in images
                        ]
                        ms_chunk.addPhotos(
                            image_paths,
                            load_xmp_accuracy=True,  # support loading RTK info
                            load_rpc_txt=True,
                        )

            QMessageBox.information(self, "Success", "Images imported successfully!")

        except Exception as e:
            QMessageBox.critical(self, "Error", f"Error during import: {str(e)}")

    def import_multi_camera(self, path, chunk_name):
        """Handle multi-camera system import (not implemented)"""
        doc = Metashape.app.document
        chunk = doc.addChunk()
        chunk.label = chunk_name

        # Show dialog to confirm multi-camera import
        reply = QMessageBox.question(
            self,
            "Multi-camera System",
            f"Folder '{chunk_name}' appears to contain images from a multi-camera system.\n"
            "Do you want to import these as a multi-camera rig?",
            QMessageBox.Yes | QMessageBox.No,
        )

        if reply == QMessageBox.Yes:
            # Import as multi-camera system <sup>2</sup>
            chunk.addPhotos(
                [
                    os.path.join(path, f)
                    for f in os.listdir(path)
                    if os.path.isdir(os.path.join(path, f))
                ],
                layout=Metashape.MultiplaneLayout,
            )
        else:
            # Import as regular images
            images = []
            for root, dirs, files in os.walk(path):
                images.extend(
                    [
                        os.path.join(root, f)
                        for f in files
                        if f.lower().endswith(
                            (".jpg", ".jpeg", ".png", ".tif", ".tiff")
                        )
                    ]
                )
            if images:
                chunk.addPhotos(images)

    def show_help(self):
        """Show help dialog"""
        help_dialog = QDialog(self)
        help_dialog.setWindowTitle("Help")
        help_dialog.setMinimumSize(500, 400)

        layout = QVBoxLayout()
        scroll = QScrollArea()
        content = QWidget()
        content_layout = QVBoxLayout()

        help_text = QLabel("""
        <h2>Batch Image Loader Help</h2>
        <p>This tool allows you to import multiple folders of images into Metashape with automatic chunk and camera group creation.</p>
        
        <h3>Basic Usage:</h3>
        <ol>
            <li>Click "Browse..." to select the root folder containing your images</li>
            <li>Check "Use second-level folders as camera groups" if you want subfolders to become camera groups</li>
            <li>Review the import structure in the preview</li>
            <li>Click "Import Images" to perform the import</li>
        </ol>
        
        <h3>Folder Structure:</h3>
        <p><b>Without camera groups:</b><br>
<pre>
Root/
├── Chunk1/ (becomes chunk)
│   ├── image1.jpg
│   └── image2.jpg
└── Chunk2/
    ├── image1.jpg
    └── image2.jpg
</pre>
        </p>
        
        <p><b>With camera groups:</b><br>
<pre>
Root/
├── Chunk1/ (becomes chunk)
│   ├── Group1/ (becomes camera group)
│   │   ├── image1.jpg
│   │   └── image2.jpg
│   └── Group2/
│       ├── image1.jpg
│       └── image2.jpg
└── Chunk2/
    ├── Group1/
    │   ├── image1.jpg
    │   └── image2.jpg
    └── Group2/
        ├── image1.jpg
        └── image2.jpg
</pre>
        </p>
        
        <h3>Multi-camera Systems:</h3>
        <p>Importing multi-camera systems have not been supported yet. This plugin is currently designed for RGB cameras only.</p>
        """)
        help_text.setWordWrap(True)
        help_text.setTextFormat(Qt.RichText)

        content_layout.addWidget(help_text)
        content.setLayout(content_layout)
        scroll.setWidget(content)
        scroll.setWidgetResizable(True)

        close_btn = QPushButton("Close")
        close_btn.clicked.connect(help_dialog.close)

        layout.addWidget(scroll)
        layout.addWidget(close_btn)
        help_dialog.setLayout(layout)

        help_dialog.exec_()


def create_batch_image_importer():
    app = QApplication.instance()  # 获取当前Qt应用实例
    window = BatchImageLoader(app.activeWindow())
    window.exec_()  # 使用exec_()而非show()确保模态性
