from PySide6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
                               QLabel, QComboBox, QScrollArea, QSplitter,
                               QMessageBox)
from PySide6.QtCore import Qt
from PySide6.QtGui import QImage, QPixmap
import fitz
from database import get_session
from models import Document, DocumentVersion
import difflib
import os
from PIL import Image
import numpy as np
import cv2

class DocumentCompare(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.doc1 = None
        self.doc2 = None
        self.current_page = 0
        self.setup_ui()
    
    def setup_ui(self):
        layout = QVBoxLayout(self)
        
        # Document selection
        select_layout = QHBoxLayout()
        
        # First document
        doc1_layout = QVBoxLayout()
        doc1_layout.addWidget(QLabel("Documento 1:"))
        self.doc1_combo = QComboBox()
        self.doc1_combo.currentIndexChanged.connect(self.load_comparison)
        doc1_layout.addWidget(self.doc1_combo)
        select_layout.addLayout(doc1_layout)
        
        # Second document
        doc2_layout = QVBoxLayout()
        doc2_layout.addWidget(QLabel("Documento 2:"))
        self.doc2_combo = QComboBox()
        self.doc2_combo.currentIndexChanged.connect(self.load_comparison)
        doc2_layout.addWidget(self.doc2_combo)
        select_layout.addLayout(doc2_layout)
        
        layout.addLayout(select_layout)
        
        # Navigation
        nav_layout = QHBoxLayout()
        
        self.prev_btn = QPushButton("Anterior")
        self.prev_btn.clicked.connect(self.prev_page)
        nav_layout.addWidget(self.prev_btn)
        
        self.page_label = QLabel("Página 0 de 0")
        nav_layout.addWidget(self.page_label)
        
        self.next_btn = QPushButton("Próxima")
        self.next_btn.clicked.connect(self.next_page)
        nav_layout.addWidget(self.next_btn)
        
        nav_layout.addStretch()
        
        self.diff_mode_combo = QComboBox()
        self.diff_mode_combo.addItems(["Visual", "Texto", "Híbrido"])
        self.diff_mode_combo.currentTextChanged.connect(self.load_comparison)
        nav_layout.addWidget(self.diff_mode_combo)
        
        layout.addLayout(nav_layout)
        
        # Comparison view
        self.splitter = QSplitter(Qt.Horizontal)
        
        # Left document
        self.left_scroll = QScrollArea()
        self.left_scroll.setWidgetResizable(True)
        self.left_label = QLabel()
        self.left_label.setAlignment(Qt.AlignCenter)
        self.left_scroll.setWidget(self.left_label)
        self.splitter.addWidget(self.left_scroll)
        
        # Right document
        self.right_scroll = QScrollArea()
        self.right_scroll.setWidgetResizable(True)
        self.right_label = QLabel()
        self.right_label.setAlignment(Qt.AlignCenter)
        self.right_scroll.setWidget(self.right_label)
        self.splitter.addWidget(self.right_scroll)
        
        # Difference view
        self.diff_scroll = QScrollArea()
        self.diff_scroll.setWidgetResizable(True)
        self.diff_label = QLabel()
        self.diff_label.setAlignment(Qt.AlignCenter)
        self.diff_scroll.setWidget(self.diff_label)
        self.splitter.addWidget(self.diff_scroll)
        
        layout.addWidget(self.splitter)
        
        # Difference summary
        self.diff_summary = QLabel()
        layout.addWidget(self.diff_summary)
        
        self.refresh_documents()
    
    def refresh_documents(self):
        try:
            session = get_session()
            documents = session.query(Document).all()
            
            # Update combo boxes
            self.doc1_combo.clear()
            self.doc2_combo.clear()
            
            for doc in documents:
                item_text = f"{doc.file_name} ({doc.marca} {doc.modelo})"
                self.doc1_combo.addItem(item_text, doc.id)
                self.doc2_combo.addItem(item_text, doc.id)
            
            session.close()
        except Exception as e:
            QMessageBox.critical(self, "Erro",
                               f"Erro ao carregar documentos: {str(e)}")
    
    def load_comparison(self):
        if (self.doc1_combo.currentData() is None or 
            self.doc2_combo.currentData() is None):
            return
        
        try:
            session = get_session()
            
            # Load documents
            doc1 = session.query(Document).get(self.doc1_combo.currentData())
            doc2 = session.query(Document).get(self.doc2_combo.currentData())
            
            if doc1 and doc2:
                self.doc1 = fitz.open(doc1.file_path)
                self.doc2 = fitz.open(doc2.file_path)
                self.current_page = 0
                self.update_comparison()
            
            session.close()
        except Exception as e:
            QMessageBox.critical(self, "Erro",
                               f"Erro ao carregar documentos: {str(e)}")
    
    def update_comparison(self):
        if not self.doc1 or not self.doc2:
            return
        
        try:
            # Update page navigation
            max_pages = min(len(self.doc1), len(self.doc2))
            self.page_label.setText(f"Página {self.current_page + 1} de {max_pages}")
            
            # Get pages
            page1 = self.doc1[self.current_page]
            page2 = self.doc2[self.current_page]
            
            # Process based on comparison mode
            mode = self.diff_mode_combo.currentText()
            
            if mode == "Visual":
                self.compare_visual(page1, page2)
            elif mode == "Texto":
                self.compare_text(page1, page2)
            else:  # Hybrid
                self.compare_hybrid(page1, page2)
            
        except Exception as e:
            QMessageBox.critical(self, "Erro",
                               f"Erro ao comparar páginas: {str(e)}")
    
    def compare_visual(self, page1, page2):
        # Convert pages to images
        pix1 = page1.get_pixmap()
        pix2 = page2.get_pixmap()
        
        # Convert to numpy arrays
        img1 = np.frombuffer(pix1.samples, dtype=np.uint8).reshape(
            pix1.height, pix1.width, 3)
        img2 = np.frombuffer(pix2.samples, dtype=np.uint8).reshape(
            pix2.height, pix2.width, 3)
        
        # Resize images to same size
        height = max(img1.shape[0], img2.shape[0])
        width = max(img1.shape[1], img2.shape[1])
        
        img1 = cv2.resize(img1, (width, height))
        img2 = cv2.resize(img2, (width, height))
        
        # Calculate difference
        diff = cv2.absdiff(img1, img2)
        diff_highlighted = img2.copy()
        diff_highlighted[diff.mean(axis=2) > 30] = [0, 0, 255]  # Red for differences
        
        # Convert to QPixmap and display
        self.display_image(img1, self.left_label)
        self.display_image(img2, self.right_label)
        self.display_image(diff_highlighted, self.diff_label)
        
        # Update summary
        diff_percentage = (diff.mean(axis=2) > 30).mean() * 100
        self.diff_summary.setText(
            f"Diferença visual: {diff_percentage:.2f}% da página")
    
    def compare_text(self, page1, page2):
        # Extract text
        text1 = page1.get_text()
        text2 = page2.get_text()
        
        # Compare text
        diff = difflib.unified_diff(
            text1.splitlines(),
            text2.splitlines(),
            lineterm=''
        )
        
        # Format differences
        diff_html = []
        for line in diff:
            if line.startswith('+'):
                diff_html.append(f'<span style="background-color: #aaffaa">{line}</span>')
            elif line.startswith('-'):
                diff_html.append(f'<span style="background-color: #ffaaaa">{line}</span>')
            else:
                diff_html.append(line)
        
        # Display original texts and differences
        self.left_label.setText(f"<pre>{text1}</pre>")
        self.right_label.setText(f"<pre>{text2}</pre>")
        self.diff_label.setText(f"<pre>{'<br>'.join(diff_html)}</pre>")
        
        # Update summary
        changes = len([l for l in diff_html if l.startswith('<span')])
        self.diff_summary.setText(
            f"Diferenças encontradas: {changes} linhas modificadas")
    
    def compare_hybrid(self, page1, page2):
        # Combine visual and text comparison
        self.compare_visual(page1, page2)
        
        # Add text comparison summary
        text1 = page1.get_text()
        text2 = page2.get_text()
        
        # Calculate text similarity
        similarity = difflib.SequenceMatcher(None, text1, text2).ratio()
        
        # Update summary
        current_summary = self.diff_summary.text()
        self.diff_summary.setText(
            f"{current_summary}\nSimilaridade do texto: {similarity:.2f}%")
    
    def display_image(self, img, label):
        height, width = img.shape[:2]
        bytes_per_line = 3 * width
        
        q_img = QImage(img.data, width, height, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(q_img)
        
        # Scale to fit while maintaining aspect ratio
        scaled_pixmap = pixmap.scaled(
            label.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        label.setPixmap(scaled_pixmap)
    
    def prev_page(self):
        if self.current_page > 0:
            self.current_page -= 1
            self.update_comparison()
    
    def next_page(self):
        if (self.doc1 and self.doc2 and 
            self.current_page < min(len(self.doc1), len(self.doc2)) - 1):
            self.current_page += 1
            self.update_comparison()
    
    def closeEvent(self, event):
        if self.doc1:
            self.doc1.close()
        if self.doc2:
            self.doc2.close()
        super().closeEvent(event)