import re
from PyQt5.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
    QFrame, QListWidget, QListWidgetItem, QGraphicsScene,
    QGraphicsView, QFileDialog, QSplitter, QWidget, QCheckBox,
    QGraphicsLineItem, QGraphicsTextItem, QGraphicsRectItem,
    QAbstractItemView, QMessageBox, QSizePolicy
)
from PyQt5.QtCore import Qt, QRectF, QPointF, QLineF
from PyQt5.QtGui import (
    QFont, QColor, QPen, QBrush, QPainter, QPolygonF, QImage, QPainterPath
)

from gui.assets.theme.styler import current_theme


class InfoCard(QFrame):
    def __init__(self, parent=None):
        super().__init__(parent)
        self._agents_count = 0
        self._messages_count = 0
        
        self.setMinimumHeight(70)
        self.setMaximumHeight(80)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
    
    def set_data(self, agents, messages):
        self._agents_count = agents
        self._messages_count = messages
        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'))
        success = QColor(current_theme.get('success', '#10B981'))
        
        bg_rect = QRectF(rect).adjusted(1, 1, -1, -1)
        path = QPainterPath()
        path.addRoundedRect(bg_rect, 8, 8)
        
        painter.setPen(QPen(border_color, 1))
        painter.setBrush(QBrush(bg_color))
        painter.drawPath(path)
        
        col_width = bg_rect.width() / 2
        
        label_font = QFont("Segoe UI", 9)
        value_font = QFont("Segoe UI", 14, QFont.Bold)
        
        x1 = bg_rect.left() + 16
        painter.setFont(label_font)
        painter.setPen(text_secondary)
        painter.drawText(QRectF(x1, bg_rect.top() + 12, col_width, 16), 
                        Qt.AlignLeft | Qt.AlignVCenter, "Selected Agents")
        
        painter.setFont(value_font)
        painter.setPen(accent)
        painter.drawText(QRectF(x1, bg_rect.top() + 32, col_width, 28), 
                        Qt.AlignLeft | Qt.AlignVCenter, str(self._agents_count))
        
        # Coluna 2: Messages
        x2 = bg_rect.left() + col_width
        painter.setFont(label_font)
        painter.setPen(text_secondary)
        painter.drawText(QRectF(x2, bg_rect.top() + 12, col_width, 16), 
                        Qt.AlignLeft | Qt.AlignVCenter, "Messages Found")
        
        painter.setFont(value_font)
        painter.setPen(success)
        painter.drawText(QRectF(x2, bg_rect.top() + 32, col_width, 28), 
                        Qt.AlignLeft | Qt.AlignVCenter, str(self._messages_count))
        
        painter.end()


class MessageDiagramDialog(QDialog):
    def __init__(self, log_store, parent=None):
        super().__init__(parent)
        self.log_store = log_store
        self.selected_agents = []
        self.all_messages = []
        
        self.setWindowTitle("Message Sequence Diagram")
        self.setMinimumSize(1100, 750)
        self.setWindowFlags(
            Qt.Window |
            Qt.WindowMinimizeButtonHint |
            Qt.WindowMaximizeButtonHint |
            Qt.WindowCloseButtonHint
        )
        
        self._apply_dialog_style()
        self._setup_ui()
        self._load_agents()
    
    def _apply_dialog_style(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')
        text_secondary = current_theme.get('text_secondary', '#9CA3AF')
        interactive = current_theme.get('interactive_primary', '#3B82F6')
        
        self.setStyleSheet(f"""
            QDialog {{
                background-color: {bg_primary};
                color: {text_primary};
            }}
            QLabel {{
                color: {text_primary};
            }}
            QLabel[class="text-secondary"] {{
                color: {text_secondary};
            }}
            QPushButton {{
                background-color: {bg_secondary};
                border: 1px solid {bg_tertiary};
                border-radius: 6px;
                padding: 8px 16px;
                color: {text_primary};
                font-size: 12px;
                min-height: 32px;
            }}
            QPushButton:hover {{
                background-color: {bg_tertiary};
                border-color: {interactive};
            }}
            QPushButton:pressed {{
                background-color: {interactive};
            }}
            QPushButton:disabled {{
                opacity: 0.5;
                color: {text_secondary};
            }}
            QPushButton#primaryButton {{
                background-color: {interactive};
                border: none;
            }}
            QPushButton#primaryButton:hover {{
                background-color: #2563EB;
            }}
            QListWidget {{
                background-color: {bg_secondary};
                border: 1px solid {bg_tertiary};
                border-radius: 8px;
                padding: 4px;
                outline: none;
            }}
            QListWidget::item {{
                background-color: transparent;
                border-radius: 4px;
                padding: 8px 12px;
                margin: 2px 0;
            }}
            QListWidget::item:selected {{
                background-color: {interactive};
            }}
            QListWidget::item:hover:!selected {{
                background-color: {bg_tertiary};
            }}
            QGraphicsView {{
                background-color: {bg_secondary};
                border: 1px solid {bg_tertiary};
                border-radius: 8px;
            }}
            QSplitter::handle {{
                background-color: {bg_tertiary};
                width: 2px;
            }}
        """)
    
    def _setup_ui(self):
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(16)
        
        header_layout = QHBoxLayout()
        
        title = QLabel("Message Sequence Diagram")
        title.setFont(QFont("Segoe UI", 20, QFont.Bold))
        
        subtitle = QLabel("Generate visual diagrams of agent communication")
        subtitle.setProperty("class", "text-secondary")
        subtitle.setFont(QFont("Segoe UI", 11))
        
        title_layout = QVBoxLayout()
        title_layout.setSpacing(2)
        title_layout.addWidget(title)
        title_layout.addWidget(subtitle)
        
        header_layout.addLayout(title_layout)
        header_layout.addStretch()
        
        main_layout.addLayout(header_layout)
        
        self.info_card = InfoCard()
        main_layout.addWidget(self.info_card)
        
        splitter = QSplitter(Qt.Horizontal)
        
        left_panel = self._create_left_panel()
        splitter.addWidget(left_panel)
        
        right_panel = self._create_right_panel()
        splitter.addWidget(right_panel)
        
        splitter.setStretchFactor(0, 1)
        splitter.setStretchFactor(1, 4)
        
        main_layout.addWidget(splitter, stretch=1)
        
        actions_layout = QHBoxLayout()
        actions_layout.setSpacing(12)
        
        self.btn_export_png = QPushButton("Export PNG")
        self.btn_export_png.setObjectName("primaryButton")
        self.btn_export_png.clicked.connect(lambda: self._export_diagram("png"))
        self.btn_export_png.setEnabled(False)
        self.btn_export_png.setMinimumWidth(120)
        
        self.btn_export_svg = QPushButton("Export SVG")
        self.btn_export_svg.clicked.connect(lambda: self._export_diagram("svg"))
        self.btn_export_svg.setEnabled(False)
        self.btn_export_svg.setMinimumWidth(120)
        
        self.btn_close = QPushButton("Close")
        self.btn_close.clicked.connect(self.accept)
        self.btn_close.setMinimumWidth(100)
        
        actions_layout.addWidget(self.btn_export_png)
        actions_layout.addWidget(self.btn_export_svg)
        actions_layout.addStretch()
        actions_layout.addWidget(self.btn_close)
        
        main_layout.addLayout(actions_layout)
    
    def _create_left_panel(self):
        panel = QFrame()
        panel.setProperty("class", "card")
        panel.setMinimumWidth(220)
        panel.setMaximumWidth(280)
        layout = QVBoxLayout(panel)
        layout.setContentsMargins(16, 16, 16, 16)
        layout.setSpacing(12)
        
        header = QLabel("Select Agents")
        header.setFont(QFont("Segoe UI", 12, QFont.Bold))
        layout.addWidget(header)
        
        hint = QLabel("Select at least 2 agents to generate diagram")
        hint.setProperty("class", "text-secondary")
        hint.setFont(QFont("Segoe UI", 9))
        hint.setWordWrap(True)
        layout.addWidget(hint)
        
        self.agents_list = QListWidget()
        self.agents_list.setSelectionMode(QAbstractItemView.MultiSelection)
        self.agents_list.itemSelectionChanged.connect(self._on_selection_changed)
        layout.addWidget(self.agents_list, 1)
        
        btns_layout = QHBoxLayout()
        btns_layout.setSpacing(8)
        
        self.btn_select_all = QPushButton("Select All")
        self.btn_select_all.clicked.connect(self._select_all_agents)
        
        self.btn_clear = QPushButton("Clear")
        self.btn_clear.clicked.connect(self._clear_selection)
        
        btns_layout.addWidget(self.btn_select_all)
        btns_layout.addWidget(self.btn_clear)
        layout.addLayout(btns_layout)
        
        return panel
    
    def _create_right_panel(self):
        panel = QFrame()
        panel.setProperty("class", "card")
        layout = QVBoxLayout(panel)
        layout.setContentsMargins(16, 16, 16, 16)
        layout.setSpacing(12)
        
        header_layout = QHBoxLayout()
        
        diagram_title = QLabel("Diagram Preview")
        diagram_title.setFont(QFont("Segoe UI", 12, QFont.Bold))
        
        header_layout.addWidget(diagram_title)
        header_layout.addStretch()
        layout.addLayout(header_layout)
        
        self.scene = QGraphicsScene()
        self.view = QGraphicsView(self.scene)
        self.view.setRenderHint(QPainter.Antialiasing)
        self.view.setRenderHint(QPainter.TextAntialiasing)
        self.view.setDragMode(QGraphicsView.ScrollHandDrag)
        layout.addWidget(self.view, 1)
        
        zoom_layout = QHBoxLayout()
        zoom_layout.setSpacing(8)
        
        zoom_label = QLabel("Zoom")
        zoom_label.setProperty("class", "text-secondary")
        zoom_label.setFont(QFont("Segoe UI", 10))
        
        self.btn_zoom_out = QPushButton("-")
        self.btn_zoom_out.setFixedSize(36, 36)
        self.btn_zoom_out.clicked.connect(lambda: self.view.scale(0.8, 0.8))
        
        self.btn_zoom_in = QPushButton("+")
        self.btn_zoom_in.setFixedSize(36, 36)
        self.btn_zoom_in.clicked.connect(lambda: self.view.scale(1.2, 1.2))
        
        self.btn_zoom_reset = QPushButton("Reset")
        self.btn_zoom_reset.clicked.connect(self._reset_zoom)
        
        self.btn_fit = QPushButton("Fit")
        self.btn_fit.clicked.connect(self._fit_to_view)
        
        zoom_layout.addWidget(zoom_label)
        zoom_layout.addWidget(self.btn_zoom_out)
        zoom_layout.addWidget(self.btn_zoom_in)
        zoom_layout.addWidget(self.btn_zoom_reset)
        zoom_layout.addWidget(self.btn_fit)
        zoom_layout.addStretch()
        layout.addLayout(zoom_layout)
        
        return panel
    
    def _load_agents(self):
        messages = self.log_store.get_all_messages()
        agents = set()
        
        for msg in messages:
            agents.add(msg.sender)
            recv_str = msg.receiver
            if "[" in recv_str:
                found = re.findall(r"[\"']?([\w\.-]+)[\"']?", recv_str)
                for x in found:
                    if x not in ['broadcast', '[]', 'N/A']:
                        agents.add(x)
            elif recv_str not in ['broadcast', '[]', 'N/A', 'Unknown']:
                agents.add(recv_str)
        
        for agent in sorted(agents):
            if agent:
                item = QListWidgetItem(agent)
                self.agents_list.addItem(item)
    
    def _on_selection_changed(self):
        selected_items = self.agents_list.selectedItems()
        self.selected_agents = [item.text() for item in selected_items]
        
        has_agents = len(self.selected_agents) >= 2
        self.btn_export_png.setEnabled(has_agents)
        self.btn_export_svg.setEnabled(has_agents)
        
        messages = self._get_filtered_messages()
        self.info_card.set_data(len(self.selected_agents), len(messages))
        
        self._update_diagram()
    
    def _select_all_agents(self):
        self.agents_list.selectAll()
    
    def _clear_selection(self):
        self.agents_list.clearSelection()
    
    def _reset_zoom(self):
        self.view.resetTransform()
    
    def _fit_to_view(self):
        self.view.fitInView(self.scene.sceneRect(), Qt.KeepAspectRatio)
    
    def _get_filtered_messages(self):
        if len(self.selected_agents) < 2:
            return []
        
        messages = self.log_store.get_all_messages()
        filtered = []
        
        for msg in messages:
            sender = msg.sender
            recv_str = msg.receiver
            
            if sender not in self.selected_agents:
                continue
            
            receivers = []
            if "[" in recv_str:
                found = re.findall(r"[\"']?([\w\.-]+)[\"']?", recv_str)
                receivers = [x for x in found if x in self.selected_agents]
            elif recv_str in self.selected_agents:
                receivers = [recv_str]
            
            for receiver in receivers:
                filtered.append({
                    'sender': sender,
                    'receiver': receiver,
                    'performative': msg.performative,
                    'content': msg.content,
                    'time': msg.system_time
                })
        
        return filtered
    
    def _update_diagram(self):
        self.scene.clear()
        
        if len(self.selected_agents) < 2:
            text = self.scene.addText("Select at least 2 agents to generate the diagram")
            text.setDefaultTextColor(QColor(current_theme.get('text_secondary', '#9CA3AF')))
            text.setFont(QFont("Segoe UI", 12))
            return
        
        messages = self._get_filtered_messages()
        
        if not messages:
            text = self.scene.addText("No messages found between selected agents")
            text.setDefaultTextColor(QColor(current_theme.get('text_secondary', '#9CA3AF')))
            text.setFont(QFont("Segoe UI", 12))
            return
        
        self._draw_sequence_diagram(messages)
    
    def _draw_sequence_diagram(self, messages):
        AGENT_WIDTH = 130
        AGENT_HEIGHT = 44
        AGENT_SPACING = 190
        MESSAGE_SPACING = 55
        TOP_MARGIN = 25
        LEFT_MARGIN = 35
        
        c_bg = QColor(current_theme.get('background_tertiary', '#374151'))
        c_border = QColor(current_theme.get('interactive_primary', '#3B82F6'))
        c_text = QColor(current_theme.get('text_primary', '#F9FAFB'))
        c_text_sec = QColor(current_theme.get('text_secondary', '#9CA3AF'))
        c_line = QColor(current_theme.get('text_disabled', '#6B7280'))
        c_arrow = QColor(current_theme.get('info', '#0EA5E9'))
        
        agents = self.selected_agents
        num_agents = len(agents)
        
        agent_positions = {}
        for i, agent in enumerate(agents):
            x = LEFT_MARGIN + i * AGENT_SPACING
            agent_positions[agent] = x + AGENT_WIDTH / 2
        
        total_height = TOP_MARGIN + AGENT_HEIGHT + 40 + len(messages) * MESSAGE_SPACING + 60
        
        for agent in agents:
            x = agent_positions[agent] - AGENT_WIDTH / 2
            y = TOP_MARGIN
            
            rect_path = QPainterPath()
            rect_path.addRoundedRect(x, y, AGENT_WIDTH, AGENT_HEIGHT, 8, 8)
            
            rect_item = self.scene.addPath(rect_path, QPen(c_border, 2), QBrush(c_bg))
            
            text = self.scene.addText(agent)
            text.setDefaultTextColor(c_text)
            text.setFont(QFont("Segoe UI", 10, QFont.Bold))
            text_width = text.boundingRect().width()
            text.setPos(x + (AGENT_WIDTH - text_width) / 2, y + 10)
        
        line_start_y = TOP_MARGIN + AGENT_HEIGHT
        line_end_y = total_height
        
        pen_lifeline = QPen(c_line, 1.5, Qt.DashLine)
        for agent in agents:
            x = agent_positions[agent]
            line = self.scene.addLine(x, line_start_y, x, line_end_y, pen_lifeline)
        
        current_y = TOP_MARGIN + AGENT_HEIGHT + 45
        
        for msg in messages:
            sender = msg['sender']
            receiver = msg['receiver']
            performative = msg['performative']
            content = msg['content']
            
            if sender not in agent_positions or receiver not in agent_positions:
                continue
            
            x1 = agent_positions[sender]
            x2 = agent_positions[receiver]
            
            going_right = x2 > x1
            
            pen_arrow = QPen(c_arrow, 2)
            line = self.scene.addLine(x1, current_y, x2, current_y, pen_arrow)
            
            arrow_size = 10
            if going_right:
                arrow_x = x2 - arrow_size
                p1 = QPointF(arrow_x, current_y - 5)
                p2 = QPointF(x2, current_y)
                p3 = QPointF(arrow_x, current_y + 5)
            else:
                arrow_x = x2 + arrow_size
                p1 = QPointF(arrow_x, current_y - 5)
                p2 = QPointF(x2, current_y)
                p3 = QPointF(arrow_x, current_y + 5)
            
            arrow_head = self.scene.addPolygon(
                QPolygonF([p1, p2, p3]),
                QPen(c_arrow),
                QBrush(c_arrow)
            )
            
            content_clean = re.sub(r'\s*\[[^\]]*\]\s*$', '', content)
            content_short = content_clean[:35] + "..." if len(content_clean) > 35 else content_clean
            msg_text = f"[{performative}] {content_short}"
            
            text_item = self.scene.addText(msg_text)
            text_item.setDefaultTextColor(c_text_sec)
            text_item.setFont(QFont("Segoe UI", 9))
            
            text_width = text_item.boundingRect().width()
            text_x = min(x1, x2) + abs(x2 - x1) / 2 - text_width / 2
            text_item.setPos(text_x, current_y - 22)
            
            current_y += MESSAGE_SPACING
        
        total_width = LEFT_MARGIN + num_agents * AGENT_SPACING + 50
        self.scene.setSceneRect(0, 0, total_width, total_height)
    
    def _export_diagram(self, format_type):
        if format_type == "png":
            file_filter = "PNG Image (*.png)"
            default_ext = ".png"
        else:
            file_filter = "SVG Image (*.svg)"
            default_ext = ".svg"
        
        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "Export Diagram",
            f"message_diagram{default_ext}",
            file_filter
        )
        
        if not file_path:
            return
        
        if format_type == "png":
            self._export_as_png(file_path)
        else:
            self._export_as_svg(file_path)
    
    def _export_as_png(self, file_path):
        scene_rect = self.scene.sceneRect()
        
        scale = 2
        image = QImage(
            int(scene_rect.width() * scale),
            int(scene_rect.height() * scale),
            QImage.Format_ARGB32
        )
        
        bg_color = QColor(current_theme.get('background_primary', '#111827'))
        image.fill(bg_color)
        
        painter = QPainter(image)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.setRenderHint(QPainter.TextAntialiasing)
        painter.scale(scale, scale)
        self.scene.render(painter)
        painter.end()
        
        if image.save(file_path):
            QMessageBox.information(self, "Export Successful", f"Diagram exported to:\n{file_path}")
        else:
            QMessageBox.warning(self, "Export Failed", "Failed to export the diagram.")
    
    def _export_as_svg(self, file_path):
        from PyQt5.QtSvg import QSvgGenerator
        
        scene_rect = self.scene.sceneRect()
        
        generator = QSvgGenerator()
        generator.setFileName(file_path)
        generator.setSize(scene_rect.size().toSize())
        generator.setViewBox(scene_rect)
        generator.setTitle("Message Sequence Diagram")
        
        painter = QPainter(generator)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.setRenderHint(QPainter.TextAntialiasing)
        self.scene.render(painter)
        painter.end()
        
        QMessageBox.information(self, "Export Successful", f"Diagram exported to:\n{file_path}")