
from PyQt6.QtWidgets import QWidget
from PyQt6.QtGui import QPainter, QColor, QFont, QImage, QBrush, QPen
from PyQt6.QtCore import Qt, QRect, QSize, pyqtSignal

from .models import Page, Packet
from .charsets import get_char, get_byte_from_char

# Helper to create a blank packet
def create_blank_packet(magazine: int, row: int) -> Packet:
    # 42 bytes
    # Bytes 0-1: Address (Hamming encoded)
    # We need to compute Hamming for M and R.
    # For now, let's just make a placeholder that "looks" okay or use 0x00 and let a future saver fix it.
    # But wait, `Packet` expects 42 bytes input.
    
    # Ideally we should implement the hamming encoder from `io.py` or similar.
    # For now, we will create a Packet with dummy address bytes and correct mag/row properties set manually if needed,
    # OR we just construct the bytes.
    
    # Let's rely on the Packet class parsing... which is annoying if we can't easily encode.
    # Hack: default to 0x00 0x00, then manually set .magazine and .row
    
    data = bytearray(42)
    # Set spaces
    data[2:] = b'\x20' * 40
    
    p = Packet(bytes(data))
    # Override
    p.magazine = magazine
    p.row = row
    p.data = bytearray(b'\x20' * 40)
    
    # We really should set the address bytes correctly so saving works "naturally".
    # But `save_t42` uses `packet.original_data[:2]`.
    # So we MUST encode properly or `save_t42` needs to regenerate headers.
    
    # Let's defer proper Hamming encoding to the Save step (as noted in TODO in io.py).
    # For in-memory editing, this is fine.
    return p


# Teletext Palette
COLORS = [
    QColor(0, 0, 0),       # Black
    QColor(255, 0, 0),     # Red
    QColor(0, 255, 0),     # Green
    QColor(255, 255, 0),   # Yellow
    QColor(0, 0, 255),     # Blue
    QColor(255, 0, 255),   # Magenta
    QColor(0, 255, 255),   # Cyan
    QColor(255, 255, 255)  # White
]

class TeletextCanvas(QWidget):
    cursorChanged = pyqtSignal(int, int, int, bool) # x, y, byte_val, is_graphics

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMouseTracking(True) # Just in case
        self.setMinimumSize(800, 600) # 40x20 * 25x24
        self.page: Page = None
        self.subset_idx = 0 # Default English
        
        # Teletext is 40 columns x 25 rows
        # We will render to a fixed size QImage and scale it
        self.cols = 40
        self.rows = 25
        self.cell_w = 20
        self.cell_h = 24
        self.img_w = self.cols * self.cell_w
        self.img_h = self.rows * self.cell_h
        
        self.buffer = QImage(self.img_w, self.img_h, QImage.Format.Format_RGB32)
        self.buffer.fill(Qt.GlobalColor.black)
        
        # Font for text
        self.font = QFont("Courier New", 18)
        self.font.setStyleHint(QFont.StyleHint.Monospace)
        self.font.setBold(True)

        # Cursor state
        self.cursor_x = 0
        self.cursor_y = 0
        self.cursor_visible = True
        self.cursor_is_graphics = False # Tracked during draw
        # Blinking cursor timer could be added, for now static inverted is fine or toggle on timer elsewhere

        # Highlight state (for search results)
        self.highlight_row = -1
        self.highlight_start_x = -1
        self.highlight_end_x = -1

    def set_highlight(self, row, start_x, end_x):
        self.highlight_row = row
        self.highlight_start_x = start_x
        self.highlight_end_x = end_x
        self.redraw()
        self.update()

    def clear_highlight(self):
        self.highlight_row = -1
        self.highlight_start_x = -1
        self.highlight_end_x = -1
        self.redraw()
        self.update()

    def get_byte_at(self, x, y):
        if not self.page: return 0
        
        # Row 0 special logic (header) - returning what is displayed or raw?
        # User probably wants RAW byte for editing.
        
        packet = None
        for p in self.page.packets:
            if p.row == y:
                packet = p
                break
        
        if packet and 0 <= x < 40:
             return packet.data[x]
        return 0

    def emit_cursor_change(self):
        val = self.get_byte_at(self.cursor_x, self.cursor_y)
        self.cursorChanged.emit(self.cursor_x, self.cursor_y, val, self.cursor_is_graphics)

    def set_cursor(self, x, y):
        self.cursor_x = max(0, min(self.cols - 1, x))
        self.cursor_y = max(0, min(self.rows - 1, y))
        self.redraw()
        self.update()
        self.emit_cursor_change()

    def move_cursor(self, dx, dy):
        self.cursor_x = max(0, min(self.cols - 1, self.cursor_x + dx))
        self.cursor_y = max(0, min(self.rows - 1, self.cursor_y + dy))
        self.redraw()
        self.update()
        self.emit_cursor_change()

    def set_byte_at_cursor(self, byte_val):
        if not self.page: return
        
        packet = None
        for p in self.page.packets:
            if p.row == self.cursor_y:
                packet = p
                break
        
        if packet:
            if 0 <= self.cursor_x < 40:
                packet.data[self.cursor_x] = byte_val
                self.redraw()
                self.emit_cursor_change()
        else:
            # Create packet if missing
            self.handle_input(chr(byte_val)) # Lazy reuse or duplicate create logic
    
    def set_page(self, page: Page):
        self.page = page
        # Set language from page header
        if page:
            self.subset_idx = page.language
        else:
            self.subset_idx = 0
            
        self.cursor_x = 0
        self.cursor_y = 0
        self.redraw()
        self.update()
        self.emit_cursor_change()

    def handle_input(self, text):
        if not self.page:
            return
            
        # Find packet for current row
        packet = None
        for p in self.page.packets:
            if p.row == self.cursor_y:
                packet = p
                break
        
        if packet:
            # Edit the data
            # Data bytes start at index 0 of packet.data corresponding to Col 0?
            # Standard Teletext row is 40 chars. Packet data is 40 bytes.
            if 0 <= self.cursor_x < 40:
                # Get ASCII value
                # We need to Reverse Map chars to Teletext Bytes if we want full suppport
                # For now, just taking ord() for basic ASCII 
                # TODO: reverse lookup for National Chars
                
                # Check if text is a single char
                if len(text) == 1:
                    byte_val = get_byte_from_char(text, self.subset_idx)
                    
                    # Simple filter
                    if byte_val > 255: byte_val = 0x3F # ?
                    
                    packet.data[self.cursor_x] = byte_val
                    self.redraw()
                    self.move_cursor(1, 0)
        else:
            # Create a packet if it doesn't exist for this row
            if 0 <= self.cursor_x < 40:
                 # Check if we should create it
                 # Only if input is valid char?
                if len(text) == 1:
                     # Create new packet
                     new_packet = create_blank_packet(self.page.magazine, self.cursor_y)
                     self.page.packets.append(new_packet)
                     # Sort packets by row? Or just append. Renderer handles unordered.
                     # But for sanity, let's just append.
                     
                     # Write the char
                     byte_val = get_byte_from_char(text, self.subset_idx)
                     if byte_val > 255: byte_val = 0x3F
                     new_packet.data[self.cursor_x] = byte_val
                     
                     self.redraw()
                     self.move_cursor(1, 0)
        
    def redraw(self):
        self.buffer.fill(Qt.GlobalColor.black)
        painter = QPainter(self.buffer)
        painter.setFont(self.font)
        
        if not self.page:
            painter.setPen(Qt.GlobalColor.white)
            painter.drawText(10, 20, "No Page Loaded")
            painter.end()
            return
            
        # Check Control Bits for "Inhibit Display" (C10)
        # In our bitmask (from parse_header):
        # C4:0, C5:1, C6:2, C7:3, C8:4, C9:5, C10:6, C11:7, C12:8, C13:9, C14:10
        inhibit_display = bool((self.page.control_bits >> 6) & 1)
        if inhibit_display:
            painter.setPen(Qt.GlobalColor.gray)
            painter.drawText(10, 20, f"Page {self.page.full_page_number} - INHIBIT DISPLAY (C10 set)")
            painter.end()
            return

        # Organize each packet by row
        grid = [None] * 26 # 0-25
        for p in self.page.packets:
            if 0 <= p.row <= 25:
                grid[p.row] = p
                
        # Pass 1: Backgrounds
        occlusion_mask = [False] * 40
        for r in range(25):
            packet = grid[r]
            occlusion_mask = self.draw_row(painter, r, packet, draw_bg=True, draw_fg=False, occlusion_mask=occlusion_mask)

        # Pass 2: Foregrounds
        occlusion_mask = [False] * 40
        for r in range(25):
            packet = grid[r]
            occlusion_mask = self.draw_row(painter, r, packet, draw_bg=False, draw_fg=True, occlusion_mask=occlusion_mask)
            
        painter.end()

    def draw_row(self, painter, row, packet, draw_bg=True, draw_fg=True, occlusion_mask=None):
        if occlusion_mask is None:
            occlusion_mask = [False] * 40
            
        # Output mask for the next row
        next_occlusion_mask = [False] * 40

        # Check for Suppress Header (C7)
        # C7:3, so bit 3 of control_bits
        suppress_header = bool((self.page.control_bits >> 3) & 1)
        
        # Default State at start of row
        fg = COLORS[7] # White
        bg = COLORS[0] # Black
        graphics_mode = False
        contiguous = True # Mosaic
        hold_graphics = False
        held_char = 0x20 # Space
        double_height = False
        
        last_visible_idx = -1
        bg_segments = [(0, bg)] # Track BG changes: (index, color)
        
        y = row * self.cell_h
        
        data = b''
        if packet:
            data = packet.data # 40 bytes
        else:
            data = b'\x20' * 40
            
        # Header string for Row 0 columns 0-7
        header_prefix = ""
        if row == 0 and self.page:
            header_prefix = f"P{self.page.magazine}{self.page.page_number:02X}"
            # Pad to 8 chars
            header_prefix = header_prefix.ljust(8)
            
        for c in range(40):
            x = c * self.cell_w
            is_occluded = occlusion_mask[c]
            
            # Decide byte value
            if row == 0:
                if c < 8:
                    # Column 0-7: Header prefix
                    byte_val = ord(header_prefix[c])
                elif suppress_header and c < 32:
                    # Column 8-31: Hide header if C7 set
                    byte_val = 0x20
                else:
                    byte_val = data[c] if c < len(data) else 0x20
            else:
                byte_val = data[c] if c < len(data) else 0x20
                
            byte_val &= 0x7F # Strip parity
            
            # Control Codes
            is_control = False
            
            if byte_val < 0x20:
                is_control = True
                # Handle control codes
                if 0x00 <= byte_val <= 0x07: # Alpha Color
                    fg = COLORS[byte_val]
                    graphics_mode = False
                elif 0x10 <= byte_val <= 0x17: # Mosaic Color
                    fg = COLORS[byte_val - 0x10]
                    graphics_mode = True
                elif byte_val == 0x1C: # Black BG
                    bg = COLORS[0]
                    bg_segments.append((c, bg))
                elif byte_val == 0x1D: # New BG
                    bg = fg
                    bg_segments.append((c, bg))
                    
                elif byte_val == 0x0C: # Normal Height
                    double_height = False
                elif byte_val == 0x0D: # Double Height
                    double_height = True
                    # Backfill Height if we are in leading controls
                    if last_visible_idx == -1:
                        # Update occlusion mask for 0..c-1
                        for k in range(c):
                            next_occlusion_mask[k] = True
                            
                        # Repaint 0..c-1 with DH
                        if draw_bg:
                            # Repaint each cell with its historical BG
                            for k in range(c):
                                # Resolve BG for k
                                cell_bg = bg_segments[0][1]
                                for idx, color in reversed(bg_segments):
                                    if idx <= k:
                                        cell_bg = color
                                        break
                                painter.fillRect(k * self.cell_w, y, self.cell_w, self.cell_h * 2, cell_bg)

                elif byte_val == 0x19: # Contiguous Graphics
                    contiguous = True
                elif byte_val == 0x1A: # Separated Graphics
                    contiguous = False
                elif byte_val == 0x1E: # Hold Graphics
                    hold_graphics = True
                elif byte_val == 0x1F: # Release Graphics
                    hold_graphics = False
            
            # Update visibility tracker
            # If it's a control code, it's "invisible" (space) UNLESS Held Graphics draws something
            visible_content = False
            if is_control:
                if hold_graphics and graphics_mode:
                    visible_content = True
            else:
                 # Treat Space (0x20) as "invisible" for backfill purposes
                 # This allows backfilling height over leading spaces in banners
                 if byte_val > 0x20:
                     visible_content = True
                 
            if visible_content:
                # If this is the first visible char, mark it
                if last_visible_idx == -1:
                    last_visible_idx = c

            # Record Double Height for next row
            if double_height and not is_occluded:
                next_occlusion_mask[c] = True

            # Capture cursor state if this is the cursor position
            if c == self.cursor_x and row == self.cursor_y:
                self.cursor_is_graphics = graphics_mode

            # If occluded, do not draw anything for this cell
            if is_occluded:
                continue

            # Draw Background
            if draw_bg:
                # If double height, draw taller background
                h_bg = self.cell_h * 2 if double_height else self.cell_h
                
                # Check for highlight
                is_highlighted = (row == self.highlight_row and self.highlight_start_x <= c <= self.highlight_end_x)
                
                if is_highlighted:
                    # Draw highlight color (e.g. Dark Yellow / Orange to contrast white text)
                    highlight_color = QColor(180, 150, 0) 
                    painter.fillRect(x, y, self.cell_w, h_bg, highlight_color)
                else:
                    painter.fillRect(x, y, self.cell_w, h_bg, bg)
            
            # Draw Foreground
            if draw_fg:
                # Calculate height
                # For Mosaics, we use the height param.
                # For Alphanumerics, we scale the painter.
                
                if is_control:
                    # "Set-at" spacing attribute? Teletext control codes occupy a space
                    # unless "Hold Graphics" replaces it with previous graphic char.
                    if hold_graphics and graphics_mode:
                        if double_height:
                             self.draw_mosaic(painter, x, y, held_char, fg, contiguous, height=self.cell_h * 2)
                        else:
                             self.draw_mosaic(painter, x, y, held_char, fg, contiguous)
                    else:
                        # Draw space (nothing, since we filled BG)
                        pass
                else:
                    if graphics_mode:
                        # Mosaic Graphics
                        h_mos = self.cell_h * 2 if double_height else self.cell_h
                        if (0x20 <= byte_val <= 0x3F) or (0x60 <= byte_val <= 0x7F):
                            self.draw_mosaic(painter, x, y, byte_val, fg, contiguous, height=h_mos)
                            held_char = byte_val
                        else:
                            # Capital letter in graphics mode? Usually shows char?
                            char = get_char(byte_val, self.subset_idx)
                            painter.setPen(fg)
                            if double_height:
                                painter.save()
                                painter.translate(x, y)
                                painter.scale(1, 2)
                                painter.drawText(QRect(0, 0, self.cell_w, self.cell_h), Qt.AlignmentFlag.AlignCenter, char)
                                painter.restore()
                            else:
                                painter.drawText(QRect(x, y, self.cell_w, self.cell_h), Qt.AlignmentFlag.AlignCenter, char)
                            held_char = 0x20
                    else:
                        # Alphanumeric
                        char = get_char(byte_val, self.subset_idx)
                        painter.setPen(fg)
                        if double_height:
                            painter.save()
                            painter.translate(x, y)
                            painter.scale(1, 2)
                            painter.drawText(QRect(0, 0, self.cell_w, self.cell_h), Qt.AlignmentFlag.AlignCenter, char)
                            painter.restore()
                        else:
                            painter.drawText(QRect(x, y, self.cell_w, self.cell_h), Qt.AlignmentFlag.AlignCenter, char)
            # Draw Cursor
            # Invert the cell at cursor position
            if draw_fg and self.cursor_visible and c == self.cursor_x and row == self.cursor_y:
                 painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Difference)
                 # Difference with white creates inversion
                 h_cursor = self.cell_h * 2 if double_height else self.cell_h
                 painter.fillRect(x, y, self.cell_w, h_cursor, QColor(255, 255, 255))
                 painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_SourceOver)
        
        return next_occlusion_mask
    def mousePressEvent(self, event):
        self.setFocus()
        # Calculate cell from mouse position
        # Need to account for scaling? 
        # The widget sizes to fit, but we render to a fixed buffer then scale up.
        # paintEvent scales buffer to self.rect().
        
        # This is a bit tricky because paintEvent uses aspect-ratio scaling and centering.
        # We need to reverse the mapping.
        
        target_rect = self.rect()
        # Re-calc scale and offset
        # Note: This logic duplicates paintEvent, commonize?
        
        # Simple approximation for MVP:
        # If the window is roughly aspect ratio correct, direct mapping works.
        # If not, we need offset.
        
        # Let's get the exact rect used in paintEvent
        scale_x = target_rect.width() / self.img_w
        scale_y = target_rect.height() / self.img_h
        scale = min(scale_x, scale_y)
        
        dw = self.img_w * scale
        dh = self.img_h * scale
        
        ox = (target_rect.width() - dw) / 2
        oy = (target_rect.height() - dh) / 2
        
        mx = event.pos().x() - ox
        my = event.pos().y() - oy
        
        if 0 <= mx < dw and 0 <= my < dh:
             col = int(mx / (self.cell_w * scale))
             row = int(my / (self.cell_h * scale))
             self.set_cursor(col, row)
             
    def draw_mosaic(self, painter, x, y, char_code, color, contiguous, height=None):
        if height is None:
            height = self.cell_h

        val = char_code & 0x7F
        bits = 0
        if val >= 0x20:
            bits = val - 0x20
            
        # Grid definitions for 2x3 grid
        x_splits = [0, int(self.cell_w / 2), self.cell_w]
        y_splits = [0, int(height / 3), int(2 * height / 3), height]
        
        # Block indices (col, row) for the 6 bits
        block_indices = [
            (0, 0), (1, 0), # Top
            (0, 1), (1, 1), # Mid
            (0, 2), (1, 2)  # Bot
        ]
        
        bit_mask = [1, 2, 4, 8, 16, 64] # 64 is bit 6
        
        painter.setPen(Qt.PenStyle.NoPen)
        painter.setBrush(QBrush(color))
        
        for i in range(6):
            if bits & bit_mask[i]:
                c, r = block_indices[i]
                
                bx_local = x_splits[c]
                by_local = y_splits[r]
                bw = x_splits[c+1] - x_splits[c]
                bh = y_splits[r+1] - y_splits[r]
                
                bx = x + bx_local
                by = y + by_local
                
                if not contiguous:
                    bx += 1
                    by += 1
                    bw -= 1
                    bh -= 1
                    
                painter.drawRect(QRect(int(bx), int(by), int(bw), int(bh)))

    def paintEvent(self, event):
        painter = QPainter(self)
        target_rect = self.rect()
        scaled = self.buffer.scaled(target_rect.size(), Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.FastTransformation)
        
        ox = (target_rect.width() - scaled.width()) // 2
        oy = (target_rect.height() - scaled.height()) // 2
        
        painter.drawImage(ox, oy, scaled)
