# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Image watermarking utilities.

Three watermark types:
- Visible overlay: semi-transparent text via PIL.ImageDraw + alpha compositing
- Invisible (LSB steganography): data embedded in least-significant bits of RGB channels
- Hash fingerprint: SHA-256 of raw file bytes for integrity verification

Pillow is lazily imported — available via the ``docs`` extra.
"""

import hashlib
import logging
import struct
from pathlib import Path

logger = logging.getLogger(__name__)

PILLOW_AVAILABLE = False
Image = None
ImageDraw = None
ImageFont = None

try:
    from PIL import Image as _Image
    from PIL import ImageDraw as _ImageDraw
    from PIL import ImageFont as _ImageFont

    Image = _Image
    ImageDraw = _ImageDraw
    ImageFont = _ImageFont
    PILLOW_AVAILABLE = True
except ImportError:
    logger.debug("Pillow not available — install with: pip install familiar-agent[docs]")


def _ensure_pillow():
    """Lazy-load Pillow, installing if necessary."""
    global PILLOW_AVAILABLE, Image, ImageDraw, ImageFont
    if PILLOW_AVAILABLE:
        return True
    from familiar.core.deps import ensure_packages

    ok, _ = ensure_packages([("PIL", "Pillow")])
    if ok:
        from PIL import Image as _Image
        from PIL import ImageDraw as _ImageDraw
        from PIL import ImageFont as _ImageFont

        globals().update(
            {"Image": _Image, "ImageDraw": _ImageDraw, "ImageFont": _ImageFont}
        )
        PILLOW_AVAILABLE = True
    return PILLOW_AVAILABLE


def compute_image_hash(image_path: str) -> str:
    """Compute SHA-256 hash of raw file bytes."""
    h = hashlib.sha256()
    with open(image_path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def apply_visible_watermark(
    image_path: str,
    text: str,
    position: str = "bottom-right",
    opacity: float = 0.3,
    font_size: int = 0,
    output_path: str = "",
) -> str:
    """
    Apply a semi-transparent text overlay watermark.

    Args:
        image_path: Path to source image.
        text: Watermark text to overlay.
        position: One of bottom-right, bottom-left, center, tiled.
        opacity: Text opacity 0.0–1.0 (default 0.3).
        font_size: Font size in pixels. 0 = auto-scale (~3% of image width).
        output_path: Where to save. Empty = auto-generate next to source.

    Returns:
        Path to the watermarked image file.
    """
    if not _ensure_pillow():
        raise RuntimeError("Pillow is required: pip install familiar-agent[docs]")

    src = Path(image_path)
    if not output_path:
        output_path = str(src.parent / f"{src.stem}_watermarked{src.suffix}")

    base = Image.open(image_path).convert("RGBA")
    overlay = Image.new("RGBA", base.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)

    if font_size <= 0:
        font_size = max(16, int(base.width * 0.03))

    try:
        font = ImageFont.truetype("DejaVuSans.ttf", font_size)
    except (OSError, IOError):
        font = ImageFont.load_default()

    alpha = int(255 * max(0.0, min(1.0, opacity)))
    fill = (255, 255, 255, alpha)

    bbox = draw.textbbox((0, 0), text, font=font)
    tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]

    if position == "tiled":
        spacing_x = tw + max(40, tw // 2)
        spacing_y = th + max(40, th // 2)
        y = 0
        while y < base.height:
            x = 0
            while x < base.width:
                draw.text((x, y), text, font=font, fill=fill)
                x += spacing_x
            y += spacing_y
    else:
        margin = max(10, font_size // 2)
        if position == "bottom-left":
            pos = (margin, base.height - th - margin)
        elif position == "center":
            pos = ((base.width - tw) // 2, (base.height - th) // 2)
        else:  # bottom-right (default)
            pos = (base.width - tw - margin, base.height - th - margin)
        draw.text(pos, text, font=font, fill=fill)

    result = Image.alpha_composite(base, overlay)
    # Save as original format — convert to RGB if saving as JPEG
    out = Path(output_path)
    if out.suffix.lower() in (".jpg", ".jpeg"):
        result = result.convert("RGB")
    result.save(output_path)
    return output_path


def apply_invisible_watermark(
    image_path: str,
    data: str,
    output_path: str = "",
) -> str:
    """
    Embed data into the least-significant bits of RGB pixel channels (LSB steganography).

    The data is encoded as UTF-8 with a 32-bit length prefix. Each bit is stored
    in the LSB of successive R, G, B channel values. Output is forced to PNG
    (lossless) to preserve the embedded data.

    Capacity: width * height * 3 bits.

    Args:
        image_path: Path to source image.
        data: String data to embed.
        output_path: Where to save. Empty = auto-generate as PNG next to source.

    Returns:
        Path to the watermarked PNG image.
    """
    if not _ensure_pillow():
        raise RuntimeError("Pillow is required: pip install familiar-agent[docs]")

    src = Path(image_path)
    if not output_path:
        output_path = str(src.parent / f"{src.stem}_lsb.png")
    else:
        # Force PNG extension
        out = Path(output_path)
        if out.suffix.lower() != ".png":
            output_path = str(out.with_suffix(".png"))

    img = Image.open(image_path).convert("RGB")
    pixels = list(img.getdata())
    width, height = img.size
    capacity = width * height * 3  # bits

    data_bytes = data.encode("utf-8")
    length_prefix = struct.pack(">I", len(data_bytes))
    payload = length_prefix + data_bytes
    total_bits = len(payload) * 8

    if total_bits > capacity:
        raise ValueError(
            f"Data too large: {total_bits} bits needed, {capacity} bits available"
        )

    # Flatten pixel channels into a list
    channels = []
    for r, g, b in pixels:
        channels.extend([r, g, b])

    # Embed each bit into LSB
    bit_idx = 0
    for byte in payload:
        for bit_pos in range(7, -1, -1):
            bit = (byte >> bit_pos) & 1
            channels[bit_idx] = (channels[bit_idx] & 0xFE) | bit
            bit_idx += 1

    # Reconstruct image
    new_pixels = []
    for i in range(0, len(channels), 3):
        new_pixels.append((channels[i], channels[i + 1], channels[i + 2]))

    new_img = Image.new("RGB", (width, height))
    new_img.putdata(new_pixels)
    new_img.save(output_path, "PNG")
    return output_path


def extract_invisible_watermark(image_path: str) -> str:
    """
    Extract data previously embedded via apply_invisible_watermark().

    Reads the 32-bit length prefix from the first pixel LSBs, then extracts
    and decodes the embedded UTF-8 string.

    Args:
        image_path: Path to a PNG image with LSB-embedded data.

    Returns:
        The extracted data string.
    """
    if not _ensure_pillow():
        raise RuntimeError("Pillow is required: pip install familiar-agent[docs]")

    img = Image.open(image_path).convert("RGB")
    pixels = list(img.getdata())

    # Flatten channels
    channels = []
    for r, g, b in pixels:
        channels.extend([r, g, b])

    # Read 32-bit length prefix (4 bytes = 32 bits)
    length_bits = []
    for i in range(32):
        length_bits.append(channels[i] & 1)

    length_val = 0
    for bit in length_bits:
        length_val = (length_val << 1) | bit

    if length_val <= 0 or length_val > len(channels) // 8:
        raise ValueError("No valid watermark data found")

    # Read data bytes
    data_bytes = bytearray()
    bit_offset = 32  # skip length prefix
    for _ in range(length_val):
        byte_val = 0
        for bit_pos in range(8):
            byte_val = (byte_val << 1) | (channels[bit_offset] & 1)
            bit_offset += 1
        data_bytes.append(byte_val)

    return data_bytes.decode("utf-8")


def generate_thumbnail(
    image_path: str,
    size: tuple = (300, 300),
    output_path: str = "",
) -> str:
    """
    Generate a thumbnail preserving aspect ratio.

    Args:
        image_path: Path to source image.
        size: Maximum (width, height) tuple.
        output_path: Where to save. Empty = auto-generate.

    Returns:
        Path to the thumbnail image.
    """
    if not _ensure_pillow():
        raise RuntimeError("Pillow is required: pip install familiar-agent[docs]")

    src = Path(image_path)
    if not output_path:
        output_path = str(src.parent / f"{src.stem}_thumb{src.suffix}")

    img = Image.open(image_path)
    img.thumbnail(size)
    out = Path(output_path)
    if out.suffix.lower() in (".jpg", ".jpeg"):
        img = img.convert("RGB")
    img.save(output_path)
    return output_path
