import sys
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QTextEdit, QPushButton, QComboBox,
                             QLabel, QFrame)
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from PyQt5.QtGui import QTextCursor
from openai import OpenAI
import bilbystats as bs


class StreamWorker(QThread):
    chunk_received = pyqtSignal(str)
    finished = pyqtSignal(str)
    error = pyqtSignal(str)
    
    def __init__(self, content, instructions, model_name, mymessages):
        super().__init__()
        self.content = content
        self.instructions = instructions
        self.model_name = model_name
        self.mymessages = mymessages
    
    def run(self):
        try:
            full_response = ""
            
            # All calls go through OpenRouter
            client = OpenAI(
                base_url="https://openrouter.ai/api/v1",
                api_key=bs.read_api_key('openrouter')
            )
            
            # Build messages list
            messages = []
            if self.instructions:
                messages.append({"role": "system", "content": self.instructions})
            if self.mymessages:
                messages.extend(self.mymessages)
            messages.append({"role": "user", "content": self.content})
            
            completion = client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                stream=True
            )
            
            for chunk in completion:
                if chunk.choices[0].delta.content:
                    text = chunk.choices[0].delta.content
                    full_response += text
                    self.chunk_received.emit(text)
            
            self.finished.emit(full_response)
            
        except Exception as e:
            self.error.emit(str(e))


class ChatWindow(QMainWindow):
    def __init__(self, instructions="You are a helpful assistant.", model_name="openai/gpt-4o"):
        super().__init__()
        self.instructions = instructions
        self.mymessages = []
        self.worker = None
        self.pending_user_input = ""
        
        # Load models from common_models.json
        model_data = bs.read_data('common_models.json')
        self.all_models = model_data['chat_models']
        
        # Parse providers and models
        self.provider_models = {}
        for model in self.all_models:
            if '/' in model:
                provider, model_short = model.split('/', 1)
                if provider not in self.provider_models:
                    self.provider_models[provider] = []
                self.provider_models[provider].append(model_short)
        
        self.providers = sorted(self.provider_models.keys())
        
        # Set initial provider and model
        if '/' in model_name:
            self.current_provider, self.current_model = model_name.split('/', 1)
        else:
            self.current_provider = self.providers[0] if self.providers else "openai"
            self.current_model = model_name
        
        self.init_ui()
        self.apply_styles()
    
    def init_ui(self):
        self.setWindowTitle("LLM Chat")
        self.setGeometry(100, 100, 800, 700)
        
        central = QWidget()
        self.setCentralWidget(central)
        layout = QVBoxLayout(central)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        
        # Header
        header = QFrame()
        header.setObjectName("header")
        header_layout = QHBoxLayout(header)
        header_layout.setContentsMargins(20, 15, 20, 15)
        
        title = QLabel("LLM Chat")
        title.setObjectName("title")
        header_layout.addWidget(title)
        
        header_layout.addStretch()
        
        # Provider selector
        provider_label = QLabel("Provider:")
        provider_label.setObjectName("selectorLabel")
        header_layout.addWidget(provider_label)
        
        self.provider_combo = QComboBox()
        self.provider_combo.addItems(self.providers)
        if self.current_provider in self.providers:
            self.provider_combo.setCurrentText(self.current_provider)
        self.provider_combo.currentTextChanged.connect(self.change_provider)
        self.provider_combo.setObjectName("selectorCombo")
        header_layout.addWidget(self.provider_combo)
        
        header_layout.addSpacing(15)
        
        # Model selector
        model_label = QLabel("Model:")
        model_label.setObjectName("selectorLabel")
        header_layout.addWidget(model_label)
        
        self.model_combo = QComboBox()
        self.model_combo.setObjectName("selectorCombo")
        self.update_model_combo()
        self.model_combo.currentTextChanged.connect(self.change_model)
        header_layout.addWidget(self.model_combo)
        
        layout.addWidget(header)
        
        # Chat area
        chat_container = QFrame()
        chat_container.setObjectName("chatContainer")
        chat_layout = QVBoxLayout(chat_container)
        chat_layout.setContentsMargins(20, 20, 20, 20)
        
        self.chat_display = QTextEdit()
        self.chat_display.setReadOnly(True)
        self.chat_display.setObjectName("chatDisplay")
        chat_layout.addWidget(self.chat_display)
        
        layout.addWidget(chat_container, 1)
        
        # Input area
        input_container = QFrame()
        input_container.setObjectName("inputContainer")
        input_layout = QHBoxLayout(input_container)
        input_layout.setContentsMargins(20, 15, 20, 20)
        input_layout.setSpacing(12)
        
        self.input_box = QTextEdit()
        self.input_box.setObjectName("inputBox")
        self.input_box.setPlaceholderText("Type your message here... (Enter to send, Shift+Enter for new line)")
        self.input_box.setMaximumHeight(100)
        self.input_box.installEventFilter(self)
        input_layout.addWidget(self.input_box, 1)
        
        btn_container = QWidget()
        btn_layout = QVBoxLayout(btn_container)
        btn_layout.setContentsMargins(0, 0, 0, 0)
        btn_layout.setSpacing(8)
        
        self.send_btn = QPushButton("Send")
        self.send_btn.setObjectName("sendBtn")
        self.send_btn.clicked.connect(self.send_message)
        self.send_btn.setMinimumWidth(80)
        btn_layout.addWidget(self.send_btn)
        
        self.clear_btn = QPushButton("Clear")
        self.clear_btn.setObjectName("clearBtn")
        self.clear_btn.clicked.connect(self.clear_chat)
        self.clear_btn.setMinimumWidth(80)
        btn_layout.addWidget(self.clear_btn)
        
        btn_layout.addStretch()
        input_layout.addWidget(btn_container)
        
        layout.addWidget(input_container)
    
    def apply_styles(self):
        self.setStyleSheet("""
            QMainWindow {
                background-color: #1a1a2e;
            }
            #header {
                background-color: #16213e;
                border-bottom: 1px solid #0f3460;
            }
            #title {
                color: #e94560;
                font-size: 22px;
                font-weight: bold;
            }
            #selectorLabel {
                color: #a0a0a0;
                font-size: 13px;
            }
            #selectorCombo {
                background-color: #0f3460;
                color: #ffffff;
                border: 1px solid #e94560;
                border-radius: 5px;
                padding: 5px 10px;
                min-width: 150px;
            }
            #selectorCombo::drop-down {
                border: none;
            }
            #selectorCombo QAbstractItemView {
                background-color: #16213e;
                color: #ffffff;
                selection-background-color: #e94560;
            }
            #chatContainer {
                background-color: #1a1a2e;
            }
            #chatDisplay {
                background-color: #16213e;
                color: #e0e0e0;
                border: 1px solid #0f3460;
                border-radius: 10px;
                padding: 15px;
                font-size: 14px;
                line-height: 1.5;
            }
            #inputContainer {
                background-color: #16213e;
                border-top: 1px solid #0f3460;
            }
            #inputBox {
                background-color: #1a1a2e;
                color: #ffffff;
                border: 1px solid #0f3460;
                border-radius: 8px;
                padding: 10px;
                font-size: 14px;
            }
            #inputBox:focus {
                border-color: #e94560;
            }
            #sendBtn {
                background-color: #e94560;
                color: #ffffff;
                border: none;
                border-radius: 6px;
                padding: 10px 20px;
                font-size: 14px;
                font-weight: bold;
            }
            #sendBtn:hover {
                background-color: #ff6b6b;
            }
            #sendBtn:pressed {
                background-color: #c73e54;
            }
            #sendBtn:disabled {
                background-color: #555555;
            }
            #clearBtn {
                background-color: transparent;
                color: #a0a0a0;
                border: 1px solid #0f3460;
                border-radius: 6px;
                padding: 10px 20px;
                font-size: 14px;
            }
            #clearBtn:hover {
                border-color: #e94560;
                color: #e94560;
            }
        """)
    
    def update_model_combo(self):
        self.model_combo.blockSignals(True)
        self.model_combo.clear()
        models = self.provider_models.get(self.current_provider, [])
        self.model_combo.addItems(models)
        if self.current_model in models:
            self.model_combo.setCurrentText(self.current_model)
        elif models:
            self.current_model = models[0]
        self.model_combo.blockSignals(False)
    
    def change_provider(self, provider):
        self.current_provider = provider
        self.update_model_combo()
    
    def change_model(self, model):
        self.current_model = model
    
    def get_full_model_name(self):
        return f"{self.current_provider}/{self.current_model}"
    
    def eventFilter(self, obj, event):
        if obj == self.input_box and event.type() == event.KeyPress:
            if event.key() == Qt.Key_Return and not (event.modifiers() & Qt.ShiftModifier):
                self.send_message()
                return True
        return super().eventFilter(obj, event)
    
    def send_message(self):
        user_input = self.input_box.toPlainText().strip()
        if not user_input:
            return
        
        self.append_message("You", user_input, "#4fc3f7")
        self.input_box.clear()
        self.send_btn.setEnabled(False)
        self.pending_user_input = user_input
        
        # Add "Assistant:" label before streaming starts
        self.chat_display.append("")
        cursor = self.chat_display.textCursor()
        cursor.movePosition(QTextCursor.End)
        cursor.insertHtml('<span style="color: #81c784; font-weight: bold;">Assistant: </span>')
        self.chat_display.setTextCursor(cursor)
        
        self.worker = StreamWorker(
            user_input, 
            self.instructions, 
            self.get_full_model_name(), 
            self.mymessages
        )
        self.worker.chunk_received.connect(self.handle_chunk)
        self.worker.finished.connect(self.handle_finished)
        self.worker.error.connect(self.handle_error)
        self.worker.start()
    
    def append_message(self, sender, message, color):
        html = f'<p><span style="color: {color}; font-weight: bold;">{sender}:</span> {message}</p>'
        self.chat_display.append(html)
    
    def handle_chunk(self, chunk):
        cursor = self.chat_display.textCursor()
        cursor.movePosition(QTextCursor.End)
        cursor.insertText(chunk)
        self.chat_display.setTextCursor(cursor)
        self.chat_display.ensureCursorVisible()
    
    def handle_finished(self, full_response):
        self.mymessages.append({"role": "user", "content": self.pending_user_input})
        self.mymessages.append({"role": "assistant", "content": full_response})
        self.chat_display.append("")  # Add spacing after response
        self.send_btn.setEnabled(True)
    
    def handle_error(self, error_msg):
        cursor = self.chat_display.textCursor()
        cursor.movePosition(QTextCursor.End)
        cursor.insertHtml(f'<span style="color: #ef5350;">Error: {error_msg}</span>')
        self.chat_display.append("")
        self.send_btn.setEnabled(True)
    
    def clear_chat(self):
        self.mymessages = []
        self.chat_display.clear()


def launch_chat_gui(instructions="You are a helpful assistant.", model_name="openai/gpt-4o"):
    """
    Launch a GUI for having an LLM conversation with memory and streaming.
    
    Parameters
    ----------
    instructions : str
        System-level instructions for the model.
    model_name : str
        Model to use in format "provider/model" (e.g., "openai/gpt-4o", "anthropic/claude-3-opus").
    """
    app = QApplication.instance()
    if app is None:
        app = QApplication(sys.argv)
    
    window = ChatWindow(instructions, model_name)
    window.show()
    window.raise_()
    window.activateWindow()
    app.exec_()


if __name__ == "__main__":
    launch_chat_gui()