from PyQt6.QtWidgets import QMainWindow, QFileDialog
from PyQt6.QtCore import Qt
from ui_main import Ui_MainWindow
import configparser
import subprocess
import os
from PyQt6.QtCore import Qt, QPoint
from PyQt6.QtWidgets import QStackedWidget
from PyQt6.QtWidgets import QMainWindow, QFileDialog

class UILoader(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.setWindowFlag(Qt.WindowType.FramelessWindowHint)
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        self.init_ui()
        self.update_batch_size_label(3)
        self.update_ctx_size_label(4)
        #self.load_config()
        self.offset = None  # Add this line

    def init_ui(self):
        self.minimize_button.clicked.connect(self.showMinimized)
        self.close_button.clicked.connect(self.close)
        self.model_browse_button.clicked.connect(self.browse_model)
        self.run_server_button.clicked.connect(self.run_server)
        self.save_config_button.clicked.connect(self.save_config)
        self.load_config_button.clicked.connect(self.load_config)

        self.top_bar.mousePressEvent = self.mouse_press_event
        self.top_bar.mouseMoveEvent = self.mouse_move_event

    def mouse_press_event(self, event):
        self.offset = event.pos()

    def update_threads_label(self, value):
        self.threads_label.setText(f"Threads: {value + 1}")

    def mouse_move_event(self, event):
        if self.offset is not None and event.buttons() == Qt.MouseButton.LeftButton:
            self.move(self.pos() + event.pos() - self.offset)

    def browse_model(self):
        model_path, _ = QFileDialog.getOpenFileName(self, "Select Model", os.getcwd(), "Model Files (*.gguf)")
        self.model_path_line_edit.setText(model_path)

    def run_server(self):
        model_path = self.model_path_line_edit.text()
        batch_size = self.batch_size_values[self.batch_size_slider.value()]
        ctx_size = self.ctx_size_values[self.ctx_size_slider.value()]
        layers_offload = self.layers_offload_spin_box.value()
        threads = self.threads_spin_box.value()
        use_flash_attention = self.flash_attention_checkbox.isChecked()
        is_moe = self.moe_checkbox.isChecked()
        moe_experts_string = ""
        if is_moe:
            experts_use = self.moe_experts_spin_box.value()
            moe_experts_string = f"--override-kv llama.expert_used_count=int:{experts_use}"

        custom_rope_freq_base = self.rope_freq_base_spin_box.value() if self.rope_freq_base_checkbox.isChecked() else None
        mlock = self.mlock_checkbox.isChecked()
        no_mmap = self.no_mmap_checkbox.isChecked()
        no_kv_offload = self.no_kv_offload_checkbox.isChecked()

        command = [
            "LlamaCPP\\llama-server.exe",
            f"--model {model_path}",
            f"--batch-size {batch_size}",
            f"--ctx-size {ctx_size}",
            f"--n-gpu-layers {layers_offload}",
            f"--threads {threads}",
        ]

        if use_flash_attention:
            command.append("--flash-attn")
        if mlock:
            command.append("--mlock")
        if no_mmap:
            command.append("--no-mmap")
        if no_kv_offload:
            command.append("--no-kv-offload")
        if custom_rope_freq_base is not None:
            command.append(f"--rope-freq-base {custom_rope_freq_base}")
        if moe_experts_string:
            command.append(moe_experts_string)

        print(" ".join(command))
        subprocess.run(" ".join(command), shell=True)

    def save_config(self):
        config = configparser.ConfigParser()
        config["Server"] = {
            "model_path": self.model_path_line_edit.text(),
            "batch_size": str(self.batch_size_slider.value()),
            "ctx_size": str(self.ctx_size_slider.value()),
            "layers_offload": str(self.layers_offload_spin_box.value()),
            "threads": str(self.threads_spin_box.value()),
            "use_flash_attention": str(self.flash_attention_checkbox.isChecked()),
            "is_moe": str(self.moe_checkbox.isChecked()),
            "moe_experts": str(self.moe_experts_spin_box.value()),
            "rope_freq_base_custom": str(self.rope_freq_base_checkbox.isChecked()),
            "rope_freq_base_value": str(self.rope_freq_base_spin_box.value()),
            "mlock": str(self.mlock_checkbox.isChecked()),
            "no_mmap": str(self.no_mmap_checkbox.isChecked()),
            "no_kv_offload": str(self.no_kv_offload_checkbox.isChecked()),
        }

        file_path, _ = QFileDialog.getSaveFileName(self, "Save Configuration", os.getcwd(), "INI Files (*.ini)")
        if file_path:
            with open(file_path, "w") as config_file:
                config.write(config_file)

    def load_config(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Load Configuration", os.getcwd(), "INI Files (*.ini)")
        if file_path:
            config = configparser.ConfigParser()
            config.read(file_path)

            self.model_path_line_edit.setText(config.get("Server", "model_path", fallback=""))
            self.batch_size_slider.setValue(config.getint("Server", "batch_size", fallback=512))
            self.ctx_size_slider.setValue(config.getint("Server", "ctx_size", fallback=8192))
            self.layers_offload_spin_box.setValue(config.getint("Server", "layers_offload", fallback=22))
            self.threads_spin_box.setValue(config.getint("Server", "threads", fallback=14))
            self.flash_attention_checkbox.setChecked(config.getboolean("Server", "use_flash_attention", fallback=False))
            self.moe_checkbox.setChecked(config.getboolean("Server", "is_moe", fallback=False))
            self.moe_experts_spin_box.setValue(config.getint("Server", "moe_experts", fallback=1))
            self.rope_freq_base_checkbox.setChecked(config.getboolean("Server", "rope_freq_base_custom", fallback=False))
            self.rope_freq_base_spin_box.setValue(config.getint("Server", "rope_freq_base_value", fallback=1))
            self.mlock_checkbox.setChecked(config.getboolean("Server", "mlock", fallback=False))
            self.no_mmap_checkbox.setChecked(config.getboolean("Server", "no_mmap", fallback=False))
            self.no_kv_offload_checkbox.setChecked(config.getboolean("Server", "no_kv_offload", fallback=False))