"""
Proof Link Embedding Module for Verisoo.

This module handles embedding proof links into images after attestation.
Two strategies are used depending on the image format:

1. JPEG: DCT steganography via stegasoo
   - Embeds in frequency domain (survives recompression)
   - Uses center region for robustness against cropping
   - Invisible to human eye

2. PNG/RAW: XMP sidecar file
   - Standard metadata format used by Lightroom/Darktable/etc.
   - Travels with the image file
   - Not steganographic (plaintext XML)

The proof link format: https://verisoo.io/v/{short_id}
"""

from __future__ import annotations

import io
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

from PIL import Image

# Stegasoo integration — imported as a pip dependency (no path hacks needed).
# Install stegasoo[dct] to enable DCT steganography for JPEG proof embedding.
try:
    from soosef.stegasoo.dct_steganography import (
        embed_in_dct,
        extract_from_dct,
        has_dct_support,
        has_jpegio_support,
        calculate_dct_capacity,
    )
    HAS_STEGASOO = True
except ImportError:
    HAS_STEGASOO = False
    has_dct_support = lambda: False
    has_jpegio_support = lambda: False


# =============================================================================
# CONSTANTS
# =============================================================================

# Fixed public seed for Verisoo proof links
# This is intentionally public - anyone should be able to extract the proof link
VERISOO_SEED = b"verisoo"

# Base URL for proof links
DEFAULT_BASE_URL = "https://verisoo.io"

# XMP namespace for Verisoo
XMP_NAMESPACE = "https://verisoo.io/ns/1.0/"

# Supported formats for DCT embedding
DCT_FORMATS = {".jpg", ".jpeg"}

# Formats that need XMP sidecar
XMP_FORMATS = {".png", ".tiff", ".tif", ".webp", ".heic", ".heif"}
RAW_FORMATS = {".cr2", ".cr3", ".nef", ".arw", ".raf", ".orf", ".rw2", ".dng", ".raw"}


# =============================================================================
# DATA CLASSES
# =============================================================================

@dataclass
class EmbedResult:
    """Result of a proof link embedding operation."""
    success: bool
    method: str  # "dct", "xmp", "none"
    output_path: Path | None
    proof_link: str
    message: str


@dataclass
class ExtractResult:
    """Result of extracting a proof link from an image."""
    success: bool
    method: str  # "dct", "xmp", "none"
    proof_link: str | None
    message: str


# =============================================================================
# XMP SIDECAR GENERATION
# =============================================================================

def generate_xmp_sidecar(
    proof_link: str,
    fingerprint: str,
    attested_at: datetime,
    image_sha256: str,
    caption: str | None = None,
) -> str:
    """
    Generate XMP sidecar XML content for a Verisoo attestation.

    This creates a standard XMP file that can be read by Lightroom,
    Darktable, and other photo management software.

    Args:
        proof_link: Full proof URL (e.g., "https://verisoo.io/v/abc123")
        fingerprint: Attestor's fingerprint
        attested_at: Attestation timestamp
        image_sha256: SHA-256 hash of the image
        caption: Optional caption/description

    Returns:
        XMP XML content as string
    """
    # Escape XML special characters
    def escape(s: str) -> str:
        return (s.replace("&", "&amp;")
                 .replace("<", "&lt;")
                 .replace(">", "&gt;")
                 .replace('"', "&quot;"))

    caption_attr = f'\n      verisoo:Caption="{escape(caption)}"' if caption else ""

    return f'''<?xml version="1.0" encoding="UTF-8"?>
<x:xmpmeta xmlns:x="adobe:ns:meta/" x:xmptk="Verisoo">
  <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">
    <rdf:Description rdf:about=""
      xmlns:verisoo="{XMP_NAMESPACE}"
      verisoo:ProofLink="{escape(proof_link)}"
      verisoo:Fingerprint="{escape(fingerprint)}"
      verisoo:AttestedAt="{attested_at.isoformat()}"
      verisoo:ImageSHA256="{escape(image_sha256)}"{caption_attr}/>
  </rdf:Description>
</x:xmpmeta>
'''


def write_xmp_sidecar(
    image_path: Path,
    proof_link: str,
    fingerprint: str,
    attested_at: datetime,
    image_sha256: str,
    caption: str | None = None,
) -> Path:
    """
    Write an XMP sidecar file alongside an image.

    The sidecar is named {image_name}.xmp and placed in the same directory.

    Args:
        image_path: Path to the original image
        proof_link: Full proof URL
        fingerprint: Attestor's fingerprint
        attested_at: Attestation timestamp
        image_sha256: SHA-256 hash of the image
        caption: Optional caption

    Returns:
        Path to the created XMP sidecar file
    """
    xmp_content = generate_xmp_sidecar(
        proof_link=proof_link,
        fingerprint=fingerprint,
        attested_at=attested_at,
        image_sha256=image_sha256,
        caption=caption,
    )

    # Standard sidecar naming: image.jpg -> image.jpg.xmp
    # Some software prefers image.xmp, but .jpg.xmp is more explicit
    xmp_path = image_path.with_suffix(image_path.suffix + ".xmp")
    xmp_path.write_text(xmp_content, encoding="utf-8")

    return xmp_path


def read_xmp_sidecar(image_path: Path) -> dict[str, str] | None:
    """
    Read Verisoo metadata from an XMP sidecar file.

    Args:
        image_path: Path to the image (sidecar path is derived)

    Returns:
        Dict with proof_link, fingerprint, etc. or None if not found
    """
    import re

    xmp_path = image_path.with_suffix(image_path.suffix + ".xmp")
    if not xmp_path.exists():
        # Try alternate naming
        xmp_path = image_path.with_suffix(".xmp")
        if not xmp_path.exists():
            return None

    try:
        content = xmp_path.read_text(encoding="utf-8")

        # Simple regex extraction (not full XML parsing for speed)
        result = {}

        patterns = {
            "proof_link": r'verisoo:ProofLink="([^"]*)"',
            "fingerprint": r'verisoo:Fingerprint="([^"]*)"',
            "attested_at": r'verisoo:AttestedAt="([^"]*)"',
            "image_sha256": r'verisoo:ImageSHA256="([^"]*)"',
            "caption": r'verisoo:Caption="([^"]*)"',
        }

        for key, pattern in patterns.items():
            match = re.search(pattern, content)
            if match:
                result[key] = match.group(1)

        return result if result else None

    except Exception:
        return None


# =============================================================================
# DCT EMBEDDING (JPEG)
# =============================================================================

def embed_proof_in_jpeg(
    image_data: bytes,
    proof_link: str,
) -> tuple[bytes, dict[str, Any]]:
    """
    Embed a proof link into a JPEG image using DCT steganography.

    Uses stegasoo's DCT embedding with:
    - Fixed public seed (b"verisoo") so anyone can extract
    - Center-biased embedding for crop resistance (TODO)
    - Minimal quality impact (only ~25 blocks needed)

    Args:
        image_data: Original JPEG bytes
        proof_link: Proof URL to embed (e.g., "https://verisoo.io/v/abc123")

    Returns:
        Tuple of (embedded_image_bytes, stats_dict)

    Raises:
        ImportError: If stegasoo is not available
        ValueError: If image is too small or embedding fails
    """
    if not HAS_STEGASOO:
        raise ImportError(
            "DCT embedding requires stegasoo. "
            "Ensure stegasoo is installed or available at ../stegasoo"
        )

    if not has_jpegio_support():
        raise ImportError(
            "Native JPEG embedding requires jpeglib. "
            "Install with: pip install jpeglib"
        )

    # Embed the proof link
    payload = proof_link.encode("utf-8")

    stego_bytes, stats = embed_in_dct(
        data=payload,
        carrier_image=image_data,
        seed=VERISOO_SEED,
        output_format="jpeg",
        color_mode="color",
    )

    return stego_bytes, {
        "blocks_used": stats.blocks_used,
        "blocks_available": stats.blocks_available,
        "usage_percent": stats.usage_percent,
        "output_format": stats.output_format,
    }


def extract_proof_from_jpeg(image_data: bytes) -> str | None:
    """
    Extract a proof link from a JPEG image.

    Args:
        image_data: JPEG image bytes

    Returns:
        Proof URL string or None if not found/invalid
    """
    if not HAS_STEGASOO:
        return None

    try:
        payload = extract_from_dct(
            stego_image=image_data,
            seed=VERISOO_SEED,
        )

        # Validate it looks like a proof link
        proof_link = payload.decode("utf-8")
        if "verisoo" in proof_link.lower() or proof_link.startswith("http"):
            return proof_link
        return None

    except Exception:
        return None


# =============================================================================
# UNIFIED EMBEDDING API
# =============================================================================

def get_embed_method(image_path: Path) -> str:
    """
    Determine the appropriate embedding method for an image.

    Args:
        image_path: Path to the image file

    Returns:
        "dct" for JPEG, "xmp" for PNG/RAW, "none" if unsupported
    """
    suffix = image_path.suffix.lower()

    if suffix in DCT_FORMATS:
        if HAS_STEGASOO and has_jpegio_support():
            return "dct"
        else:
            return "xmp"  # Fallback to XMP if stegasoo unavailable

    if suffix in XMP_FORMATS or suffix in RAW_FORMATS:
        return "xmp"

    return "none"


def embed_proof_link(
    image_path: Path,
    proof_link: str,
    fingerprint: str,
    attested_at: datetime,
    image_sha256: str,
    caption: str | None = None,
    output_path: Path | None = None,
) -> EmbedResult:
    """
    Embed a proof link into an image using the appropriate method.

    For JPEG: Uses DCT steganography (invisible, survives compression)
    For PNG/RAW: Creates XMP sidecar file

    Args:
        image_path: Path to the original image
        proof_link: Full proof URL to embed
        fingerprint: Attestor's fingerprint
        attested_at: Attestation timestamp
        image_sha256: SHA-256 hash of the image
        caption: Optional caption
        output_path: Optional output path for embedded image (JPEG only)

    Returns:
        EmbedResult with success status, method used, and output path
    """
    method = get_embed_method(image_path)

    if method == "dct":
        try:
            image_data = image_path.read_bytes()
            stego_bytes, stats = embed_proof_in_jpeg(image_data, proof_link)

            # Determine output path
            if output_path is None:
                stem = image_path.stem
                output_path = image_path.parent / f"{stem}_attested.jpg"

            output_path.write_bytes(stego_bytes)

            return EmbedResult(
                success=True,
                method="dct",
                output_path=output_path,
                proof_link=proof_link,
                message=f"Embedded via DCT ({stats['usage_percent']:.2f}% capacity used)",
            )

        except Exception as e:
            # Fall back to XMP on DCT failure
            return _embed_xmp_fallback(
                image_path, proof_link, fingerprint, attested_at,
                image_sha256, caption, str(e)
            )

    elif method == "xmp":
        try:
            xmp_path = write_xmp_sidecar(
                image_path=image_path,
                proof_link=proof_link,
                fingerprint=fingerprint,
                attested_at=attested_at,
                image_sha256=image_sha256,
                caption=caption,
            )

            return EmbedResult(
                success=True,
                method="xmp",
                output_path=xmp_path,
                proof_link=proof_link,
                message=f"Created XMP sidecar: {xmp_path.name}",
            )

        except Exception as e:
            return EmbedResult(
                success=False,
                method="xmp",
                output_path=None,
                proof_link=proof_link,
                message=f"XMP sidecar failed: {e}",
            )

    else:
        return EmbedResult(
            success=False,
            method="none",
            output_path=None,
            proof_link=proof_link,
            message=f"Unsupported format: {image_path.suffix}",
        )


def _embed_xmp_fallback(
    image_path: Path,
    proof_link: str,
    fingerprint: str,
    attested_at: datetime,
    image_sha256: str,
    caption: str | None,
    dct_error: str,
) -> EmbedResult:
    """Fallback to XMP when DCT fails."""
    try:
        xmp_path = write_xmp_sidecar(
            image_path=image_path,
            proof_link=proof_link,
            fingerprint=fingerprint,
            attested_at=attested_at,
            image_sha256=image_sha256,
            caption=caption,
        )

        return EmbedResult(
            success=True,
            method="xmp",
            output_path=xmp_path,
            proof_link=proof_link,
            message=f"DCT failed ({dct_error}), created XMP sidecar instead",
        )

    except Exception as e:
        return EmbedResult(
            success=False,
            method="none",
            output_path=None,
            proof_link=proof_link,
            message=f"Both DCT and XMP failed: {dct_error}; {e}",
        )


def extract_proof_link(image_path: Path) -> ExtractResult:
    """
    Extract a proof link from an image (DCT or XMP).

    Tries DCT extraction for JPEG, XMP sidecar for all formats.

    Args:
        image_path: Path to the image

    Returns:
        ExtractResult with proof link if found
    """
    suffix = image_path.suffix.lower()

    # Try DCT for JPEG
    if suffix in DCT_FORMATS and HAS_STEGASOO:
        try:
            image_data = image_path.read_bytes()
            proof_link = extract_proof_from_jpeg(image_data)
            if proof_link:
                return ExtractResult(
                    success=True,
                    method="dct",
                    proof_link=proof_link,
                    message="Extracted from DCT steganography",
                )
        except Exception:
            pass

    # Try XMP sidecar
    xmp_data = read_xmp_sidecar(image_path)
    if xmp_data and "proof_link" in xmp_data:
        return ExtractResult(
            success=True,
            method="xmp",
            proof_link=xmp_data["proof_link"],
            message="Extracted from XMP sidecar",
        )

    return ExtractResult(
        success=False,
        method="none",
        proof_link=None,
        message="No embedded proof link found",
    )


# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

def can_embed_dct() -> bool:
    """Check if DCT embedding is available."""
    return HAS_STEGASOO and has_jpegio_support()


def get_embed_capabilities() -> dict[str, Any]:
    """Get information about available embedding capabilities."""
    return {
        "dct_available": HAS_STEGASOO and has_dct_support(),
        "jpeg_native": HAS_STEGASOO and has_jpegio_support(),
        "xmp_available": True,  # Always available
        "supported_dct_formats": list(DCT_FORMATS) if can_embed_dct() else [],
        "supported_xmp_formats": list(XMP_FORMATS | RAW_FORMATS),
    }
