import logging
from pathlib import Path

import PIL.Image

log = logging.getLogger(__name__)


class MetadataBuilder:
    @staticmethod
    def embed_metadata(metadata: dict, path: str | Path) -> None:
        # Check if file is PNG format
        path_obj = Path(path) if isinstance(path, str) else path
        if path_obj.suffix.lower() != ".png":
            log.warning(f"XMP/IPTC metadata embedding is only supported for PNG files, skipping: {path_obj}")
            return

        try:
            from PIL import PngImagePlugin

            # Load the image preserving existing metadata
            image = PIL.Image.open(path)

            # Get existing PNG info to preserve it
            existing_info = image.info if hasattr(image, "info") else {}

            # Preserve existing EXIF separately (if it exists)
            existing_exif = existing_info.get("exif")

            # Create new PngInfo preserving existing data
            pnginfo = PngImagePlugin.PngInfo()

            # Copy existing metadata
            for key, value in existing_info.items():
                if key not in ["XML:com.adobe.xmp", "IPTC", "exif"]:  # Handle these separately
                    pnginfo.add_text(key, str(value))

            # Build XMP and IPTC metadata using builder methods
            xmp_packet = MetadataBuilder.build_xmp_packet(metadata)
            iptc_binary = MetadataBuilder.build_iptc_binary(metadata)

            # Add XMP and IPTC to PNG info
            pnginfo.add_text("XML:com.adobe.xmp", xmp_packet)
            if iptc_binary:
                pnginfo.add_text("IPTC", iptc_binary.hex())

            # Save preserving ALL existing metadata + adding XMP/IPTC
            # Pass exif separately to preserve it correctly
            if existing_exif:
                image.save(path, pnginfo=pnginfo, exif=existing_exif)
            else:
                image.save(path, pnginfo=pnginfo)

        except Exception as e:  # noqa: BLE001
            log.error(f"Error embedding XMP/IPTC metadata: {e}")

    @staticmethod
    def build_xmp_packet(metadata: dict) -> str:
        # Escape prompt for XML
        prompt_escaped = metadata.get("prompt", "").replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")

        # Build LoRA info for XMP
        lora_info = MetadataBuilder._build_lora_string(metadata)

        # Get version from metadata
        version = metadata.get("mflux_version", "unknown")

        xmp_packet = f"""<?xpacket begin="﻿" id="W5M0MpCehiHzreSzNTczkc9d"?>
<x:xmpmeta xmlns:x="adobe:ns:meta/">
<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">
<rdf:Description rdf:about=""
    xmlns:dc="http://purl.org/dc/elements/1.1/"
    xmlns:xmp="http://ns.adobe.com/xap/1.0/"
    xmlns:photoshop="http://ns.adobe.com/photoshop/1.0/"
    xmlns:mflux="http://ns.mflux.ai/1.0/">
    <dc:description><rdf:Alt><rdf:li xml:lang="x-default">{prompt_escaped}</rdf:li></rdf:Alt></dc:description>
    <dc:creator><rdf:Seq><rdf:li>MFLUX</rdf:li></rdf:Seq></dc:creator>
    <dc:rights><rdf:Alt><rdf:li xml:lang="x-default">AI Generated Content</rdf:li></rdf:Alt></dc:rights>
    <xmp:CreatorTool>MFLUX {version}</xmp:CreatorTool>
    <photoshop:Category>ART</photoshop:Category>
    <photoshop:Credit>Generated by MFLUX</photoshop:Credit>"""

        # Add technical parameters to XMP
        if "seed" in metadata:
            xmp_packet += f"\n    <mflux:seed>{metadata['seed']}</mflux:seed>"
        if "steps" in metadata:
            xmp_packet += f"\n    <mflux:steps>{metadata['steps']}</mflux:steps>"
        if "guidance" in metadata:
            xmp_packet += f"\n    <mflux:guidance>{metadata['guidance']}</mflux:guidance>"
        if "model_config" in metadata:
            xmp_packet += f"\n    <mflux:model>{metadata['model_config']}</mflux:model>"
        if lora_info:
            xmp_packet += f"\n    <mflux:loras>{lora_info}</mflux:loras>"
        if "generation_time" in metadata:
            xmp_packet += f"\n    <mflux:generationTime>{metadata['generation_time']}</mflux:generationTime>"

        xmp_packet += """
</rdf:Description>
</rdf:RDF>
</x:xmpmeta>
<?xpacket end="w"?>"""

        return xmp_packet

    @staticmethod
    def build_iptc_binary(metadata: dict) -> bytes:
        # Build LoRA info for IPTC
        lora_info = MetadataBuilder._build_lora_string(metadata)

        iptc_data = {}

        # Add prompt information
        if "prompt" in metadata:
            prompt = metadata["prompt"]
            prompt_encoded = prompt.encode("utf-8")
            if len(prompt_encoded) > 2000:
                log.warning(f"Prompt is too long ({len(prompt_encoded)} bytes), truncating to 2000 bytes for IPTC")
                iptc_data[120] = prompt_encoded[:2000]  # Caption/Description
            else:
                iptc_data[120] = prompt_encoded  # Caption/Description
            iptc_data[5] = f"AI: {prompt[:50]}...".encode("utf-8")  # Object Name/Title
            iptc_data[105] = f"AI Generated: {prompt[:80]}...".encode("utf-8")  # Headline

        # Add standard fields
        iptc_data[80] = b"MFLUX"  # By-line (Creator)
        iptc_data[85] = b"AI Artist"  # By-line Title
        iptc_data[15] = b"ART"  # Category
        iptc_data[110] = b"Generated by MFLUX"  # Credit
        iptc_data[115] = b"AI Generation"  # Source
        iptc_data[116] = b"AI Generated Content"  # Copyright Notice
        iptc_data[118] = b"AI Generated using MFLUX"  # Contact
        iptc_data[103] = b"AI"  # Instructions/Special Instructions

        # Add seed and model info in specific fields
        if "seed" in metadata:
            iptc_data[122] = f"Seed: {metadata['seed']}".encode("utf-8")  # Writer/Editor

        if "model_config" in metadata:
            iptc_data[90] = f"Model: {metadata['model_config']}".encode("utf-8")  # City

        # Add LoRA info in Province/State field
        if lora_info:
            iptc_data[95] = f"LoRA: {lora_info}".encode("utf-8")  # Province/State

        # Add generation parameters in Country field
        if "steps" in metadata and "guidance" in metadata:
            iptc_data[101] = f"Steps:{metadata['steps']} CFG:{metadata['guidance']}".encode("utf-8")  # Country

        # Build keywords including LoRA info
        keywords = ["AI", "Generated", "MFLUX"]
        if "seed" in metadata:
            keywords.append(f"seed-{metadata['seed']}")
        if "steps" in metadata:
            keywords.append(f"steps-{metadata['steps']}")
        if "guidance" in metadata:
            keywords.append(f"guidance-{metadata['guidance']}")
        if "model_config" in metadata:
            keywords.append(f"model-{metadata['model_config']}")
        if lora_info:
            keywords.append(f"loras-{lora_info}")

        iptc_data[25] = ";".join(keywords).encode("utf-8")  # Keywords

        # Build IPTC binary (IPTC field length limit is 32767 bytes)
        iptc_binary = b""
        for tag_id, value in iptc_data.items():
            length = len(value)
            if length < 32768:
                iptc_binary += bytes([0x1C, 0x02, tag_id]) + length.to_bytes(2, "big") + value

        return iptc_binary

    @staticmethod
    def _build_lora_string(metadata: dict) -> str:
        # Check if lora_paths exists and is not None/empty
        if "lora_paths" not in metadata or metadata["lora_paths"] is None or not metadata["lora_paths"]:
            return ""

        lora_list = []
        lora_paths = metadata["lora_paths"]
        lora_scales = metadata.get("lora_scales", [])

        for i, lora_path in enumerate(lora_paths):
            # Use Path for OS-agnostic path handling (works on Windows, Mac, Linux)
            lora_name = Path(lora_path).name
            scale = lora_scales[i] if i < len(lora_scales) else "1.0"
            lora_list.append(f"{lora_name}:{scale}")

        return ", ".join(lora_list)
