"""Write artwork (ithmb + ArtworkDB) to iPod.

Ported from iOpenPod ArtworkDB_Writer/artwork_writer.py, simplified
for our use case (Nano 7, always-new tracks, no incremental sync).
"""

from __future__ import annotations

import hashlib
import logging
import os

from PIL import Image

from .presets import NANO_7G_COVER_ART_FORMATS, ArtworkFormat
from .types import ArtworkEntry, EncodedFormatPayload, ExistingFormatRef, IthmbLocation
from .codecs import encode_for_ipod
from .chunks import build_artworkdb
from .extractor import extract_cover_with_fallback

logger = logging.getLogger(__name__)

ITHMB_MAX_SIZE_BYTES = 256 * 1000 * 1000


def _ithmb_filename(format_id: int, index: int = 1) -> str:
    return f"F{int(format_id)}_{int(index)}.ithmb"


def _art_hash(art_bytes: bytes) -> str:
    return hashlib.md5(art_bytes).hexdigest()


def image_from_bytes(art_bytes: bytes) -> Image.Image | None:
    """Load a PIL Image from raw bytes (JPEG/PNG), ensuring RGB mode."""
    from io import BytesIO
    try:
        img = Image.open(BytesIO(art_bytes))
        if img.mode != "RGB":
            img = img.convert("RGBA").convert("RGB")
        return img
    except Exception:
        return None


def prepare_artwork_for_track(source_path: str) -> tuple[str, dict[int, EncodedFormatPayload]] | None:
    """Extract cover art from a source file and encode as iPod formats.

    Returns (art_hash, {format_id: EncodedFormatPayload}) or None.
    """
    art_bytes = extract_cover_with_fallback(source_path)
    if art_bytes is None:
        return None

    img = image_from_bytes(art_bytes)
    if img is None:
        return None

    formats: dict[int, EncodedFormatPayload] = {}
    for fmt in NANO_7G_COVER_ART_FORMATS:
        try:
            encoded = encode_for_ipod(img, fmt.format_id, fmt_override=fmt)
            formats[fmt.format_id] = encoded
        except Exception as exc:
            logger.warning("Failed to encode format %d: %s", fmt.format_id, exc)

    if not formats:
        return None

    return _art_hash(art_bytes), formats


def write_artworkdb(
    ipod_path: str,
    entries: list[ArtworkEntry],
    start_img_id: int = 100,
) -> dict[int, tuple[int, int]]:
    """Write ithmb files and ArtworkDB binary to the iPod.

    Args:
        ipod_path: iPod mount point
        entries: List of ArtworkEntry (one per unique album art)
        start_img_id: Starting image ID

    Returns:
        Dict mapping db_track_id → (img_id, src_img_size)
    """
    artwork_dir = os.path.join(ipod_path, "iPod_Control", "Artwork")
    os.makedirs(artwork_dir, exist_ok=True)

    # Collect all unique format IDs across all entries
    format_ids = sorted({fmt_id for entry in entries for fmt_id in entry.formats.keys()})
    writable_format_ids = sorted({
        fmt_id for entry in entries
        for fmt_id, payload in entry.formats.items()
        if isinstance(payload, EncodedFormatPayload)
    })

    # Write ithmb files
    ithmb_state: dict[int, dict[str, int]] = {
        fmt_id: {"index": 0, "offset": 0}
        for fmt_id in writable_format_ids
    }
    ithmb_files: dict[int, object] = {}
    format_locations_map: dict[int, dict[int, IthmbLocation]] = {}
    image_sizes: dict[int, int] = {}

    def _close_current(fmt_id: int) -> None:
        handle = ithmb_files.pop(fmt_id, None)
        if handle is not None:
            handle.close()

    def _open_next(fmt_id: int) -> None:
        _close_current(fmt_id)
        state = ithmb_state[fmt_id]
        state["index"] += 1
        state["offset"] = 0
        filename = _ithmb_filename(fmt_id, state["index"])
        ithmb_files[fmt_id] = open(os.path.join(artwork_dir, filename), "ab")

    def _write_payload(fmt_id: int, data: bytes) -> IthmbLocation:
        state = ithmb_state[fmt_id]
        if fmt_id not in ithmb_files:
            _open_next(fmt_id)
        else:
            f = ithmb_files[fmt_id]
            f.seek(0, 2)
            cur_size = f.tell()
            if cur_size > 0 and cur_size + len(data) > ITHMB_MAX_SIZE_BYTES:
                _open_next(fmt_id)
        f = ithmb_files[fmt_id]
        f.seek(0, 2)
        offset = f.tell()
        f.write(data)
        state["offset"] = offset + len(data)
        filename = _ithmb_filename(fmt_id, state["index"])
        return IthmbLocation(filename, offset)

    try:
        for entry in entries:
            locations: dict[int, IthmbLocation] = {}
            for fmt_id in format_ids:
                if fmt_id not in entry.formats:
                    continue
                payload = entry.formats[fmt_id]
                if isinstance(payload, EncodedFormatPayload):
                    locations[fmt_id] = _write_payload(fmt_id, payload.data)
                elif isinstance(payload, ExistingFormatRef):
                    fname = _ithmb_filename(fmt_id)
                    locations[fmt_id] = IthmbLocation(fname, payload.ithmb_offset)
            format_locations_map[entry.img_id] = locations
    finally:
        for fmt_id in list(ithmb_files.keys()):
            _close_current(fmt_id)

    # Calculate image sizes for MHIF
    for fmt_id in format_ids:
        sizes = [
            int(entry.formats[fmt_id].size)
            for entry in entries
            if fmt_id in entry.formats and int(entry.formats[fmt_id].size) > 0
        ]
        if sizes:
            from collections import Counter
            image_sizes[fmt_id] = Counter(sizes).most_common(1)[0][0]

    # Build and write ArtworkDB binary
    ref_mhfd = None
    existing_artdb = os.path.join(artwork_dir, "ArtworkDB")
    if os.path.exists(existing_artdb):
        with open(existing_artdb, "rb") as f:
            ref_mhfd = f.read()

    next_id = start_img_id + len(entries)
    db_data = build_artworkdb(
        entries,
        format_locations_map,
        format_ids,
        image_sizes,
        next_id,
        ref_mhfd,
    )

    with open(existing_artdb, "wb") as f:
        f.write(db_data)
        f.flush()
        os.fsync(f.fileno())

    # Return track_id → (img_id, src_img_size) mapping
    return {
        entry.db_track_id: (entry.img_id, entry.src_img_size)
        for entry in entries
    }
