import sys
import requests
import json
from PyQt6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QTextEdit, QLineEdit, QLabel, QComboBox
)
from PyQt6.QtCore import QTimer, Qt
from PyQt6.QtGui import QFont, QCursor, QIcon, QColor, QTextCursor

API_BASE = "http://127.0.0.1:7860/api"

class ChatClient(QWidget):
    def __init__(self, token):
        super().__init__()
        self.token = token  # Token comes from main.py login
        self.ai_typing_timer = None
        self.ai_text_buffer = ""
        self.setWindowTitle("AI HUB Desktop Client")
        self.setWindowIcon(QIcon("favicon.png"))

        self.init_ui()
        self.showMaximized()
        self.setCursor(QCursor(Qt.CursorShape.ArrowCursor))

    def init_ui(self):
        self.layout = QVBoxLayout(self)
        self.setLayout(self.layout)

        # Chat Display Area
        self.chat_area = QTextEdit()
        self.chat_area.setReadOnly(True)
        self.chat_area.setFont(QFont("Consolas", 12))
        self.chat_area.setStyleSheet("""
            QTextEdit {
                background-color: #1e1e1e;
                color: #f1f1f1;
                border-radius: 10px;
                padding: 10px;
            }
        """)
        self.layout.addWidget(self.chat_area)

        # Typing Indicator
        self.typing_label = QLabel("")
        self.typing_label.setStyleSheet("color: #00ff00; font-style: italic;")
        self.layout.addWidget(self.typing_label)

        # Input Area
        input_layout = QHBoxLayout()
        self.user_input = QLineEdit()
        self.user_input.setPlaceholderText("Type your message here...")
        self.user_input.returnPressed.connect(self.send_message)
        self.user_input.setStyleSheet("""
            QLineEdit {
                border-radius: 8px;
                padding: 6px;
                font-size: 14px;
            }
        """)
        input_layout.addWidget(self.user_input)

        self.send_button = QPushButton("Send")
        self.send_button.setStyleSheet("""
            QPushButton {
                background-color: #0078d7;
                color: white;
                border-radius: 8px;
                padding: 6px 12px;
            }
            QPushButton:hover {
                background-color: #005a9e;
            }
        """)
        self.send_button.clicked.connect(self.send_message)
        input_layout.addWidget(self.send_button)

        # Model selection
        self.model_select = QComboBox()
        self.model_select.addItems(["qwen", "deepseek"])
        self.model_select.setStyleSheet("""
            QComboBox {
                border-radius: 8px;
                padding: 4px;
                font-size: 14px;
            }
        """)
        input_layout.addWidget(self.model_select)

        self.layout.addLayout(input_layout)

    # --------------------
    # Display AI typing animation
    # --------------------
    def type_ai_message(self, full_text, delay=30):
        self.ai_text_buffer = full_text
        self.typing_label.setText("🤖 AI is typing...")
        self.ai_typing_timer = QTimer()
        self.ai_typing_timer.timeout.connect(self._type_next_char)
        self.ai_typing_timer.start(delay)

    def _type_next_char(self):
        if self.ai_text_buffer:
            char = self.ai_text_buffer[0]
            self.chat_area.moveCursor(QTextCursor.MoveOperation.End)
            self.chat_area.insertPlainText(char)
            self.ai_text_buffer = self.ai_text_buffer[1:]
        else:
            self.ai_typing_timer.stop()
            self.typing_label.setText("")

    # --------------------
    # Send message to backend
    # --------------------
    def send_message(self):
        message = self.user_input.text().strip()
        if not message:
            return
        self.chat_area.append(f"🧑 {message}")
        self.user_input.clear()
        self.typing_label.setText("🤖 AI is thinking...")

        model_choice = self.model_select.currentText()
        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.token}"}

        try:
            res = requests.post(
                f"{API_BASE}/chat",
                headers=headers,
                data=json.dumps({"message": message, "model": model_choice, "max_tokens": 200}),
                timeout=30
            )
            data = res.json()
            if "response" in data:
                self.chat_area.append("🤖 ")
                self.type_ai_message(data["response"])
            else:
                self.chat_area.append("Error: " + data.get("error", "Unknown error"))
        except Exception as e:
            self.chat_area.append(f"Error sending message: {e}")

# --------------------
# Run App
# --------------------
if __name__ == "__main__":
    # Example: Pass token from main.py login
    token = "YOUR_JWT_TOKEN_HERE"
    app = QApplication(sys.argv)
    window = ChatClient(token)
    window.show()
    sys.exit(app.exec())