"""Modern, stunning mask annotator with alpha blending."""

import cv2
import numpy as np
from typing import Optional
from eye.detection.core import Detections
from eye.draw.color import ColorPalette


class MaskAnnotator:
    """Draw segmentation masks with transparency and effects.
    
    Stunning features: Gradient fills, edge glow, pattern fills.
    """
    
    def __init__(
        self,
        color_palette: Optional[ColorPalette] = None,
        opacity: float = 0.5,
        border_thickness: int = 2,
        edge_glow: bool = False,
        glow_intensity: float = 0.3
    ):
        """
        Args:
            color_palette: Color palette for masks
            opacity: Mask transparency (0-1)
            border_thickness: Border thickness (0 = no border)
            edge_glow: Add soft glow around edges
            glow_intensity: Glow intensity (0-1)
        """
        from eye.draw.color import PredefinedPalettes
        self.colors = color_palette or PredefinedPalettes.bright()
        self.opacity = opacity
        self.border_thickness = border_thickness
        self.edge_glow = edge_glow
        self.glow_intensity = glow_intensity
    
    def annotate(
        self,
        image: np.ndarray,
        detections: Detections
    ) -> np.ndarray:
        """Draw masks on image."""
        if detections.mask is None:
            return image
        
        annotated = image.copy()
        masks = detections.mask
        
        if len(masks) == 0:
            return annotated
        
        # Create overlay
        overlay = image.copy()
        
        for i in range(len(detections)):
            if i >= len(masks):
                break
            
            mask = masks[i]
            if mask is None:
                continue
            
            # Get color
            if detections.class_id is not None:
                color = self.colors.by_id(detections.class_id[i])
            elif detections.tracker_id is not None:
                color = self.colors.by_id(detections.tracker_id[i])
            else:
                color = self.colors[i]
            
            bgr_color = color.as_bgr()
            
            # Ensure mask is binary
            if mask.dtype != bool:
                mask = mask > 0.5
            
            # Fill mask
            overlay[mask] = bgr_color
            
            # Edge glow effect
            if self.edge_glow:
                kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
                dilated = cv2.dilate(mask.astype(np.uint8), kernel, iterations=2)
                edge = dilated - mask.astype(np.uint8)
                edge_mask = edge > 0
                
                # Soft glow
                glow_color = tuple(min(255, int(c * (1 + self.glow_intensity))) for c in bgr_color)
                overlay[edge_mask] = glow_color
            
            # Border
            if self.border_thickness > 0:
                contours, _ = cv2.findContours(
                    mask.astype(np.uint8),
                    cv2.RETR_EXTERNAL,
                    cv2.CHAIN_APPROX_SIMPLE
                )
                cv2.drawContours(overlay, contours, -1, bgr_color, self.border_thickness)
        
        # Blend
        annotated = cv2.addWeighted(overlay, self.opacity, annotated, 1 - self.opacity, 0)
        
        return annotated


class GradientBoxAnnotator:
    """Draw boxes with gradient fills - Modern look!"""
    
    def __init__(
        self,
        color_palette: Optional[ColorPalette] = None,
        thickness: int = 3,
        corner_radius: int = 12,
        gradient: bool = True
    ):
        """
        Args:
            color_palette: Color palette
            thickness: Border thickness
            corner_radius: Corner radius for rounded boxes
            gradient: Fill with gradient
        """
        from eye.draw.color import PredefinedPalettes
        self.colors = color_palette or PredefinedPalettes.bright()
        self.thickness = thickness
        self.corner_radius = corner_radius
        self.gradient = gradient
    
    def annotate(
        self,
        image: np.ndarray,
        detections: Detections
    ) -> np.ndarray:
        """Draw gradient boxes."""
        annotated = image.copy()
        
        for i in range(len(detections)):
            x1, y1, x2, y2 = detections.xyxy[i].astype(int)
            
            # Get color
            if detections.class_id is not None:
                color = self.colors.by_id(detections.class_id[i])
            elif detections.tracker_id is not None:
                color = self.colors.by_id(detections.tracker_id[i])
            else:
                color = self.colors[i]
            
            bgr = color.as_bgr()
            
            # Gradient fill
            if self.gradient:
                overlay = annotated.copy()
                for y in range(y1, y2):
                    alpha = (y - y1) / max(1, (y2 - y1))
                    row_color = tuple(int(c * (0.3 + 0.7 * alpha)) for c in bgr)
                    cv2.rectangle(overlay, (x1, y), (x2, y+1), row_color, -1)
                cv2.addWeighted(overlay, 0.3, annotated, 0.7, 0, annotated)
            
            # Rounded rectangle
            self._draw_rounded_rect(annotated, (x1, y1), (x2, y2), bgr, self.corner_radius, self.thickness)
        
        return annotated
    
    @staticmethod
    def _draw_rounded_rect(img, pt1, pt2, color, radius, thickness):
        """Draw rounded rectangle."""
        x1, y1 = pt1
        x2, y2 = pt2
        
        # Draw rectangles
        cv2.rectangle(img, (x1 + radius, y1), (x2 - radius, y2), color, thickness)
        cv2.rectangle(img, (x1, y1 + radius), (x2, y2 - radius), color, thickness)
        
        # Draw circles at corners
        cv2.circle(img, (x1 + radius, y1 + radius), radius, color, thickness)
        cv2.circle(img, (x2 - radius, y1 + radius), radius, color, thickness)
        cv2.circle(img, (x1 + radius, y2 - radius), radius, color, thickness)
        cv2.circle(img, (x2 - radius, y2 - radius), radius, color, thickness)


class NeonTraceAnnotator:
    """Neon-style trails - Ultra modern!"""
    
    def __init__(
        self,
        color_palette: Optional[ColorPalette] = None,
        thickness: int = 4,
        trace_length: int = 50,
        glow: bool = True
    ):
        """
        Args:
            color_palette: Color palette
            thickness: Line thickness
            trace_length: Max trail points
            glow: Add neon glow effect
        """
        from eye.draw.color import PredefinedPalettes
        from collections import deque
        self.colors = color_palette or PredefinedPalettes.bright()
        self.thickness = thickness
        self.trace_length = trace_length
        self.glow = glow
        self.history = {}
        # Position smoothing alpha when appending to trace history (0-1).
        # Lower = more smoothing. Default 0.6 is moderate smoothing.
        self.position_smoothing_alpha = 0.6
        # When objects move faster than this (pixels/frame), use `fast_position_smoothing_alpha`
        # to apply stronger smoothing for jitter reduction on fast objects.
        self.position_speed_threshold = 8.0
        # Alpha to use when speed > threshold (lower -> more smoothing)
        self.fast_position_smoothing_alpha = 0.3
    
    def annotate(
        self,
        image: np.ndarray,
        detections: Detections
    ) -> np.ndarray:
        """Draw neon trails."""
        if detections.tracker_id is None or len(detections) == 0:
            return image
        
        annotated = image.copy()
        
        # Update history
        from collections import deque
        from eye.geometry.core import Position
        
        # Compute centers using get_anchors_coordinates
        centers = detections.get_anchors_coordinates(anchor=Position.CENTER)
        
        for i in range(len(detections)):
            tid = detections.tracker_id[i]
            center = centers[i]

            if tid not in self.history:
                self.history[tid] = deque(maxlen=self.trace_length)

            # Apply lightweight temporal smoothing to the appended center
            if len(self.history[tid]) > 0 and self.position_smoothing_alpha is not None:
                last = np.array(self.history[tid][-1], dtype=float)
                newc = np.array(center, dtype=float)
                dist = float(np.linalg.norm(newc - last))
                # choose alpha based on speed
                if self.position_speed_threshold is not None and dist > float(self.position_speed_threshold):
                    alpha = float(self.fast_position_smoothing_alpha)
                else:
                    alpha = float(self.position_smoothing_alpha)
                center = alpha * newc + (1.0 - alpha) * last

            self.history[tid].append(center)
        
        # Draw trails
        for i in range(len(detections)):
            tid = detections.tracker_id[i]
            if tid not in self.history or len(self.history[tid]) < 2:
                continue
            
            points = list(self.history[tid])
            color = self.colors.by_id(detections.class_id[i] if detections.class_id is not None else tid)
            
            # Draw with fade and glow
            for j in range(len(points) - 1):
                alpha = (j + 1) / len(points)
                pt1 = tuple(points[j].astype(int))
                pt2 = tuple(points[j + 1].astype(int))
                
                # Glow effect (thicker, semi-transparent)
                if self.glow:
                    glow_color = tuple(int(c * 0.5) for c in color.as_bgr())
                    cv2.line(annotated, pt1, pt2, glow_color, self.thickness + 4, cv2.LINE_AA)
                
                # Main line
                main_color = tuple(int(c * alpha) for c in color.as_bgr())
                cv2.line(annotated, pt1, pt2, main_color, self.thickness, cv2.LINE_AA)
        
        return annotated
    
    def reset(self):
        """Clear history."""
        self.history.clear()


class ShadowBoxAnnotator:
    """Boxes with drop shadows - Depth effect!"""
    
    def __init__(
        self,
        color_palette: Optional[ColorPalette] = None,
        thickness: int = 2,
        shadow_offset: int = 5,
        corner_radius: int = 8
    ):
        """
        Args:
            color_palette: Color palette
            thickness: Border thickness
            shadow_offset: Shadow offset in pixels
            corner_radius: Corner radius
        """
        from eye.draw.color import PredefinedPalettes
        self.colors = color_palette or PredefinedPalettes.bright()
        self.thickness = thickness
        self.shadow_offset = shadow_offset
        self.corner_radius = corner_radius
    
    def annotate(
        self,
        image: np.ndarray,
        detections: Detections
    ) -> np.ndarray:
        """Draw boxes with shadows."""
        annotated = image.copy()
        
        for i in range(len(detections)):
            x1, y1, x2, y2 = detections.xyxy[i].astype(int)
            
            color = self.colors.by_id(
                detections.class_id[i] if detections.class_id is not None 
                else detections.tracker_id[i] if detections.tracker_id is not None 
                else i
            )
            
            # Draw shadow
            shadow_color = (30, 30, 30)  # Dark gray
            cv2.rectangle(
                annotated,
                (x1 + self.shadow_offset, y1 + self.shadow_offset),
                (x2 + self.shadow_offset, y2 + self.shadow_offset),
                shadow_color,
                self.thickness
            )
            
            # Draw main box
            GradientBoxAnnotator._draw_rounded_rect(
                annotated, (x1, y1), (x2, y2),
                color.as_bgr(), self.corner_radius, self.thickness
            )
        
        return annotated


class CornerBoxAnnotator:
    """Draw boxes using only corner segments (no full borders)."""

    def __init__(
        self,
        color_palette: Optional[ColorPalette] = None,
        thickness: int = 3,
        corner_length: int = 20,
        corner_radius: int = 6,
    ):
        from eye.draw.color import PredefinedPalettes
        self.colors = color_palette or PredefinedPalettes.bright()
        self.thickness = thickness
        self.corner_length = corner_length
        self.corner_radius = corner_radius

    def annotate(self, image: np.ndarray, detections: Detections) -> np.ndarray:
        annotated = image.copy()
        for i in range(len(detections)):
            x1, y1, x2, y2 = detections.xyxy[i].astype(int)
            color = self.colors.by_id(
                detections.class_id[i] if detections.class_id is not None else (
                    detections.tracker_id[i] if detections.tracker_id is not None else i
                )
            ).as_bgr()

            # Top-left
            cv2.line(annotated, (x1, y1), (x1 + self.corner_length, y1), color, self.thickness)
            cv2.line(annotated, (x1, y1), (x1, y1 + self.corner_length), color, self.thickness)
            # Top-right
            cv2.line(annotated, (x2, y1), (x2 - self.corner_length, y1), color, self.thickness)
            cv2.line(annotated, (x2, y1), (x2, y1 + self.corner_length), color, self.thickness)
            # Bottom-left
            cv2.line(annotated, (x1, y2), (x1 + self.corner_length, y2), color, self.thickness)
            cv2.line(annotated, (x1, y2), (x1, y2 - self.corner_length), color, self.thickness)
            # Bottom-right
            cv2.line(annotated, (x2, y2), (x2 - self.corner_length, y2), color, self.thickness)
            cv2.line(annotated, (x2, y2), (x2, y2 - self.corner_length), color, self.thickness)

        return annotated


class FPSAnnotator:
    """Annotator that shows processing FPS on the frame.

    Usage: call `update(fps)` from the processing loop to set the current value.
    """

    def __init__(self, position: str = "top-left", font_scale: float = 0.6, color=(0, 255, 0), bg_color=(0, 0, 0), padding: int = 6):
        self.position = position
        self.font_scale = font_scale
        self.color = color
        self.bg_color = bg_color
        self.padding = padding
        self.fps = None

    def update(self, fps: float):
        try:
            self.fps = float(fps)
        except Exception:
            self.fps = None

    def annotate(self, image: np.ndarray, detections: Optional[Detections] = None) -> np.ndarray:
        annotated = image.copy()
        if self.fps is None:
            return annotated

        text = f"FPS: {self.fps:.1f}"
        font = cv2.FONT_HERSHEY_SIMPLEX
        thickness = 1
        (w, h), _ = cv2.getTextSize(text, font, self.font_scale, thickness)

        if self.position == "top-left":
            x, y = 10, 10 + h
        elif self.position == "top-right":
            x, y = annotated.shape[1] - w - 10, 10 + h
        elif self.position == "bottom-left":
            x, y = 10, annotated.shape[0] - 10
        else:
            x, y = annotated.shape[1] - w - 10, annotated.shape[0] - 10

        # Background rectangle
        cv2.rectangle(annotated, (x - self.padding, y - h - self.padding), (x + w + self.padding, y + self.padding), self.bg_color, -1)
        cv2.putText(annotated, text, (x, y), font, self.font_scale, self.color, thickness)
        return annotated


class InfoAnnotator:
    """Flexible information/text/table annotator.

    - Reads info from `detections.data.get('info')` if present, otherwise from `self.info`.
    - `info` can be a dict (key->value), a list of (k,v) pairs, or a list of rows (for tables).
    """

    def __init__(
        self,
        position: str = "top-left",
        as_table: bool = False,
        cols: int = 1,
        show_border: bool = False,
        bg_color=(0, 0, 0),
        text_color=(255, 255, 255),
        font_scale: float = 0.5,
        padding: int = 6,
    ):
        self.position = position
        self.as_table = as_table
        self.cols = max(1, int(cols))
        self.show_border = show_border
        self.bg_color = bg_color
        self.text_color = text_color
        self.font_scale = font_scale
        self.padding = padding
        self.info = None

    def annotate(self, image: np.ndarray, detections: Optional[Detections] = None) -> np.ndarray:
        annotated = image.copy()
        info = None
        if detections is not None and isinstance(detections, Detections):
            info = detections.data.get('info') if isinstance(detections.data, dict) else None
        if info is None:
            info = self.info
        if not info:
            return annotated

        # Normalize info into list of strings per row
        rows = []
        if isinstance(info, dict):
            for k, v in info.items():
                rows.append(f"{k}: {v}")
        elif isinstance(info, list):
            # Could be list of tuples or list of rows
            for item in info:
                if isinstance(item, (list, tuple)):
                    rows.append(" | ".join(str(x) for x in item))
                else:
                    rows.append(str(item))
        else:
            rows = [str(info)]

        # When as_table and cols>1, arrange rows into grid
        if self.as_table and self.cols > 1:
            grid = []
            for i in range(0, len(rows), self.cols):
                grid.append(rows[i:i + self.cols])
            # compute cell sizes
            cell_texts = ["\t".join(r) for r in grid]
            rows = cell_texts

        # compute text block size
        font = cv2.FONT_HERSHEY_SIMPLEX
        thickness = 1
        line_heights = []
        max_w = 0
        for r in rows:
            (w, h), _ = cv2.getTextSize(r, font, self.font_scale, thickness)
            line_heights.append(h)
            max_w = max(max_w, w)

        total_h = sum(line_heights) + self.padding * 2 + (len(rows) - 1) * 4
        total_w = max_w + self.padding * 2

        # choose position
        if self.position == 'top-left':
            x0, y0 = 10, 10
        elif self.position == 'top-right':
            x0, y0 = annotated.shape[1] - total_w - 10, 10
        elif self.position == 'bottom-left':
            x0, y0 = 10, annotated.shape[0] - total_h - 10
        else:
            x0, y0 = annotated.shape[1] - total_w - 10, annotated.shape[0] - total_h - 10

        # background
        cv2.rectangle(annotated, (x0, y0), (x0 + total_w, y0 + total_h), self.bg_color, -1)
        if self.show_border:
            cv2.rectangle(annotated, (x0, y0), (x0 + total_w, y0 + total_h), (200, 200, 200), 1)

        # draw lines
        y = y0 + self.padding + line_heights[0]
        for idx, r in enumerate(rows):
            cv2.putText(annotated, r, (x0 + self.padding, y), font, self.font_scale, self.text_color, thickness)
            if idx + 1 < len(rows):
                y += line_heights[idx + 1] + 4

        return annotated
