PROJECT-CONTORL/backend/app/services/watermark_service.py

"""
Watermark Service for MED-009: Dynamic Watermark for Downloads

This service provides functions to add watermarks to image and PDF files
containing user information for audit and tracking purposes.

Watermark content includes:
- User name
- Employee ID (or email as fallback)
- Download timestamp
"""

import io
import logging
import math
from datetime import datetime
from typing import Optional, Tuple

import fitz  # PyMuPDF
from PIL import Image, ImageDraw, ImageFont

logger = logging.getLogger(__name__)


class WatermarkService:
    """Service for adding watermarks to downloaded files."""

    # Watermark configuration
    WATERMARK_OPACITY = 0.3  # 30% opacity for semi-transparency
    WATERMARK_ANGLE = -45  # Diagonal angle in degrees
    WATERMARK_FONT_SIZE = 24
    WATERMARK_COLOR = (128, 128, 128)  # Gray color for watermark
    WATERMARK_SPACING = 200  # Spacing between repeated watermarks

    @staticmethod
    def _format_watermark_text(
        user_name: str,
        employee_id: Optional[str] = None,
        download_time: Optional[datetime] = None
    ) -> str:
        """
        Format the watermark text with user information.

        Args:
            user_name: Name of the user
            employee_id: Employee ID (工號) - uses 'N/A' if not provided
            download_time: Time of download (defaults to now)

        Returns:
            Formatted watermark text
        """
        if download_time is None:
            download_time = datetime.now()
        time_str = download_time.strftime("%Y-%m-%d %H:%M:%S")
        emp_id = employee_id if employee_id else "N/A"
        return f"{user_name} ({emp_id}) - {time_str}"

    @staticmethod
    def _get_font(size: int = 24) -> ImageFont.FreeTypeFont:
        """Get a font for the watermark. Falls back to default if custom font not available."""
        try:
            # Try to use a common system font (macOS)
            return ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", size)
        except (OSError, IOError):
            try:
                # Try Linux font
                return ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", size)
            except (OSError, IOError):
                try:
                    # Try Windows font
                    return ImageFont.truetype("C:/Windows/Fonts/arial.ttf", size)
                except (OSError, IOError):
                    # Fall back to default bitmap font
                    return ImageFont.load_default()

    def add_image_watermark(
        self,
        image_bytes: bytes,
        user_name: str,
        employee_id: Optional[str] = None,
        download_time: Optional[datetime] = None
    ) -> Tuple[bytes, str]:
        """
        Add a semi-transparent diagonal watermark to an image.

        Args:
            image_bytes: The original image as bytes
            user_name: Name of the user downloading the file
            employee_id: Employee ID of the user (工號)
            download_time: Time of download (defaults to now)

        Returns:
            Tuple of (watermarked image bytes, output format)

        Raises:
            Exception: If watermarking fails
        """
        # Open the image
        original = Image.open(io.BytesIO(image_bytes))

        # Convert to RGBA if necessary for transparency support
        if original.mode != 'RGBA':
            image = original.convert('RGBA')
        else:
            image = original.copy()

        # Create a transparent overlay for the watermark
        watermark_layer = Image.new('RGBA', image.size, (255, 255, 255, 0))
        draw = ImageDraw.Draw(watermark_layer)

        # Get watermark text
        watermark_text = self._format_watermark_text(user_name, employee_id, download_time)

        # Get font
        font = self._get_font(self.WATERMARK_FONT_SIZE)

        # Calculate text size
        bbox = draw.textbbox((0, 0), watermark_text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]

        # Create a larger canvas for the rotated text pattern
        diagonal = int(math.sqrt(image.size[0]**2 + image.size[1]**2))
        pattern_size = diagonal * 2

        # Create pattern layer
        pattern = Image.new('RGBA', (pattern_size, pattern_size), (255, 255, 255, 0))
        pattern_draw = ImageDraw.Draw(pattern)

        # Draw repeated watermark text across the pattern
        opacity = int(255 * self.WATERMARK_OPACITY)
        watermark_color = (*self.WATERMARK_COLOR, opacity)

        y = 0
        row = 0
        while y < pattern_size:
            x = -text_width if row % 2 else 0  # Offset alternate rows
            while x < pattern_size:
                pattern_draw.text((x, y), watermark_text, font=font, fill=watermark_color)
                x += text_width + self.WATERMARK_SPACING
            y += text_height + self.WATERMARK_SPACING
            row += 1

        # Rotate the pattern
        rotated_pattern = pattern.rotate(
            self.WATERMARK_ANGLE,
            expand=False,
            center=(pattern_size // 2, pattern_size // 2)
        )

        # Crop to original image size (centered)
        crop_x = (pattern_size - image.size[0]) // 2
        crop_y = (pattern_size - image.size[1]) // 2
        cropped_pattern = rotated_pattern.crop((
            crop_x, crop_y,
            crop_x + image.size[0],
            crop_y + image.size[1]
        ))

        # Composite the watermark onto the image
        watermarked = Image.alpha_composite(image, cropped_pattern)

        # Determine output format
        original_format = original.format or 'PNG'
        if original_format.upper() == 'JPEG':
            # Convert back to RGB for JPEG (no alpha channel)
            watermarked = watermarked.convert('RGB')
            output_format = 'JPEG'
        else:
            output_format = 'PNG'

        # Save to bytes
        output = io.BytesIO()
        watermarked.save(output, format=output_format, quality=95)
        output.seek(0)

        logger.info(
            f"Image watermark applied successfully for user {user_name} "
            f"(employee_id: {employee_id})"
        )

        return output.getvalue(), output_format.lower()

    def add_pdf_watermark(
        self,
        pdf_bytes: bytes,
        user_name: str,
        employee_id: Optional[str] = None,
        download_time: Optional[datetime] = None
    ) -> bytes:
        """
        Add a semi-transparent diagonal watermark to a PDF using PyMuPDF.

        Args:
            pdf_bytes: The original PDF as bytes
            user_name: Name of the user downloading the file
            employee_id: Employee ID of the user (工號)
            download_time: Time of download (defaults to now)

        Returns:
            Watermarked PDF as bytes

        Raises:
            Exception: If watermarking fails
        """
        # Get watermark text
        watermark_text = self._format_watermark_text(user_name, employee_id, download_time)

        # Open the PDF with PyMuPDF
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        page_count = len(doc)

        # Process each page
        for page_num in range(page_count):
            page = doc[page_num]
            page_rect = page.rect
            page_width = page_rect.width
            page_height = page_rect.height

            # Calculate text width for spacing estimation
            text_length = fitz.get_text_length(
                watermark_text,
                fontname="helv",
                fontsize=self.WATERMARK_FONT_SIZE
            )

            # Calculate diagonal for watermark coverage
            diagonal = math.sqrt(page_width**2 + page_height**2)

            # Set watermark color with opacity (gray with 30% opacity)
            color = (0.5, 0.5, 0.5)  # Gray

            # Calculate rotation angle in radians
            angle_rad = math.radians(self.WATERMARK_ANGLE)

            # Draw watermark pattern using shape with proper rotation
            # We use insert_textbox with a morph transform for rotation
            spacing_x = text_length + self.WATERMARK_SPACING
            spacing_y = self.WATERMARK_FONT_SIZE + self.WATERMARK_SPACING

            # Create watermark by drawing rotated text lines
            # We'll use a simpler approach: draw text and apply rotation via morph
            shape = page.new_shape()

            # Calculate grid positions to cover the page when rotated
            center = fitz.Point(page_width / 2, page_height / 2)

            # Calculate start and end points for coverage
            start = -diagonal
            end = diagonal * 2

            y = start
            row = 0
            while y < end:
                x = start + (spacing_x / 2 if row % 2 else 0)
                while x < end:
                    # Create text position
                    text_point = fitz.Point(x, y)

                    # Apply rotation around center
                    cos_a = math.cos(angle_rad)
                    sin_a = math.sin(angle_rad)

                    # Translate to origin, rotate, translate back
                    rx = text_point.x - center.x
                    ry = text_point.y - center.y

                    new_x = rx * cos_a - ry * sin_a + center.x
                    new_y = rx * sin_a + ry * cos_a + center.y

                    # Check if the rotated point is within page bounds (with margin)
                    margin = 50
                    if (-margin <= new_x <= page_width + margin and
                            -margin <= new_y <= page_height + margin):
                        # Insert text using shape with rotation via morph
                        text_rect = fitz.Rect(new_x, new_y, new_x + text_length + 10, new_y + 30)

                        # Use insert_textbox with morph for rotation
                        pivot = fitz.Point(new_x, new_y)
                        morph = (pivot, fitz.Matrix(1, 0, 0, 1, 0, 0).prerotate(self.WATERMARK_ANGLE))

                        shape.insert_textbox(
                            text_rect,
                            watermark_text,
                            fontname="helv",
                            fontsize=self.WATERMARK_FONT_SIZE,
                            color=color,
                            fill_opacity=self.WATERMARK_OPACITY,
                            morph=morph
                        )

                    x += spacing_x
                y += spacing_y
                row += 1

            # Commit the shape drawings
            shape.commit(overlay=True)

        # Save to bytes
        output = io.BytesIO()
        doc.save(output)
        doc.close()
        output.seek(0)

        logger.info(
            f"PDF watermark applied successfully for user {user_name} "
            f"(employee_id: {employee_id}), pages: {page_count}"
        )

        return output.getvalue()

    def is_supported_image(self, mime_type: str) -> bool:
        """Check if the mime type is a supported image format."""
        supported_types = {'image/png', 'image/jpeg', 'image/jpg'}
        return mime_type.lower() in supported_types

    def is_supported_pdf(self, mime_type: str) -> bool:
        """Check if the mime type is a PDF."""
        return mime_type.lower() == 'application/pdf'

    def supports_watermark(self, mime_type: str) -> bool:
        """Check if the file type supports watermarking."""
        return self.is_supported_image(mime_type) or self.is_supported_pdf(mime_type)


# Singleton instance
watermark_service = WatermarkService()