from PyQt5.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel,
    QFrame, QListWidget, QListWidgetItem, QTableWidget,
    QTableWidgetItem, QHeaderView, QSplitter, QAbstractItemView,
    QSizePolicy
)
from PyQt5.QtCore import Qt, QRectF
from PyQt5.QtGui import QColor, QFont, QBrush, QPainter, QPen, QPainterPath

from gui.assets.theme.styler import current_theme
from gui.widgets.help_button import HelpButton


class ModelInfoCard(QFrame):
    def __init__(self, parent=None):
        super().__init__(parent)
        self._model_name = "Select a model"
        self._environment = "-"
        self._agents = "-"
        self._status = "-"
        self._status_color = "#9CA3AF"
        
        self.setMinimumHeight(140)
        self.setMaximumHeight(160)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
    
    def set_data(self, model_name, environment, agents, status, is_complete):
        self._model_name = model_name
        self._environment = environment
        self._agents = agents
        self._status = status
        
        if is_complete:
            self._status_color = current_theme.get('success', '#10B981')
        else:
            self._status_color = current_theme.get('warning', '#F59E0B')
        
        self.update()
    
    def clear(self):
        self._model_name = "Select a model"
        self._environment = "-"
        self._agents = "-"
        self._status = "-"
        self._status_color = "#9CA3AF"
        self.update()
    
    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        rect = self.rect()
        
        bg_color = QColor(current_theme.get('background_secondary', '#1F2937'))
        border_color = QColor(current_theme.get('background_tertiary', '#4B5563'))
        text_primary = QColor(current_theme.get('text_primary', '#F9FAFB'))
        text_secondary = QColor(current_theme.get('text_secondary', '#9CA3AF'))
        accent = QColor(current_theme.get('interactive_primary', '#3B82F6'))
        
        bg_rect = QRectF(rect).adjusted(1, 1, -1, -1)
        path = QPainterPath()
        path.addRoundedRect(bg_rect, 10, 10)
        
        painter.setPen(QPen(border_color, 1))
        painter.setBrush(QBrush(bg_color))
        painter.drawPath(path)
        
        accent_rect = QRectF(bg_rect.left() + 4, bg_rect.top() + 12, 4, bg_rect.height() - 24)
        accent_path = QPainterPath()
        accent_path.addRoundedRect(accent_rect, 2, 2)
        painter.setPen(Qt.NoPen)
        painter.setBrush(QBrush(accent))
        painter.drawPath(accent_path)
        
        content_left = bg_rect.left() + 20
        content_width = bg_rect.width() - 32
        
        title_font = QFont("Segoe UI", 16, QFont.Bold)
        painter.setFont(title_font)
        painter.setPen(text_primary)
        painter.drawText(QRectF(content_left, bg_rect.top() + 16, content_width, 28), 
                        Qt.AlignLeft | Qt.AlignVCenter, self._model_name)
        
        info_font = QFont("Segoe UI", 10)
        label_font = QFont("Segoe UI", 9)
        
        y_info = bg_rect.top() + 52
        col_width = content_width / 2
        
        painter.setFont(label_font)
        painter.setPen(text_secondary)
        painter.drawText(QRectF(content_left, y_info, col_width, 16), 
                        Qt.AlignLeft | Qt.AlignVCenter, "Environment")
        
        painter.setFont(info_font)
        painter.setPen(text_primary)
        painter.drawText(QRectF(content_left, y_info + 18, col_width - 10, 20), 
                        Qt.AlignLeft | Qt.AlignVCenter, self._environment)
        
        col2_left = content_left + col_width
        painter.setFont(label_font)
        painter.setPen(text_secondary)
        painter.drawText(QRectF(col2_left, y_info, col_width, 16), 
                        Qt.AlignLeft | Qt.AlignVCenter, "Training Status")
        
        painter.setFont(info_font)
        painter.setPen(QColor(self._status_color))
        painter.drawText(QRectF(col2_left, y_info + 18, col_width, 20), 
                        Qt.AlignLeft | Qt.AlignVCenter, self._status)
        
        y_agents = y_info + 48
        painter.setFont(label_font)
        painter.setPen(text_secondary)
        painter.drawText(QRectF(content_left, y_agents, content_width, 16), 
                        Qt.AlignLeft | Qt.AlignVCenter, "Connected Agents")
        
        painter.setFont(info_font)
        painter.setPen(text_primary)
        
        fm = painter.fontMetrics()
        agents_text = fm.elidedText(self._agents, Qt.ElideRight, int(content_width))
        painter.drawText(QRectF(content_left, y_agents + 18, content_width, 20), 
                        Qt.AlignLeft | Qt.AlignVCenter, agents_text)
        
        painter.end()


class PageLearning(QWidget):
    def __init__(self, log_store, parent=None):
        super().__init__(parent)
        self.log_store = log_store
        self.setObjectName("PageLearning")
        
        self.current_model = None
        self._last_model_count = 0
        
        self._setup_ui()
        
        self.log_store.store_updated.connect(self.on_store_updated)
    
    def _setup_ui(self):
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(15)

        header_layout = QHBoxLayout()
        title = QLabel("Learning Models")
        title.setProperty("class", "h1")
        title.setFont(QFont("Segoe UI", 20, QFont.Bold))
        
        self.help_button = HelpButton(
            "<div style='max-width: 450px;'>"
            "<h3 style='margin: 0 0 10px 0; color: #3B82F6;'>Learning Models</h3>"
            "<p style='margin: 0 0 10px 0;'>This page displays Q-Learning models used by MASPY agents for reinforcement learning</p>"
            "<p style='margin: 0 0 5px 0;'><b>Model Name:</b> Identifier of the learning model</p>"
            "<p style='margin: 0 0 5px 0;'><b>Environment:</b> The environment where the agent is learning</p>"
            "<p style='margin: 0 0 5px 0;'><b>Connected Agents:</b> Agents using this learning model</p>"
            "<p style='margin: 0 0 10px 0;'><b>Training Status:</b> Whether the training is complete or still in progress</p>"
            "<p style='margin: 0 0 5px 0;'><b>Q-Table (Policy):</b> The learned policy showing state-action values</p>"
            "<p style='margin: 0 0 5px 0;'><b>States (rows):</b> Possible states the agent can be in</p>"
            "<p style='margin: 0 0 5px 0;'><b>Actions (columns):</b> Available actions for each state</p>"
            "<p style='margin: 0 0 5px 0;'><b>Values:</b> Expected reward for each action</p>"
            "</div>",
            icon_size=28
        )
        
        header_layout.addWidget(title)
        header_layout.addWidget(self.help_button)
        header_layout.addStretch()
        main_layout.addLayout(header_layout)

        splitter = QSplitter(Qt.Horizontal)

        left_frame = QFrame()
        left_frame.setProperty("class", "card")
        left_frame.setMinimumWidth(220)
        left_frame.setMaximumWidth(320)
        
        left_layout = QVBoxLayout(left_frame)
        left_layout.setContentsMargins(16, 16, 16, 16)
        left_layout.setSpacing(12)
        
        models_label = QLabel("Available Models")
        models_label.setFont(QFont("Segoe UI", 12, QFont.Bold))
        left_layout.addWidget(models_label)
        
        self.models_list = QListWidget()
        self.models_list.setObjectName("ModelsList")
        self.models_list.setFrameShape(QFrame.NoFrame)
        self.models_list.setSpacing(4)
        self.models_list.currentItemChanged.connect(self._on_model_selected)
        self._style_list()
        left_layout.addWidget(self.models_list)
        
        splitter.addWidget(left_frame)

        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)
        right_layout.setContentsMargins(0, 0, 0, 0)
        right_layout.setSpacing(15)
        
        self.model_info_card = ModelInfoCard()
        right_layout.addWidget(self.model_info_card)

        qtable_frame = QFrame()
        qtable_frame.setProperty("class", "card")
        qtable_layout = QVBoxLayout(qtable_frame)
        qtable_layout.setContentsMargins(16, 16, 16, 16)
        
        qtable_title = QLabel("Q-Table (Policy)")
        qtable_title.setFont(QFont("Segoe UI", 12, QFont.Bold))
        qtable_layout.addWidget(qtable_title)
        
        self.qtable_widget = QTableWidget()
        self.qtable_widget.setObjectName("QTableWidget")
        self.qtable_widget.setAlternatingRowColors(False)
        self.qtable_widget.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.qtable_widget.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.qtable_widget.horizontalHeader().setStretchLastSection(True)
        self.qtable_widget.verticalHeader().setVisible(False)
        self.qtable_widget.setFrameShape(QFrame.NoFrame)
        qtable_layout.addWidget(self.qtable_widget, 1)
        
        right_layout.addWidget(qtable_frame, 1)
        
        splitter.addWidget(right_widget)
        splitter.setStretchFactor(0, 1)
        splitter.setStretchFactor(1, 3)
        
        main_layout.addWidget(splitter, 1)
    
    def _style_list(self):
        bg_primary = current_theme.get('background_primary', '#111827')
        bg_secondary = current_theme.get('background_secondary', '#1F2937')
        bg_tertiary = current_theme.get('background_tertiary', '#4B5563')
        text_primary = current_theme.get('text_primary', '#F9FAFB')
        interactive = current_theme.get('interactive_primary', '#3B82F6')
        
        self.models_list.setStyleSheet(f"""
            QListWidget {{
                background-color: transparent;
                border: none;
                outline: none;
            }}
            QListWidget::item {{
                background-color: {bg_secondary};
                border: 1px solid {bg_tertiary};
                border-radius: 6px;
                padding: 12px 14px;
                margin: 2px 0px;
                color: {text_primary};
                font-size: 11px;
            }}
            QListWidget::item:selected {{
                background-color: {interactive};
                border: 1px solid {interactive};
            }}
            QListWidget::item:hover:!selected {{
                background-color: {bg_tertiary};
                border: 1px solid {bg_tertiary};
            }}
        """)
    
    def on_store_updated(self):
        learning_models = self.log_store.get_learning_models()
        
        if len(learning_models) != self._last_model_count:
            self._update_models_list(learning_models)
            self._last_model_count = len(learning_models)

        if self.current_model and self.current_model in learning_models:
            self._update_model_details(self.current_model, learning_models[self.current_model])
    
    def _update_models_list(self, learning_models):
        current_selection = self.current_model
        
        self.models_list.clear()
        
        for model_name in sorted(learning_models.keys()):
            item = QListWidgetItem(model_name)
            item.setData(Qt.UserRole, model_name)
            self.models_list.addItem(item)
            
            if model_name == current_selection:
                item.setSelected(True)
                self.models_list.setCurrentItem(item)
        
        if not current_selection and self.models_list.count() > 0:
            self.models_list.setCurrentRow(0)
    
    def _on_model_selected(self, current, previous):
        if not current:
            self.current_model = None
            self._clear_details()
            return
        
        model_name = current.data(Qt.UserRole)
        self.current_model = model_name
        
        learning_models = self.log_store.get_learning_models()
        if model_name in learning_models:
            self._update_model_details(model_name, learning_models[model_name])
    
    def _update_model_details(self, model_name, model_data):
        environment = model_data.get('environment', '-')
        
        agents = model_data.get('agents', [])
        if agents:
            agents_str = ', '.join(sorted(agents))
        else:
            agents_str = 'No agents connected'
        
        is_complete = model_data.get('training_complete', False)
        status = "Training Complete" if is_complete else "In Progress"
        
        self.model_info_card.set_data(model_name, environment, agents_str, status, is_complete)
        self._update_qtable(model_data.get('policy', {}))
    
    def _update_qtable(self, policy):
        self.qtable_widget.clear()
        
        if not policy:
            self.qtable_widget.setRowCount(0)
            self.qtable_widget.setColumnCount(1)
            self.qtable_widget.setHorizontalHeaderLabels(["No data"])
            return

        first_state = next(iter(policy.values()))
        num_actions = len(first_state) if isinstance(first_state, list) else 1

        columns = ["State"] + [f"Action {i}" for i in range(num_actions)]
        self.qtable_widget.setColumnCount(len(columns))
        self.qtable_widget.setHorizontalHeaderLabels(columns)

        row_color_1 = QColor(current_theme.get('background_secondary', '#1F2937'))
        row_color_2 = QColor(current_theme.get('background_tertiary', '#4B5563'))
        text_color = QColor(current_theme.get('text_primary', '#F9FAFB'))

        best_value_color = QColor(current_theme.get('success', '#10B981'))
        best_value_color.setAlpha(100)

        self.qtable_widget.setRowCount(len(policy))
        
        for row, (state, values) in enumerate(sorted(policy.items())):
            row_bg = row_color_1 if row % 2 == 0 else row_color_2

            state_item = QTableWidgetItem(str(state))
            state_item.setTextAlignment(Qt.AlignCenter)
            state_item.setBackground(QBrush(row_bg))
            state_item.setForeground(QBrush(text_color))
            self.qtable_widget.setItem(row, 0, state_item)

            if isinstance(values, list):
                max_val = max(values)
                for col, val in enumerate(values):
                    val_item = QTableWidgetItem(f"{val:.4f}")
                    val_item.setTextAlignment(Qt.AlignCenter)
                    val_item.setForeground(QBrush(text_color))

                    if val == max_val:
                        val_item.setBackground(QBrush(best_value_color))
                        val_item.setFont(QFont("Segoe UI", 10, QFont.Bold))
                    else:
                        val_item.setBackground(QBrush(row_bg))
                    
                    self.qtable_widget.setItem(row, col + 1, val_item)
            else:
                val_item = QTableWidgetItem(str(values))
                val_item.setTextAlignment(Qt.AlignCenter)
                val_item.setBackground(QBrush(row_bg))
                val_item.setForeground(QBrush(text_color))
                self.qtable_widget.setItem(row, 1, val_item)

        self.qtable_widget.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeToContents)
        for i in range(1, len(columns)):
            self.qtable_widget.horizontalHeader().setSectionResizeMode(i, QHeaderView.Stretch)

        self._apply_header_style()
    
    def _apply_header_style(self):
        header = self.qtable_widget.horizontalHeader()
        bg_tertiary = current_theme.get('background_tertiary', '#4B5563')
        text_primary = current_theme.get('text_primary', '#F9FAFB')
        interactive = current_theme.get('interactive_primary', '#3B82F6')
        
        header.setStyleSheet(f"""
            QHeaderView::section {{
                background-color: {bg_tertiary};
                color: {text_primary};
                padding: 8px;
                border: none;
                border-bottom: 2px solid {interactive};
                font-weight: bold;
            }}
        """)
    
    def _clear_details(self):
        self.model_info_card.clear()
        self.qtable_widget.clear()
        self.qtable_widget.setRowCount(0)
        self.qtable_widget.setColumnCount(0)
    
    def update_theme(self):
        self._style_list()
        self.model_info_card.update()
        
        if self.current_model:
            learning_models = self.log_store.get_learning_models()
            if self.current_model in learning_models:
                self._update_model_details(self.current_model, learning_models[self.current_model])