from __future__ import annotations

import atexit
import json
import logging
import os
import pathlib
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import ClassVar

from PySide2.QtCore import QProcess, QTextStream, QUrl
from PySide2.QtGui import (
    QBrush,
    QColor,
    QDesktopServices,
    QFont,
    QMoveEvent,
    QResizeEvent,
    QTextCharFormat,
    QTextCursor,
)
from PySide2.QtWidgets import (
    QApplication,
    QFileDialog,
    QGroupBox,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QMainWindow,
    QPushButton,
    QSpinBox,
    QTextEdit,
    QVBoxLayout,
    QWidget,
)

CONFIG_FILE = Path.home() / ".pysfi" / "llmserver.json"
logging.basicConfig(level="INFO", format="%(message)s")
logger = logging.getLogger(__name__)


@dataclass
class LLMServerConfig:
    """Llama local model server configuration."""

    TITLE: str = "Llama Local Model Server"
    WIN_SIZE: ClassVar[list[int]] = [800, 800]
    WIN_POS: ClassVar[list[int]] = [200, 200]
    MODEL_PATH: str = ""
    URL: str = "http://127.0.0.1"
    LISTEN_PORT: int = 8080
    LISTEN_PORT_RNG: ClassVar[list[int]] = [1024, 65535]
    THREAD_COUNT_RNG: ClassVar[list[int]] = [1, 24]
    THREAD_COUNT: int = 4

    _loaded_from_file: bool = False

    def __post_init__(self) -> None:
        if CONFIG_FILE.exists():
            logger.info("Loading configuration from %s", CONFIG_FILE)
            try:
                config_data = json.loads(CONFIG_FILE.read_text())
                # Update instance attributes with loaded values
                for key, value in config_data.items():
                    if hasattr(self, key):
                        setattr(self, key, value)
                self._loaded_from_file = True
            except (json.JSONDecodeError, TypeError, AttributeError) as e:
                logger.error(f"Error loading config from {CONFIG_FILE}: {e}")
        else:
            logger.info("Using default configuration")

    def save(self) -> None:
        """Save configuration."""
        try:
            CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
            # Convert dataclass to dict for JSON serialization
            config_dict = {}
            for attr_name in dir(self):
                if not attr_name.startswith("_"):
                    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))
            logger.info(f"Configuration saved to {CONFIG_FILE}")
        except Exception as e:
            logger.error(f"Failed to save configuration: {e}")
            raise


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


class LlamaServerGUI(QMainWindow):
    """Llama local model server GUI."""

    # Color constants
    COLOR_ERROR = QColor(255, 0, 0)
    COLOR_WARNING = QColor(255, 165, 0)
    COLOR_INFO = QColor(0, 0, 0)

    # Process constants
    PROCESS_TERMINATE_TIMEOUT_MS = 2000

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

        self.process: QProcess
        self.init_ui()
        self.setup_process()

        # Apply loaded configuration to UI controls
        self.apply_config_to_ui()

    def init_ui(self) -> None:
        """Initialize user interface."""
        # Main layout
        main_widget = QWidget()
        main_layout = QVBoxLayout()  # type: ignore

        # Configuration panel
        config_group = QGroupBox("Server Configuration")
        config_layout = QVBoxLayout(config_group)

        # Model path selection
        model_path_layout = QHBoxLayout()  # type: ignore
        model_path_layout.addWidget(QLabel("Model Path:"))
        self.model_path_input = QLineEdit()

        model_path_layout.addWidget(self.model_path_input)
        self.load_model_btn = QPushButton("Browse...")
        self.load_model_btn.clicked.connect(self.on_load_model)  # type: ignore
        model_path_layout.addWidget(self.load_model_btn)
        config_layout.addLayout(model_path_layout)

        # Server parameters
        params_layout = QHBoxLayout()  # type: ignore
        params_layout.addStretch(1)
        params_layout.addWidget(QLabel("Port:"))
        self.port_spin = QSpinBox()
        self.port_spin.setRange(*conf.LISTEN_PORT_RNG)
        self.port_spin.setValue(conf.LISTEN_PORT)
        params_layout.addWidget(self.port_spin)
        self.port_spin.valueChanged.connect(self.on_config_changed)  # type: ignore

        params_layout.addWidget(QLabel("Threads:"))
        self.threads_spin = QSpinBox()
        self.threads_spin.setRange(*conf.THREAD_COUNT_RNG)
        self.threads_spin.setValue(conf.THREAD_COUNT)
        params_layout.addWidget(self.threads_spin)
        config_layout.addLayout(params_layout)
        self.threads_spin.valueChanged.connect(self.on_config_changed)  # type: ignore

        config_group.setLayout(config_layout)
        main_layout.addWidget(config_group)

        # Control buttons
        control_layout = QHBoxLayout()  # type: ignore
        self.start_btn = QPushButton("Start Server")
        self.start_btn.clicked.connect(self.toggle_server)  # type: ignore
        self.browser_btn = QPushButton("Start Browser")
        self.browser_btn.setEnabled(False)
        self.browser_btn.clicked.connect(self.on_start_browser)  # type: ignore
        control_layout.addWidget(self.start_btn)
        control_layout.addWidget(self.browser_btn)
        main_layout.addLayout(control_layout)

        # Output display
        output_group = QGroupBox("Server Output")
        output_layout = QVBoxLayout(output_group)
        self.output_area = QTextEdit("")
        self.output_area.setReadOnly(True)
        self.output_area.setLineWrapMode(QTextEdit.NoWrap)  # type: ignore

        # Set colors for different message types
        self.error_format = self.create_text_format(self.COLOR_ERROR)
        self.warning_format = self.create_text_format(self.COLOR_WARNING)
        self.info_format = self.create_text_format(self.COLOR_INFO)

        output_layout.addWidget(self.output_area)
        output_group.setLayout(output_layout)
        main_layout.addWidget(output_group)

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

    @staticmethod
    def create_text_format(color: QColor) -> QTextCharFormat:
        """Create text format.

        Args:
            color: Text color.

        Returns:
            Text format.
        """
        text_format = QTextCharFormat()
        text_format.setForeground(QBrush(color))
        return text_format

    def setup_process(self) -> None:
        """Initialize process."""
        self.process = QProcess(self)
        self.process.readyReadStandardOutput.connect(self.handle_stdout)  # type: ignore
        self.process.readyReadStandardError.connect(self.handle_stderr)  # type: ignore
        self.process.finished.connect(self.on_process_finished)  # type: ignore

    def apply_config_to_ui(self) -> None:
        """Apply loaded configuration to UI controls."""
        # Set model path
        model_path = conf.MODEL_PATH
        if model_path:
            self.model_path_input.setText(str(model_path))
        else:
            self.model_path_input.setPlaceholderText("Choose model file...")

        # Set port and thread values
        self.port_spin.setValue(conf.LISTEN_PORT)
        self.threads_spin.setValue(conf.THREAD_COUNT)

    def on_config_changed(self) -> None:
        """Configuration changed."""
        conf.MODEL_PATH = self.model_path_input.text().strip()
        conf.LISTEN_PORT = self.port_spin.value()
        conf.THREAD_COUNT = self.threads_spin.value()
        # Configuration will be saved on exit via atexit.register

    def on_load_model(self) -> None:
        """Select model file."""
        initial_dir = (
            conf.MODEL_PATH if os.path.exists(conf.MODEL_PATH) else str(Path.home())
        )
        path, _ = QFileDialog.getOpenFileName(
            self,
            "Select Model File",
            initial_dir,
            "Model Files (*.bin *.gguf)",
        )

        if path:
            conf.MODEL_PATH = path
            self.model_path_input.setText(os.path.normpath(path))

    def toggle_server(self) -> None:
        """Start or stop server."""
        if self.process.state() == QProcess.Running:
            self.stop_server()
        else:
            self.start_server()

    def start_server(self) -> None:
        """Start server."""
        model_path = pathlib.Path(self.model_path_input.text().strip())
        if not model_path.exists():
            self.append_output(
                "Error: Invalid model file path\n",
                self.error_format,
            )
            return

        os.chdir(str(model_path.parent))
        cmd = [
            "llama-server",
            "--model",
            model_path.name,
            "--port",
            str(self.port_spin.value()),
            "--threads",
            str(self.threads_spin.value()),
        ]

        self.append_output(f"Start: {' '.join(cmd)}\n", self.info_format)

        try:
            self.process.start(cmd[0], cmd[1:])
            self.update_ui_state(running=True)
        except Exception as e:
            self.append_output(f"Start failed: {e!s}\n", self.error_format)

    def stop_server(self) -> None:
        """Stop server."""
        if self.process.state() == QProcess.Running:
            self.append_output("Stopping server...\n", self.info_format)
            self.process.terminate()
            if not self.process.waitForFinished(self.PROCESS_TERMINATE_TIMEOUT_MS):
                self.process.kill()

    @staticmethod
    def on_start_browser() -> None:
        """Start browser."""
        QDesktopServices.openUrl(QUrl(f"{conf.URL}:{conf.LISTEN_PORT}"))

    def on_process_finished(self, exit_code: int, exit_status: int) -> None:
        """Process finished."""
        self.append_output(
            f"\nServer stopped, Exit code: {exit_code}, Status: {exit_status}\n",
            self.info_format,
        )
        self.update_ui_state(running=False)

    def handle_stdout(self) -> None:
        """Handle standard output."""
        data = self.process.readAllStandardOutput()
        text = QTextStream(data).readAll()
        self.append_output(text, self.info_format)

    def handle_stderr(self) -> None:
        """Handle standard error."""
        data = self.process.readAllStandardError()
        text = QTextStream(data).readAll()
        self.append_output(text, self.error_format)

    def append_output(
        self,
        text: str,
        text_format: QTextCharFormat | None = None,
    ) -> None:
        """Append output."""
        cursor: QTextCursor = self.output_area.textCursor()
        cursor.movePosition(QTextCursor.MoveOperation.End)

        if text_format:
            cursor.setCharFormat(text_format)

        cursor.insertText(text)

        self.output_area.setTextCursor(cursor)
        self.output_area.ensureCursorVisible()

    def update_ui_state(self, *, running: bool) -> None:
        """Update UI state."""
        self.model_path_input.setEnabled(not running)
        self.load_model_btn.setEnabled(not running)
        self.port_spin.setEnabled(not running)
        self.threads_spin.setEnabled(not running)
        self.browser_btn.setEnabled(running)

        if running:
            self.start_btn.setText("Stop Server")
        else:
            self.start_btn.setText("Start Server")

    def moveEvent(self, event: QMoveEvent) -> None:  # noqa: N802
        """Handle window move event."""
        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
        """Handle window resize event."""
        geometry = self.geometry()
        conf.WIN_SIZE = [geometry.width(), geometry.height()]
        return super().resizeEvent(event)

    def closeEvent(self, event) -> None:  # noqa: N802
        """Handle window close event to ensure configuration is saved."""
        try:
            # Save current configuration
            conf.save()
            logger.info("Configuration saved successfully on exit")
        except Exception as e:
            logger.error(f"Failed to save configuration on exit: {e}")
        finally:
            # Also stop server if running
            if hasattr(self, "process") and self.process.state() == QProcess.Running:
                self.process.terminate()
                self.process.waitForFinished(2000)  # Wait up to 2 seconds
            event.accept()


def main() -> None:
    """Main entry point."""
    app = QApplication(sys.argv)
    app.setFont(QFont("Consolas", 12))
    window = LlamaServerGUI()
    window.show()
    sys.exit(app.exec_())
