"""LLM Chat client application.

Provides a graphical interface client for streaming conversations with LLM servers.
Supports real-time streaming response display, connection testing, and parameter adjustment.
"""

from __future__ import annotations

import atexit
import json
import logging
import sys
from codecs import getincrementaldecoder
from dataclasses import dataclass
from pathlib import Path
from typing import ClassVar
from urllib.error import URLError
from urllib.request import Request, urlopen

from PySide2.QtCore import Qt, QThread, Signal
from PySide2.QtGui import QMoveEvent, QResizeEvent, QTextCursor
from PySide2.QtWidgets import (
    QApplication,
    QDoubleSpinBox,
    QGroupBox,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QMainWindow,
    QPushButton,
    QSpinBox,
    QTextEdit,
    QVBoxLayout,
    QWidget,
)

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

CONNECTION_TIMEOUT = 5


@dataclass
class LLMClientConfig:
    """LLM Chat client configuration."""

    TITLE: str = "Llama Local Model Tool"
    WIN_SIZE: ClassVar[list[int]] = [800, 600]
    WIN_POS: ClassVar[list[int]] = [100, 100]
    SERVER_URL: str = "http://localhost:8080"
    MAX_TOKENS: int = 256
    TEMPERATURE: float = 0.7
    TOP_P: float = 0.9
    TOP_K: int = 40
    MAX_TOKENS_RANGE: ClassVar[list[int]] = [1, 4096]
    TEMPERATURE_RANGE: ClassVar[list[float]] = [0.0, 2.0]
    TOP_P_RANGE: ClassVar[list[float]] = [0.0, 1.0]
    TOP_K_RANGE: ClassVar[list[int]] = [1, 100]
    _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.warning("Failed to load configuration: %s", e)
                logger.info("Using default configuration")
        else:
            logger.info("Using default configuration")

    def save(self) -> None:
        """Save configuration."""
        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("_") and attr_name not in [
                "WIN_SIZE",
                "WIN_POS",
                "MAX_TOKENS_RANGE",
                "TEMPERATURE_RANGE",
                "TOP_P_RANGE",
                "TOP_K_RANGE",
            ]:
                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 = LLMClientConfig()
atexit.register(conf.save)


class LLMWorker(QThread):
    """LLM server communication worker thread.

    Handles HTTP streaming requests in background thread to avoid blocking UI main thread.
    Uses incremental UTF-8 decoder to correctly handle multi-byte characters across lines,
    preventing garbled text from truncated multi-byte characters.

    Signals:
        response_received: Emitted when response content is received, carries response text
        error_occurred: Emitted when error occurs, carries error message
        is_finished: Emitted when request completes
    """

    response_received = Signal(str)
    error_occurred = Signal(str)
    is_finished = Signal()

    def __init__(
        self,
        prompt: str,
        server_url: str,
        max_tokens: int,
        temperature: float,
        top_p: float,
        top_k: int,
    ) -> None:
        """Initialize LLM worker thread.

        Args:
            prompt: User input prompt text
            server_url: LLM server address
            max_tokens: Maximum number of tokens to generate
            temperature: Temperature parameter controlling randomness (0.0-2.0)
            top_p: Nucleus sampling parameter (0.0-1.0)
            top_k: Number of candidate tokens to retain
        """
        super().__init__()
        self.prompt = prompt
        self.server_url = server_url
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self._is_running = True

    def run(self) -> None:
        """Execute streaming HTTP request and process response.

        Receives streaming response using Server-Sent Events (SSE) format.
        Uses incremental UTF-8 decoder to avoid garbled text from truncated multi-byte characters,
        ensuring full-chain character encoding consistency (request/response both use UTF-8).
        """
        try:
            headers = {"Content-Type": "application/json; charset=utf-8"}
            data = {
                "prompt": self.prompt,
                "max_tokens": self.max_tokens,
                "temperature": self.temperature,
                "top_p": self.top_p,
                "top_k": self.top_k,
                "stream": True,
            }

            request = Request(
                f"{self.server_url}/completion",
                data=json.dumps(data, ensure_ascii=False).encode("utf-8"),
                headers=headers,
            )

            with urlopen(request) as response:
                if response.status != 200:
                    error_text = response.read().decode("utf-8")
                    self.error_occurred.emit(
                        f"Error: {response.status} - {error_text}",
                    )
                    return

                decoder = getincrementaldecoder("utf-8")(errors="replace")
                buffer = ""

                for line in response:
                    if not self._is_running:
                        break

                    if line:
                        try:
                            decoded_line = decoder.decode(line, False).strip()
                            if not decoded_line:
                                continue

                            if decoded_line.startswith("data: "):
                                json_str = decoded_line[6:]
                                try:
                                    json_data = json.loads(json_str)
                                    content = json_data.get("content", "")
                                    if content:
                                        buffer += content
                                        self.response_received.emit(buffer)
                                except json.JSONDecodeError as e:
                                    logger.debug(
                                        f"JSON parsing failed: {json_str}, error: {e}"
                                    )
                                    continue
                        except Exception as e:
                            logger.debug(f"Line processing failed: {e}")
                            continue

                decoder.decode(b"", True)

        except URLError as e:
            logger.error(f"Connection error: {e.reason}")
            self.error_occurred.emit(f"Connection error: {e.reason}")
        except Exception as e:
            logger.error(f"Request error: {e!s}")
            self.error_occurred.emit(f"Request error: {e!s}")
        finally:
            self.is_finished.emit()

    def stop(self) -> None:
        """Stop worker thread execution."""
        self._is_running = False


class ConnectionTestWorker(QThread):
    """Server connection test worker thread.

    Tests whether the LLM server's health check endpoint is accessible in background thread.

    Signals:
        result_ready: Emitted when test completes, passes (success_flag, message_text)
    """

    result_ready = Signal(bool, str)

    def __init__(self, server_url: str) -> None:
        """Initialize connection test worker thread.

        Args:
            server_url: LLM server address
        """
        super().__init__()
        self.server_url = server_url

    def run(self) -> None:
        """Execute connection test by accessing server health check endpoint."""
        try:
            request = Request(f"{self.server_url}/health")
            with urlopen(request, timeout=CONNECTION_TIMEOUT) as response:
                if response.status == 200:
                    self.result_ready.emit(True, "Connection successful!")
                else:
                    self.result_ready.emit(
                        False, f"Connection failed: {response.status}"
                    )
        except URLError as e:
            logger.error(f"Connection error: {e.reason}")
            self.result_ready.emit(False, f"Connection error: {e.reason}")
        except Exception as e:
            logger.error(f"Request error: {e!s}")
            self.result_ready.emit(False, f"Request error: {e!s}")


class LLMChatApp(QMainWindow):
    """LLM Chat client main window.

    Provides graphical interface for interacting with LLM server, supports:
    - Real-time streaming response display (uses incremental updates to avoid repeated rendering)
    - Connection testing
    - Model parameter adjustment (temperature, top-p, top-k, etc.)
    - Conversation history tracking
    """

    def __init__(self) -> None:
        """Initialize LLM Chat main window."""
        super().__init__()
        self.setWindowTitle(conf.TITLE)
        self.setGeometry(*conf.WIN_POS, *conf.WIN_SIZE)

        self.init_ui()

        self.worker_thread: LLMWorker | None = None
        self.test_thread: ConnectionTestWorker | None = None
        self.current_ai_start_pos = -1

    def init_ui(self) -> None:
        """Initialize user interface components."""
        main_widget = QWidget()
        main_layout = QVBoxLayout()
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

        # 构建界面各部分
        main_layout.addWidget(self._create_server_group())
        main_layout.addWidget(self._create_params_group())
        main_layout.addWidget(self._create_chat_display())
        main_layout.addLayout(self._create_input_layout())

        self.statusBar().showMessage("Ready")

    def _create_server_group(self) -> QGroupBox:
        """Create server settings group.

        Returns:
            Group box containing server address input and test connection button
        """
        server_group = QGroupBox("Server Settings")
        server_layout = QHBoxLayout()

        self.server_url_input = QLineEdit(conf.SERVER_URL)
        self.server_url_input.setPlaceholderText("Enter llama-server address")
        self.server_url_input.textChanged.connect(self.on_config_changed)

        self.test_connection_btn = QPushButton("Test Connection")
        self.test_connection_btn.clicked.connect(self.test_connection)

        server_layout.addWidget(QLabel("Server Address:"))
        server_layout.addWidget(self.server_url_input)
        server_layout.addWidget(self.test_connection_btn)
        server_group.setLayout(server_layout)

        return server_group

    def _create_params_group(self) -> QGroupBox:
        """Create model parameter settings group.

        Returns:
            Group box containing all model parameter adjustment controls
        """
        params_group = QGroupBox("Model Parameters")
        params_layout = QHBoxLayout()

        self.max_tokens_spin = QSpinBox()
        self.max_tokens_spin.setRange(*conf.MAX_TOKENS_RANGE)
        self.max_tokens_spin.setValue(conf.MAX_TOKENS)
        self.max_tokens_spin.valueChanged.connect(self.on_config_changed)

        self.temperature_spin = QDoubleSpinBox()
        self.temperature_spin.setRange(*conf.TEMPERATURE_RANGE)
        self.temperature_spin.setSingleStep(0.1)
        self.temperature_spin.setValue(conf.TEMPERATURE)
        self.temperature_spin.valueChanged.connect(self.on_config_changed)

        self.top_p_spin = QDoubleSpinBox()
        self.top_p_spin.setRange(*conf.TOP_P_RANGE)
        self.top_p_spin.setSingleStep(0.05)
        self.top_p_spin.setValue(conf.TOP_P)
        self.top_p_spin.valueChanged.connect(self.on_config_changed)

        self.top_k_spin = QSpinBox()
        self.top_k_spin.setRange(*conf.TOP_K_RANGE)
        self.top_k_spin.setValue(conf.TOP_K)
        self.top_k_spin.valueChanged.connect(self.on_config_changed)

        params_layout.addWidget(QLabel("Max Tokens:"))
        params_layout.addWidget(self.max_tokens_spin)
        params_layout.addWidget(QLabel("Temperature:"))
        params_layout.addWidget(self.temperature_spin)
        params_layout.addWidget(QLabel("Top P:"))
        params_layout.addWidget(self.top_p_spin)
        params_layout.addWidget(QLabel("Top K:"))
        params_layout.addWidget(self.top_k_spin)
        params_group.setLayout(params_layout)

        return params_group

    def _create_chat_display(self) -> QTextEdit:
        """Create chat display area.

        Returns:
            Read-only text display widget
        """
        self.chat_display = QTextEdit()
        self.chat_display.setReadOnly(True)
        self.chat_display.setStyleSheet("font-family: monospace;")
        return self.chat_display

    def _create_input_layout(self) -> QHBoxLayout:
        """Create user input area layout.

        Returns:
            Layout containing input box and send/stop buttons
        """
        input_layout = QHBoxLayout()

        self.user_input = QLineEdit()
        self.user_input.setPlaceholderText("Enter your prompt...")
        self.user_input.returnPressed.connect(self.send_prompt)

        self.send_btn = QPushButton("Send")
        self.send_btn.clicked.connect(self.send_prompt)

        self.stop_btn = QPushButton("Stop")
        self.stop_btn.clicked.connect(self.stop_generation)
        self.stop_btn.setEnabled(False)

        input_layout.addWidget(self.user_input)
        input_layout.addWidget(self.send_btn)
        input_layout.addWidget(self.stop_btn)

        return input_layout

    def on_config_changed(self) -> None:
        """Configuration changed handler."""
        conf.SERVER_URL = self.server_url_input.text().strip()
        conf.MAX_TOKENS = self.max_tokens_spin.value()
        conf.TEMPERATURE = self.temperature_spin.value()
        conf.TOP_P = self.top_p_spin.value()
        conf.TOP_K = self.top_k_spin.value()

    def test_connection(self) -> None:
        """Test connection with LLM server.

        Sends health check request in background thread to avoid blocking UI.
        """
        server_url = self.server_url_input.text().strip()
        if not server_url:
            self.statusBar().showMessage("Please enter server address")
            return

        if self.test_thread and self.test_thread.isRunning():
            self.statusBar().showMessage("Test in progress...")
            return

        self.test_thread = ConnectionTestWorker(server_url)
        self.test_thread.result_ready.connect(self.on_connection_test_result)
        self.test_thread.finished.connect(self.on_test_thread_finished)

        self.statusBar().showMessage("Testing connection...")
        self.test_connection_btn.setEnabled(False)

        self.test_thread.start()

    def on_test_thread_finished(self) -> None:
        """Handle test thread completion event and clean up resources."""
        if self.test_thread:
            self.test_thread.quit()
            self.test_thread.wait()
            self.test_thread = None

    def on_connection_test_result(self, success: bool, message: str) -> None:
        """Handle connection test result.

        Args:
            success: Whether connection succeeded
            message: Result message text
        """
        self.statusBar().showMessage(message)
        self.test_connection_btn.setEnabled(True)

    def send_prompt(self) -> None:
        """Send user input prompt to LLM server.

        Creates worker thread to handle streaming response and updates UI to display conversation.
        Resets AI reply start position to avoid repeated rendering issues.
        """
        if self.worker_thread and self.worker_thread.isRunning():
            self.statusBar().showMessage("Please wait for current request to complete")
            return

        prompt = self.user_input.text().strip()
        if not prompt:
            self.statusBar().showMessage("Please enter prompt")
            return

        self.current_ai_start_pos = -1

        server_url = self.server_url_input.text().strip()
        max_tokens = self.max_tokens_spin.value()
        temperature = self.temperature_spin.value()
        top_p = self.top_p_spin.value()
        top_k = self.top_k_spin.value()

        self._display_user_input(prompt, server_url)
        self.user_input.clear()

        logger.info(f"Sending prompt: {prompt}")
        logger.info(
            f"Parameters: max_tokens={max_tokens}, temperature={temperature}, top_p={top_p}, top_k={top_k}"
        )

        self.worker_thread = LLMWorker(
            prompt=prompt,
            server_url=server_url,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
        )

        self.worker_thread.response_received.connect(self.update_response)
        self.worker_thread.error_occurred.connect(self.handle_error)
        self.worker_thread.is_finished.connect(self.on_finished)

        self.send_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.statusBar().showMessage("Generating response...")

        self.worker_thread.start()

    def _display_user_input(self, prompt: str, server_url: str) -> None:
        """Display user input and target in chat area.

        Args:
            prompt: User input prompt text
            server_url: Target server address
        """
        cursor = self.chat_display.textCursor()
        cursor.movePosition(QTextCursor.End)
        self.chat_display.setTextCursor(cursor)

        self.chat_display.setTextColor(Qt.blue)
        self.chat_display.insertPlainText(f"You: {prompt}\n")

        self.chat_display.setTextColor(Qt.darkGray)
        self.chat_display.insertPlainText(f"[Sending to {server_url}]\n")

        self.chat_display.setTextColor(Qt.black)
        self.chat_display.insertPlainText("AI:")

    def stop_generation(self) -> None:
        """Stop current ongoing response generation."""
        if self.worker_thread and self.worker_thread.isRunning():
            self.worker_thread.stop()
            self.statusBar().showMessage("Generation stopped")

    def update_response(self, text: str) -> None:
        """Update AI response content in chat display area.

        Uses incremental update strategy to replace only AI reply portion of text,
        avoiding performance issues and repeated rendering from frequent redrawing of entire text area.

        Args:
            text: Complete AI generated response text
        """
        cursor = self.chat_display.textCursor()

        if self.current_ai_start_pos == -1:
            cursor.movePosition(QTextCursor.End)
            self.current_ai_start_pos = cursor.position()
            self.chat_display.insertPlainText(f" {text}")
        else:
            cursor.setPosition(self.current_ai_start_pos)
            cursor.movePosition(QTextCursor.End, QTextCursor.KeepAnchor)
            cursor.removeSelectedText()
            self.chat_display.setTextCursor(cursor)
            self.chat_display.insertPlainText(f" {text}")

        self.chat_display.ensureCursorVisible()

    def append_to_chat(self, text: str, *, is_user: bool = False) -> None:
        """Append text to chat area.

        Args:
            text: Text content to append
            is_user: Whether this is a user message (for color setting)
        """
        cursor = self.chat_display.textCursor()
        cursor.movePosition(QTextCursor.End)
        self.chat_display.setTextCursor(cursor)

        if is_user:
            self.chat_display.setTextColor(Qt.blue)
        else:
            self.chat_display.setTextColor(Qt.black)

        self.chat_display.insertPlainText(text + "\n")
        self.chat_display.setTextColor(Qt.black)

    def handle_error(self, error_msg: str) -> None:
        """Handle error messages.

        Args:
            error_msg: Error message text
        """
        self.append_to_chat(f"Error: {error_msg}")
        self.statusBar().showMessage(error_msg)

    def on_finished(self) -> None:
        """Handle response generation completion event, restore UI state and clean up resources."""
        self.send_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.statusBar().showMessage("Generation complete")
        self.append_to_chat("")

        if self.worker_thread:
            self.worker_thread.quit()
            self.worker_thread.wait()
            self.worker_thread = None

    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 main() -> None:
    """Application entry point."""
    app = QApplication(sys.argv)
    app.setStyle("Fusion")

    window = LLMChatApp()
    window.show()

    sys.exit(app.exec_())
