"""GGUF量化转换GUI工具, 用于将F16格式的GGUF文件转换为其他主流量化格式."""

from __future__ import annotations

import atexit
import json
import logging
import pathlib
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

from PySide2.QtCore import QThread, Signal, Slot
from PySide2.QtGui import QMoveEvent, QResizeEvent
from PySide2.QtWidgets import (
    QApplication,
    QCheckBox,
    QFileDialog,
    QGridLayout,
    QGroupBox,
    QHBoxLayout,
    QLabel,
    QMainWindow,
    QProgressBar,
    QPushButton,
    QTextEdit,
    QVBoxLayout,
    QWidget,
)

CONFIG_FILE = Path.home() / ".pysfi" / "llmquantize.json"

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

__version__ = "1.0.0"
__build__ = "20260204"


@dataclass
class QuantizerConfig:
    """GGUF量化转换工具配置."""

    TITLE: str = "GGUF量化转换工具"
    WIN_SIZE: list[int] = None
    WIN_POS: list[int] = None
    LAST_INPUT_FILE: str = ""
    SELECTED_QUANTS: list[str] = None

    def __post_init__(self) -> None:
        """初始化默认值并加载配置文件."""
        # 初始化默认值
        if self.WIN_SIZE is None:
            self.WIN_SIZE = [600, 500]
        if self.WIN_POS is None:
            self.WIN_POS = [100, 100]
        if self.SELECTED_QUANTS is None:
            self.SELECTED_QUANTS = ["Q4_K_M", "Q5_K_M"]

        if CONFIG_FILE.exists():
            logger.info("Loading configuration from %s", CONFIG_FILE)
            try:
                config_data = json.loads(CONFIG_FILE.read_text())
                # 更新实例属性，只更新存在的属性
                for key, value in config_data.items():
                    if hasattr(self, key):
                        setattr(self, key, value)
            except (json.JSONDecodeError, TypeError, AttributeError) as e:
                logger.warning("Failed to load configuration: %s", e)
                logger.info("Using default configuration")
        else:
            logger.info("Using default configuration")

    def save(self) -> None:
        """保存配置到文件."""
        CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
        # 将数据类转换为字典进行JSON序列化
        config_dict = {}
        for attr_name in dir(self):
            if not attr_name.startswith("_") and attr_name not in ["TITLE"]:
                try:
                    attr_value = getattr(self, attr_name)
                    if not callable(attr_value):
                        config_dict[attr_name] = attr_value
                except AttributeError:
                    continue
        CONFIG_FILE.write_text(json.dumps(config_dict, indent=4))


conf = QuantizerConfig()
atexit.register(conf.save)


def _process_gguf_stem(filename: str) -> str:
    """处理文件名, 移除可能的F16后缀.

    Args:
        filename: 输入的文件名(不含扩展名)

    Returns:
        str: 处理后的文件名, 移除了F16后缀(如果存在)
    """
    if filename.upper().endswith("-F16"):
        filename = filename[:-4]  # 移除-F16后缀
    return filename


class QuantizationWorker(QThread):
    """量化执行线程Worker.

    Attributes:
        progress_msg_updated: 进度消息更新信号
        progress_count_updated: 进度数值更新信号
        is_finished: 完成信号
    """

    progress_msg_updated = Signal(str)
    progress_count_updated = Signal(int)
    is_finished = Signal()

    def __init__(
        self,
        input_file: pathlib.Path,
        quant_types: list[str],
    ) -> None:
        """初始化量化Worker.

        Args:
            input_file: 输入的F16 GGUF文件路径
            quant_types: 需要转换的量化类型列表
        """
        super().__init__()

        self.input_file = input_file
        self.quant_types = quant_types
        self.input_dir = input_file.parent
        self.base_name = _process_gguf_stem(input_file.stem)
        self.total_files = len(quant_types)
        self.completed_files = 0
        self.success = True  # 记录整体转换状态

    def run(self) -> None:
        """执行量化转换任务."""
        try:
            for quant_type in self.quant_types:
                output_file: pathlib.Path = (
                    self.input_dir / f"{self.base_name}-{quant_type}.gguf"
                )

                self.progress_msg_updated.emit(
                    f"正在转换到 {quant_type} 格式...",
                )

                # 构建命令行参数
                cmd = [
                    "llama-quantize",
                    str(self.input_file.name),
                    str(output_file.name),
                    quant_type,
                ]

                # 执行转换命令
                try:
                    process = subprocess.Popen(
                        cmd,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.STDOUT,
                        universal_newlines=True,
                        bufsize=1,  # 行缓冲以获得实时输出
                        cwd=str(self.input_file.parent),
                    )
                except (FileNotFoundError, PermissionError) as e:
                    self.progress_msg_updated.emit(f"无法启动量化进程: {e!s}")
                    self.success = False
                    continue

                # 实时输出进度
                if not process.stdout:
                    logger.error("无法获取进度信息")
                    process.wait()
                    self.success = False
                    continue

                try:
                    for line in process.stdout:
                        self.progress_msg_updated.emit(line.strip())
                finally:
                    # 等待进程结束，设置超时防止永久挂起
                    try:
                        process.wait(timeout=3600)  # 单个文件最多1小时
                    except subprocess.TimeoutExpired:
                        self.progress_msg_updated.emit(f"转换 {quant_type} 超时")
                        process.kill()
                        process.wait()
                        self.success = False
                        continue
                    finally:
                        # 确保关闭管道以避免资源泄漏
                        if process.stdout:
                            process.stdout.close()

                # 只有在成功时才更新进度
                if process.returncode == 0:
                    self.completed_files += 1
                    progress = int((self.completed_files / self.total_files) * 100)
                    self.progress_count_updated.emit(progress)
                    self.progress_msg_updated.emit(
                        f"成功生成: {output_file!s}",
                    )
                else:
                    self.progress_msg_updated.emit(f"转换 {quant_type} 失败")
                    self.success = False

            self.is_finished.emit()
        except Exception as e:
            logger.exception("量化转换过程中发生异常")
            self.progress_msg_updated.emit(f"发生错误: {e!s}")
            self.success = False
            self.is_finished.emit()


class GGUFQuantizerGUI(QMainWindow):
    """GGUF量化转换工具GUI界面."""

    def __init__(self) -> None:
        super().__init__()
        self.setWindowTitle(conf.TITLE)
        self.setGeometry(*conf.WIN_POS, *conf.WIN_SIZE)

        self.input_file: pathlib.Path = pathlib.Path(conf.LAST_INPUT_FILE)
        self.worker: QuantizationWorker | None = None
        self.quant_types = {
            "Q2_K": "Q2_K (极低精度, 最小尺寸)",
            "Q3_K_S": "Q3_K_S (低精度, 小尺寸)",
            "Q3_K_M": "Q3_K_M (低精度, 中等尺寸)",
            "Q3_K_L": "Q3_K_L (低精度, 大尺寸)",
            "Q4_0": "Q4_0 (基本4位)",
            "Q4_K_S": "Q4_K_S (4位, 小尺寸)",
            "Q4_K_M": "Q4_K_M (4位, 中等尺寸)",
            "Q5_0": "Q5_0 (基本5位)",
            "Q5_K_S": "Q5_K_S (5位, 小尺寸)",
            "Q5_K_M": "Q5_K_M (5位, 中等尺寸)",
            "Q6_K": "Q6_K (6位, 高质量)",
            "Q8_0": "Q8_0 (8位, 最高质量)",
        }
        self.quant_checks = {}  # 存储量化类型对应的checkbox

        self.init_ui()

        # 恢复上次选择的文件
        if conf.LAST_INPUT_FILE and pathlib.Path(conf.LAST_INPUT_FILE).exists():
            self.input_file = pathlib.Path(conf.LAST_INPUT_FILE)
            self.file_label.setText(self.input_file.name)
            self.check_existing_quant_files()
            self.convert_btn.setEnabled(True)

    def init_ui(self) -> None:
        """初始化界面."""
        main_widget = QWidget()
        main_layout = QVBoxLayout()

        # 文件选择部分
        file_group = QGroupBox("选择F16格式的GGUF文件")
        file_layout = QVBoxLayout()

        self.file_label = QLabel("未选择文件")
        file_btn = QPushButton("选择文件")
        file_btn.clicked.connect(self.select_file)

        file_layout.addWidget(self.file_label)
        file_layout.addWidget(file_btn)
        file_group.setLayout(file_layout)

        # 量化选项部分
        quant_group = QGroupBox("选择量化类型")
        quant_layout = QGridLayout()

        for i, (quant_type, label) in enumerate(self.quant_types.items()):
            check = QCheckBox(label)
            # 连接信号，选择时保存配置
            check.stateChanged.connect(self.on_quant_type_changed)
            self.quant_checks[quant_type] = check
            row = i // 2
            col = i % 2
            quant_layout.addWidget(check, row, col)

        # 恢复上次选择的量化类型
        for quant_type in conf.SELECTED_QUANTS:
            if quant_type in self.quant_checks:
                self.quant_checks[quant_type].setChecked(True)

        quant_group.setLayout(quant_layout)

        # 进度显示部分
        progress_group = QGroupBox("转换进度")
        progress_layout = QVBoxLayout()

        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)

        self.output_text = QTextEdit()
        self.output_text.setReadOnly(True)

        progress_layout.addWidget(self.progress_bar)
        progress_layout.addWidget(self.output_text)
        progress_group.setLayout(progress_layout)

        # 操作按钮
        btn_layout = QHBoxLayout()
        self.convert_btn = QPushButton("开始转换")
        self.convert_btn.clicked.connect(self.start_conversion)
        self.convert_btn.setEnabled(False)

        btn_layout.addWidget(self.convert_btn)

        # 组装主界面
        main_layout.addWidget(file_group)
        main_layout.addWidget(quant_group)
        main_layout.addWidget(progress_group)
        main_layout.addLayout(btn_layout)

        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

    def select_file(self) -> None:
        """选择F16 GGUF文件."""
        # 使用上次选择的目录作为初始目录
        initial_dir = ""
        if conf.LAST_INPUT_FILE and pathlib.Path(conf.LAST_INPUT_FILE).exists():
            initial_dir = str(pathlib.Path(conf.LAST_INPUT_FILE).parent)
        elif self.input_file.exists():
            initial_dir = str(self.input_file.parent)

        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "选择F16格式的GGUF文件",
            initial_dir,
            "GGUF Files (*.gguf)",
        )

        if file_path:
            self.input_file = pathlib.Path(file_path)
            filename = self.input_file.name
            self.file_label.setText(filename)

            # 保存最后选择的文件路径
            conf.LAST_INPUT_FILE = file_path

            # 检查文件名是否包含F16
            if "-F16" not in filename.upper():
                self.output_text.append(
                    "注意: 输入文件名不包含F16后缀,输出文件名将直接添加量化类型",
                )
                self._scroll_to_bottom()

            # 检查已存在的量化文件
            self.check_existing_quant_files()

            self.convert_btn.setEnabled(True)
            self.output_text.clear()
            self.progress_bar.setValue(0)
            self._scroll_to_bottom()

    def check_existing_quant_files(self) -> None:
        """检查当前目录下已存在的量化文件, 并更新checkbox状态."""
        if not self.input_file:
            return

        dir_path = self.input_file.parent

        for quant_type in self.quant_types:
            filename = f"{_process_gguf_stem(self.input_file.stem)}-{quant_type}.gguf"
            expected_file = dir_path / filename
            if expected_file.exists():
                # 文件已存在，标记并禁用，防止重复生成
                self.quant_checks[quant_type].setText(
                    f"{self.quant_types[quant_type]} (已存在)",
                )
                self.quant_checks[quant_type].setStyleSheet("color: orange;")
                self.quant_checks[quant_type].setChecked(True)
                self.quant_checks[quant_type].setEnabled(False)
            else:
                self.quant_checks[quant_type].setText(
                    self.quant_types[quant_type],
                )
                self.quant_checks[quant_type].setStyleSheet("")
                self.quant_checks[quant_type].setEnabled(True)

    def _scroll_to_bottom(self) -> None:
        """滚动输出框到底部."""
        scrollbar = self.output_text.verticalScrollBar()
        if scrollbar:
            scrollbar.setValue(scrollbar.maximum())

    def on_quant_type_changed(self, _state: int) -> None:
        """量化类型变更时保存配置."""
        selected_quants = [
            q for q, check in self.quant_checks.items() if check.isChecked()
        ]
        conf.SELECTED_QUANTS = selected_quants

    def moveEvent(self, event: QMoveEvent) -> None:  # noqa: N802
        """处理窗口移动事件，保存窗口位置."""
        top_left = self.geometry().topLeft()
        conf.WIN_POS = [top_left.x(), top_left.y()]
        return super().moveEvent(event)

    def resizeEvent(self, event: QResizeEvent) -> None:  # noqa: N802
        """处理窗口调整大小事件，保存窗口大小."""
        geometry = self.geometry()
        conf.WIN_SIZE = [geometry.width(), geometry.height()]
        return super().resizeEvent(event)

    def start_conversion(self) -> None:
        """开始转换量化任务."""
        # 检查是否已有任务在运行
        if self.worker and self.worker.isRunning():
            self.output_text.append("已有转换任务正在进行, 请等待完成")
            self._scroll_to_bottom()
            return

        selected_quants: list[str] = [
            q
            for q, check in self.quant_checks.items()
            if check.isChecked() and check.isEnabled()
        ]

        if not selected_quants:
            self.output_text.append("请至少选择一种量化类型")
            self._scroll_to_bottom()
            return

        if not self.input_file:
            self.output_text.append("请先选择输入文件")
            self._scroll_to_bottom()
            return

        # 注意：由于selected_quants只包含启用的复选框，
        # 所以不会包含已存在的禁用文件，无需额外检查覆盖

        self.convert_btn.setEnabled(False)
        self.progress_bar.setValue(0)
        self.output_text.append(f"开始转换: {self.input_file!s}")
        self.output_text.append(f"选择的量化类型: {', '.join(selected_quants)}")
        self.output_text.append(f"将生成 {len(selected_quants)} 个量化文件")

        self.worker = QuantizationWorker(self.input_file, selected_quants)
        # 连接信号
        self.worker.progress_msg_updated.connect(self.update_progress_msg)
        self.worker.is_finished.connect(self.conversion_finished)
        self.worker.progress_count_updated.connect(self.update_progress_value)
        self.worker.start()

    @Slot(str)
    def update_progress_msg(self, message: str) -> None:
        """更新进度信息."""
        self.output_text.append(message)
        self.output_text.ensureCursorVisible()
        self._scroll_to_bottom()

    @Slot(int)
    def update_progress_value(self, value: int) -> None:
        """更新进度条."""
        self.progress_bar.setValue(value)

    @Slot()
    def conversion_finished(self) -> None:
        """转换完成回调函数."""
        success = self.worker.success if self.worker else False

        # 断开信号连接，允许 worker 被垃圾回收
        if self.worker:
            try:
                self.worker.progress_msg_updated.disconnect(self.update_progress_msg)
                self.worker.is_finished.disconnect(self.conversion_finished)
                self.worker.progress_count_updated.disconnect(
                    self.update_progress_value
                )
            except RuntimeError:
                # Worker 可能已经被销毁
                pass

        self.convert_btn.setEnabled(True)

        if success:
            self.output_text.append("所有量化转换完成!")
            self.progress_bar.setValue(100)
            # 重新检查已存在的量化文件
            self.check_existing_quant_files()
        else:
            self.output_text.append("量化转换过程中出现错误!")
            self._scroll_to_bottom()


def main() -> None:
    """主程序入口."""
    app = QApplication(sys.argv)

    # 检查是否安装了llama.cpp
    try:
        subprocess.run(
            ["llama-quantize", "--help"],
            capture_output=True,
            check=False,
        )
    except FileNotFoundError:
        logger.error("错误: 未找到llama.cpp/quantize工具")
        logger.error(
            "请确保已编译llama.cpp并将quantize工具放在llama.cpp/目录下",
        )
        sys.exit(1)

    window = GGUFQuantizerGUI()
    window.show()
    sys.exit(app.exec_())
