from collections.abc import Iterable
from dataclasses import dataclass
from importlib.resources.abc import Traversable
from itertools import starmap
from typing import Self

from .document import MarkdownDocument


@dataclass(frozen=True, slots=True)
class PromptFilesBundle:
    """Represents a prompt and all of its associated markdown files."""

    slug: str
    prompt: MarkdownDocument
    other: tuple[MarkdownDocument, ...]

    @classmethod
    def from_root(cls, root: Traversable) -> list[Self]:
        # Should only be used for testing
        return list(cls.from_docs(MarkdownDocument.from_root(root).values()))

    @classmethod
    def from_docs(cls, docs: Iterable[MarkdownDocument]) -> tuple[Self, ...]:
        return tuple(starmap(cls.from_paths, group_prompt_paths(docs).items()))

    @classmethod
    def from_paths(cls, slug: str, paths: set[MarkdownDocument]) -> Self:
        prompt = next(iter(p for p in paths if p.relpath.name == "PROMPT.md"))
        sorted_paths = tuple(sorted(paths, key=lambda p: p.relpath.as_posix()))
        other = tuple(p for p in sorted_paths if p != prompt)
        return cls(
            slug=slug,
            prompt=prompt,
            other=other,
        )


def group_prompt_paths(docs: Iterable[MarkdownDocument]) -> dict[str, set[MarkdownDocument]]:
    """Group prompts from a list of markdown documents by their prompt slug."""
    grouped: dict[str, set[MarkdownDocument]] = {}
    for doc in sorted(
        filter(lambda d: d.prompt_slug is not None, docs),
        key=lambda d: d.relpath.as_posix(),
    ):
        if doc.prompt_slug:
            grouped.setdefault(doc.prompt_slug, set()).add(doc)
    return grouped