import sys
import os
import json
from pathlib import Path
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QLabel,
    QPushButton, QFileDialog, QLineEdit, QListWidget, QMessageBox, QHBoxLayout, QInputDialog
)
from PySide6.QtCore import Qt, QMimeData
from PySide6.QtGui import QDragEnterEvent, QDropEvent

CONFIG_FILE = "category_names.json"

class FileRenamer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("批量文件重命名工具")
        self.setMinimumSize(600, 400)
        self.setAcceptDrops(True)

        self.files = []
        self.initial_path = ""
        self.prefix = ""
        self.categories = self.load_categories()

        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        layout = QVBoxLayout(central_widget)

        self.info_label = QLabel("拖拽文件到此区域")
        self.info_label.setAlignment(Qt.AlignCenter)
        self.info_label.setStyleSheet("border: 2px dashed #aaa; padding: 20px; font-size: 16px;")
        layout.addWidget(self.info_label)

        path_layout = QHBoxLayout()
        self.path_btn = QPushButton("设置初始路径")
        self.path_btn.clicked.connect(self.set_initial_path)
        self.path_label = QLabel("未设置")
        path_layout.addWidget(self.path_btn)
        path_layout.addWidget(self.path_label)
        layout.addLayout(path_layout)

        prefix_layout = QHBoxLayout()
        self.prefix_btn = QPushButton("设置前缀")
        self.prefix_btn.clicked.connect(self.set_prefix)
        self.prefix_label = QLabel("当前前缀: (无)")
        prefix_layout.addWidget(self.prefix_btn)
        prefix_layout.addWidget(self.prefix_label)
        layout.addLayout(prefix_layout)

        self.name_input = QLineEdit()
        self.name_input.setPlaceholderText("输入新文件名（不含前缀和序号）")
        layout.addWidget(self.name_input)

        self.suggestion_list = QListWidget()
        self.suggestion_list.itemClicked.connect(self.select_suggestion)
        layout.addWidget(self.suggestion_list)

        btn_layout = QHBoxLayout()
        self.rename_btn = QPushButton("重命名")
        self.rename_btn.clicked.connect(self.rename_files)
        btn_layout.addWidget(self.rename_btn)
        layout.addLayout(btn_layout)

    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event: QDropEvent):
        for url in event.mimeData().urls():
            file_path = url.toLocalFile()
            if os.path.isfile(file_path):
                self.files.append(file_path)
        self.info_label.setText(f"已添加 {len(self.files)} 个文件")

        if self.initial_path:
            categories_found = set()
            for f in self.files:
                rel_path = os.path.relpath(f, self.initial_path)
                category = os.path.dirname(rel_path).replace("\\", "/")
                if category:
                    categories_found.add(category)
            if categories_found:
                self.show_suggestions(categories_found)

    def set_initial_path(self):
        folder = QFileDialog.getExistingDirectory(self, "选择初始路径")
        if folder:
            self.initial_path = folder
            self.path_label.setText(folder)

    def set_prefix(self):
        prefix, ok = QInputDialog.getText(self, "设置前缀", "输入前缀:", text=self.prefix if self.prefix else "")
        if ok:
            self.prefix = prefix
            self.prefix_label.setText(f"当前前缀: {self.prefix if self.prefix else '(无)'}")

    def rename_files(self):
        if not self.files:
            QMessageBox.warning(self, "警告", "请先拖拽文件")
            return

        new_name = self.name_input.text().strip()
        if not new_name:
            QMessageBox.warning(self, "警告", "请输入新文件名")
            return

        for idx, file_path in enumerate(self.files, start=1):
            folder = os.path.dirname(file_path)
            ext = os.path.splitext(file_path)[1]
            new_filename = f"{self.prefix}{new_name}{idx:02}{ext}"
            new_path = os.path.join(folder, new_filename)
            os.rename(file_path, new_path)

        if self.initial_path:
            for f in self.files:
                rel_path = os.path.relpath(f, self.initial_path)
                category = os.path.dirname(rel_path).replace("\\", "/")
                if category:
                    self.categories.setdefault(category, set()).add(new_name)

            self.save_categories()

        QMessageBox.information(self, "完成", "重命名完成")
        self.files.clear()
        self.info_label.setText("拖拽文件到此区域")

    def load_categories(self):
        if os.path.exists(CONFIG_FILE):
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
                return {k: set(v) for k, v in data.items()}
        return {}

    def save_categories(self):
        data = {k: list(v) for k, v in self.categories.items()}
        with open(CONFIG_FILE, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    def show_suggestions(self, categories_found):
        self.suggestion_list.clear()
        suggestions = set()
        for cat in categories_found:
            if cat in self.categories:
                suggestions.update(self.categories[cat])
        self.suggestion_list.addItems(sorted(suggestions))

    def select_suggestion(self, item):
        self.name_input.setText(item.text())

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