from __future__ import annotations

import hashlib
import json
import os
import tempfile
from collections.abc import Sequence
from dataclasses import dataclass
from functools import cached_property, lru_cache, partial
from itertools import combinations
from math import gcd, log
from pathlib import Path
from typing import Any, Literal

from lektor.pluginsystem import Plugin

LAYOUT_ALGORITHM_VERSION = 1

MAX_CANDIDATES_PER_SUBSET = 32
RATIO_PRECISION = 5

T_JUNCTION_WEIGHT = 0.10
SIBLING_IMBALANCE_WEIGHT = 0.015
ORIENTATION_SWITCH_WEIGHT = 0.01
DEPTH_WEIGHT = 0.005


type Direction = Literal["row", "column"]
type AspectRatio = tuple[int, int]
type AspectRatios = tuple[AspectRatio, ...]
type Signature = tuple[Any, ...]


@dataclass(frozen=True)
class ImageNode:
    index: int
    width_factor: float
    block: Any = None
    gap_factor: float = 0.0

    @property
    def ratio(self) -> float:
        return self.width_factor

    @property
    def count(self) -> int:
        return 1

    @property
    def depth(self) -> int:
        return 0

    @property
    def signature(self) -> Signature:
        return ("image", self.index)

    def materialize(self, blocks: Sequence[Any]) -> ImageNode:
        return ImageNode(
            index=self.index,
            width_factor=self.width_factor,
            block=blocks[self.index],
        )

    def to_data(self) -> dict[str, Any]:
        return {
            "type": "image",
            "index": self.index,
            "width_factor": self.width_factor,
        }


@dataclass(frozen=True)
class GroupNode:
    direction: Direction
    first: Node
    second: Node
    width_factor: float
    gap_factor: float

    @property
    def ratio(self) -> float:
        return self.width_factor

    @cached_property
    def count(self) -> int:
        return self.first.count + self.second.count

    @cached_property
    def depth(self) -> int:
        return 1 + max(
            self.first.depth,
            self.second.depth,
        )

    @cached_property
    def signature(self) -> Signature:
        return (
            self.direction,
            self.first.signature,
            self.second.signature,
        )

    @property
    def first_width_share(self) -> float:
        return self.first.width_factor / self.width_factor

    @property
    def second_width_share(self) -> float:
        return self.second.width_factor / self.width_factor

    @property
    def first_gap_coefficient(self) -> float:
        return self.first.gap_factor - self.first_width_share * self.gap_factor

    @property
    def second_gap_coefficient(self) -> float:
        return self.second.gap_factor - self.second_width_share * self.gap_factor

    def materialize(self, blocks: Sequence[Any]) -> GroupNode:
        return GroupNode(
            direction=self.direction,
            first=self.first.materialize(blocks),
            second=self.second.materialize(blocks),
            width_factor=self.width_factor,
            gap_factor=self.gap_factor,
        )

    def to_data(self) -> dict[str, Any]:
        return {
            "type": "group",
            "direction": self.direction,
            "width_factor": self.width_factor,
            "gap_factor": self.gap_factor,
            "first": self.first.to_data(),
            "second": self.second.to_data(),
        }


type Node = ImageNode | GroupNode


def gallery_layout(
    blocks,
    record,
    *,
    target_ratio: float = 1.5,
    cache_dir: Path | None = None,
) -> Node | None:
    """Create a deterministic rectangular gallery layout."""
    target_ratio = target_ratio if target_ratio > 0 else 1.5

    valid_blocks = []
    aspect_ratios = []

    for block in blocks:
        attachment = record.attachments.get(block["image"])

        if attachment is None or not attachment.width or not attachment.height:
            continue

        divisor = gcd(
            attachment.width,
            attachment.height,
        )

        aspect_ratios.append(
            (
                attachment.width // divisor,
                attachment.height // divisor,
            )
        )
        valid_blocks.append(block)

    if not aspect_ratios:
        return None

    ratios = tuple(aspect_ratios)

    if len(ratios) == 1:
        width, height = ratios[0]

        return ImageNode(
            index=0,
            width_factor=width / height,
            block=valid_blocks[0],
        )

    normalized_target = round(
        target_ratio,
        RATIO_PRECISION,
    )

    if cache_dir is None:
        layout = select_cached_layout(
            LAYOUT_ALGORITHM_VERSION,
            ratios,
            normalized_target,
        )
    else:
        layout = LayoutCache(cache_dir).get(
            version=LAYOUT_ALGORITHM_VERSION,
            aspect_ratios=ratios,
            target_ratio=normalized_target,
        )

    return layout.materialize(valid_blocks)


@lru_cache(maxsize=128)
def build_cached_candidates(
    version: int,
    aspect_ratios: AspectRatios,
) -> tuple[Node, ...]:
    """Return candidate layouts for the given sequence of aspect ratios."""
    images = [
        ImageNode(
            index=index,
            width_factor=width / height,
        )
        for index, (width, height) in enumerate(aspect_ratios)
    ]

    return tuple(build_candidates(images))


@lru_cache(maxsize=256)
def select_cached_layout(
    version: int,
    aspect_ratios: AspectRatios,
    target_ratio: float,
) -> Node:
    """Select the best layout and cache it for the current process."""
    candidates = build_cached_candidates(
        version,
        aspect_ratios,
    )

    return min(
        candidates,
        key=lambda candidate: (
            final_score(
                candidate,
                target_ratio=target_ratio,
            ),
            candidate.signature,
        ),
    )


@dataclass(frozen=True)
class LayoutCache:
    directory: Path

    def get(
        self,
        *,
        version: int,
        aspect_ratios: AspectRatios,
        target_ratio: float,
    ) -> Node:
        """Return a cached layout, calculating it when necessary."""
        path = self.path_for(
            version=version,
            aspect_ratios=aspect_ratios,
            target_ratio=target_ratio,
        )

        if layout := self.read(path):
            return layout

        layout = select_cached_layout(
            version,
            aspect_ratios,
            target_ratio,
        )

        self.write(path, layout)

        return layout

    def path_for(
        self,
        *,
        version: int,
        aspect_ratios: AspectRatios,
        target_ratio: float,
    ) -> Path:
        payload = {
            "version": version,
            "aspect_ratios": aspect_ratios,
            "target_ratio": target_ratio,
        }

        encoded = json.dumps(
            payload,
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")

        digest = hashlib.sha256(encoded).hexdigest()

        return self.directory / f"{digest}.json"

    @staticmethod
    def read(path: Path) -> Node | None:
        try:
            with path.open(
                "r",
                encoding="utf-8",
            ) as file:
                return node_from_data(json.load(file))
        except (
            OSError,
            ValueError,
            KeyError,
            TypeError,
        ):
            return None

    @staticmethod
    def write(
        path: Path,
        node: Node,
    ) -> None:
        temporary_path: Path | None = None

        try:
            path.parent.mkdir(
                parents=True,
                exist_ok=True,
            )

            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                dir=path.parent,
                delete=False,
            ) as file:
                json.dump(
                    node.to_data(),
                    file,
                    separators=(",", ":"),
                )
                temporary_path = Path(file.name)

            os.replace(
                temporary_path,
                path,
            )

        except OSError:
            if temporary_path is not None:
                temporary_path.unlink(
                    missing_ok=True,
                )


def node_from_data(data: dict[str, Any]) -> Node:
    """Create a layout node from its JSON representation."""
    match data["type"]:
        case "image":
            return ImageNode(
                index=int(data["index"]),
                width_factor=float(data["width_factor"]),
            )

        case "group":
            direction = data["direction"]

            if direction not in ("row", "column"):
                raise ValueError(f"Invalid gallery direction: {direction!r}")

            return GroupNode(
                direction=direction,
                first=node_from_data(data["first"]),
                second=node_from_data(data["second"]),
                width_factor=float(data["width_factor"]),
                gap_factor=float(data["gap_factor"]),
            )

        case node_type:
            raise ValueError(f"Unknown gallery node type: {node_type!r}")


def build_candidates(
    images: Sequence[ImageNode],
) -> list[Node]:
    """Build slicing-layout candidates using dynamic programming."""
    count = len(images)
    all_indices = range(count)

    candidates: dict[int, list[Node]] = {
        1 << index: [image] for index, image in enumerate(images)
    }

    for subset_size in range(2, count + 1):
        for indices in combinations(
            all_indices,
            subset_size,
        ):
            mask = mask_for(indices)
            first_bit = mask & -mask

            generated: dict[float, Node] = {}

            partition = (mask - 1) & mask

            while partition:
                if partition & first_bit:
                    other = mask ^ partition

                    if other:
                        combine_partitions(
                            generated,
                            candidates[partition],
                            candidates[other],
                        )

                partition = (partition - 1) & mask

            candidates[mask] = prune_candidates(
                generated.values(),
                limit=MAX_CANDIDATES_PER_SUBSET,
            )

    return candidates[(1 << count) - 1]


def mask_for(indices) -> int:
    """Return the bit mask representing a sequence of image indices."""
    return sum(1 << index for index in indices)


def combine_partitions(
    generated: dict[float, Node],
    first_candidates: Sequence[Node],
    second_candidates: Sequence[Node],
) -> None:
    """Combine every candidate from two partitions."""
    for first in first_candidates:
        for second in second_candidates:
            retain_candidate(
                generated,
                combine_row(first, second),
            )
            retain_candidate(
                generated,
                combine_column(first, second),
            )


def combine_row(
    first: Node,
    second: Node,
) -> GroupNode:
    """Place two rectangles beside each other."""
    return GroupNode(
        direction="row",
        first=first,
        second=second,
        width_factor=(first.width_factor + second.width_factor),
        gap_factor=(first.gap_factor + second.gap_factor + 1),
    )


def combine_column(
    first: Node,
    second: Node,
) -> GroupNode:
    """Stack two rectangles vertically."""

    first_width = first.width_factor
    second_width = second.width_factor
    combined_width = first_width + second_width

    width_factor = first_width * second_width / combined_width

    gap_factor = (
        second_width * first.gap_factor
        + first_width * second.gap_factor
        - first_width * second_width
    ) / combined_width

    return GroupNode(
        direction="column",
        first=first,
        second=second,
        width_factor=width_factor,
        gap_factor=gap_factor,
    )


def retain_candidate(
    candidates: dict[float, Node],
    candidate: Node,
) -> None:
    """Retain the best candidate for an approximate outer ratio."""
    key = round(
        candidate.ratio,
        RATIO_PRECISION,
    )

    existing = candidates.get(key)

    if existing is None or candidate_preference(candidate) < candidate_preference(
        existing
    ):
        candidates[key] = candidate


def candidate_preference(
    node: Node,
) -> tuple[float, Signature]:
    """Return the ordering used for equivalent-ratio candidates."""
    return (
        structure_score(node),
        node.signature,
    )


def prune_candidates(
    candidates,
    *,
    limit: int,
) -> list[Node]:
    """Keep a deterministic range of candidate aspect ratios."""
    ordered = sorted(
        candidates,
        key=lambda candidate: (
            log(candidate.ratio),
            *candidate_preference(candidate),
        ),
    )

    if len(ordered) <= limit:
        return ordered

    last_index = len(ordered) - 1
    step = last_index / (limit - 1)

    return [ordered[round(index * step)] for index in range(limit)]


@lru_cache(maxsize=None)
def t_junction_penalty(node: Node) -> float:
    """Score seams that terminate against neighbouring groups."""
    if isinstance(node, ImageNode):
        return 0.0

    opposite = "column" if node.direction == "row" else "row"

    penalty = t_junction_penalty(node.first) + t_junction_penalty(node.second)

    for child in (node.first, node.second):
        if isinstance(child, GroupNode) and child.direction == opposite:
            penalty += child.count - 1

    return penalty


@lru_cache(maxsize=None)
def sibling_imbalance_penalty(node: Node) -> float:
    """Score differences in sibling subtree sizes."""
    if isinstance(node, ImageNode):
        return 0.0

    return (
        abs(node.first.count - node.second.count)
        + sibling_imbalance_penalty(node.first)
        + sibling_imbalance_penalty(node.second)
    )


@lru_cache(maxsize=None)
def orientation_switch_penalty(node: Node) -> float:
    """Score changes in split direction within the tree."""
    if isinstance(node, ImageNode):
        return 0.0

    penalty = orientation_switch_penalty(node.first) + orientation_switch_penalty(
        node.second
    )

    for child in (node.first, node.second):
        if isinstance(child, GroupNode) and child.direction != node.direction:
            penalty += 1

    return penalty


@lru_cache(maxsize=None)
def structure_score(node: Node) -> float:
    """Return the structural penalty for a layout."""
    return (
        node.depth * DEPTH_WEIGHT
        + t_junction_penalty(node) * T_JUNCTION_WEIGHT
        + sibling_imbalance_penalty(node) * SIBLING_IMBALANCE_WEIGHT
        + orientation_switch_penalty(node) * ORIENTATION_SWITCH_WEIGHT
    )


def final_score(
    node: Node,
    *,
    target_ratio: float,
) -> float:
    """Score a complete layout against the requested gallery ratio."""
    return abs(log(node.ratio / target_ratio)) + structure_score(node)


class GalleryPlugin(Plugin):
    name = "Gallery"
    description = "Static aspect-ratio-aware image gallery layouts."

    def on_setup_env(
        self,
        **extra: Any,
    ) -> None:
        cache_dir = Path(self.env.root_path) / ".cache" / "lektor-gallery"

        self.env.jinja_env.filters["gallery_layout"] = partial(
            gallery_layout,
            cache_dir=cache_dir,
        )