# 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


# ── QR Code Watermark ────────────────────────────────────────────────


def apply_qr_watermark(
    image_path: str,
    data: str,
    position: str = "bottom-right",
    size_ratio: float = 0.15,
    opacity: float = 0.4,
    output_path: str = None,
) -> str:
    """Apply a QR code watermark to an image.

    Args:
        image_path: Source image
        data: URL or text to encode in QR code
        position: bottom-right, bottom-left, top-right, top-left, center
        size_ratio: QR size as fraction of image width (default 15%)
        opacity: Transparency (0-1)
        output_path: Output path (default: overwrites original)
    """
    _ensure_pillow()
    from PIL import Image

    if not output_path:
        output_path = image_path

    img = Image.open(image_path).convert("RGBA")

    # Generate QR code
    try:
        import qrcode

        qr = qrcode.QRCode(version=1, box_size=10, border=1)
        qr.add_data(data)
        qr.make(fit=True)
        qr_img = qr.make_image(fill_color="black", back_color="white").convert("RGBA")
    except ImportError:
        # Fallback: create a simple text box if qrcode not installed
        from PIL import ImageDraw, ImageFont

        qr_img = Image.new("RGBA", (200, 200), (255, 255, 255, 255))
        draw = ImageDraw.Draw(qr_img)
        try:
            font = ImageFont.truetype("DejaVuSans.ttf", 10)
        except (IOError, OSError):
            font = ImageFont.load_default()
        draw.text((10, 80), f"QR: {data[:30]}", fill="black", font=font)

    # Resize QR to fit
    qr_size = int(img.width * size_ratio)
    qr_img = qr_img.resize((qr_size, qr_size), Image.LANCZOS)

    # Apply opacity
    alpha = qr_img.split()[-1]
    alpha = alpha.point(lambda p: int(p * opacity))
    qr_img.putalpha(alpha)

    # Position
    margin = int(img.width * 0.02)
    positions = {
        "bottom-right": (img.width - qr_size - margin, img.height - qr_size - margin),
        "bottom-left": (margin, img.height - qr_size - margin),
        "top-right": (img.width - qr_size - margin, margin),
        "top-left": (margin, margin),
        "center": ((img.width - qr_size) // 2, (img.height - qr_size) // 2),
    }
    pos = positions.get(position, positions["bottom-right"])

    # Composite
    img.paste(qr_img, pos, qr_img)
    img = img.convert("RGB")
    img.save(output_path)
    return output_path


# ── Logo Watermark ───────────────────────────────────────────────────


def apply_logo_watermark(
    image_path: str,
    logo_path: str,
    position: str = "bottom-right",
    size_ratio: float = 0.12,
    opacity: float = 0.5,
    output_path: str = None,
) -> str:
    """Overlay a logo image as watermark.

    Args:
        image_path: Source image
        logo_path: Path to logo PNG (transparent background recommended)
        position: bottom-right, bottom-left, top-right, top-left, center, tiled
        size_ratio: Logo size as fraction of image width
        opacity: Transparency (0-1)
        output_path: Output path
    """
    _ensure_pillow()
    from PIL import Image

    if not output_path:
        output_path = image_path

    img = Image.open(image_path).convert("RGBA")
    logo = Image.open(logo_path).convert("RGBA")

    # Resize logo
    logo_width = int(img.width * size_ratio)
    aspect = logo.height / logo.width
    logo_height = int(logo_width * aspect)
    logo = logo.resize((logo_width, logo_height), Image.LANCZOS)

    # Apply opacity
    alpha = logo.split()[-1]
    alpha = alpha.point(lambda p: int(p * opacity))
    logo.putalpha(alpha)

    if position == "tiled":
        # Tile across entire image
        for y in range(0, img.height, logo_height + 20):
            for x in range(0, img.width, logo_width + 20):
                img.paste(logo, (x, y), logo)
    else:
        margin = int(img.width * 0.02)
        positions = {
            "bottom-right": (img.width - logo_width - margin, img.height - logo_height - margin),
            "bottom-left": (margin, img.height - logo_height - margin),
            "top-right": (img.width - logo_width - margin, margin),
            "top-left": (margin, margin),
            "center": ((img.width - logo_width) // 2, (img.height - logo_height) // 2),
        }
        pos = positions.get(position, positions["bottom-right"])
        img.paste(logo, pos, logo)

    img = img.convert("RGB")
    img.save(output_path)
    return output_path


# ── Watermark Templates ──────────────────────────────────────────────

WATERMARK_TEMPLATES_FILE = Path.home() / ".familiar" / "data" / "gallery" / "watermark_templates.json"


def save_watermark_template(name: str, config: dict) -> dict:
    """Save a watermark preset template."""
    import json

    templates = load_watermark_templates()
    templates[name] = {**config, "name": name}
    WATERMARK_TEMPLATES_FILE.parent.mkdir(parents=True, exist_ok=True)
    WATERMARK_TEMPLATES_FILE.write_text(json.dumps(templates, indent=2))
    return templates[name]


def load_watermark_templates() -> dict:
    """Load all saved watermark templates."""
    import json

    if WATERMARK_TEMPLATES_FILE.exists():
        try:
            return json.loads(WATERMARK_TEMPLATES_FILE.read_text())
        except Exception:
            pass
    return {}


def delete_watermark_template(name: str) -> bool:
    """Delete a watermark template by name."""
    import json

    templates = load_watermark_templates()
    if name in templates:
        del templates[name]
        WATERMARK_TEMPLATES_FILE.write_text(json.dumps(templates, indent=2))
        return True
    return False
