"""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 typing import TYPE_CHECKING

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

if TYPE_CHECKING:
    from PySide2.QtGui import QMoveEvent, QResizeEvent

CONFIG_FILE = Path.home() / ".pytola" / "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 != "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:
        """处理窗口移动事件, 保存窗口位置."""
        top_left = self.geometry().topLeft()
        conf.WIN_POS = [top_left.x(), top_left.y()]
        return super().moveEvent(event)

    def resizeEvent(self, event: QResizeEvent) -> None:
        """处理窗口调整大小事件, 保存窗口大小."""
        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

        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

        # Disconnect signals to allow worker to be garbage collected
        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.exception("错误: 未找到llama.cpp/quantize工具")
        logger.exception(
            "请确保已编译llama.cpp并将quantize工具放在llama.cpp/目录下",
        )
        sys.exit(1)

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