"""Data models for POMA SDK results."""

from __future__ import annotations

import warnings
from dataclasses import dataclass, field, asdict
from typing import Any

from poma.utils import _resolve_images


@dataclass
class PomaChunk:
    """A single chunk of document content.

    Attributes:
        tag: Deprecated. Use ``file_id`` instead. Will be removed in a
            future version.
    """

    chunk_index: int
    content: str
    depth: int
    file_id: str | None = None
    tag: str | None = None
    #: Identifier for the code block this chunk belongs to, or ``None``
    #: for non-code chunks. Chunks sharing the same ``code`` value are
    #: part of the same code block; their depth is rebased relative to
    #: the block's minimum depth for correct indentation in cheatsheets.
    code: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PomaChunk:
        file_id = data.get("file_id")
        tag = data.get("tag")
        if not file_id and not tag:
            warnings.warn(
                "PomaChunk: neither 'file_id' nor 'tag' is set. "
                "This may cause issues with cheatsheet generation.",
                UserWarning,
                stacklevel=2,
            )
        if tag and not file_id:
            warnings.warn(
                "PomaChunk: 'tag' is deprecated, use 'file_id' instead.",
                DeprecationWarning,
                stacklevel=2,
            )
        return cls(
            chunk_index=data["chunk_index"],
            content=data["content"],
            file_id=file_id or tag,
            tag=tag,
            depth=data["depth"],
            code=data.get("code"),
        )


@dataclass
class PomaChunkSet:
    """A set of related chunks with embedding text.

    Attributes:
        tag: Deprecated. Use ``file_id`` instead. Will be removed in a
            future version.
    """

    chunkset_index: int
    chunks: list[int]
    to_embed: str
    file_id: str | None = None
    tag: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PomaChunkSet:
        file_id = data.get("file_id")
        tag = data.get("tag")
        if not file_id and not tag:
            warnings.warn(
                "PomaChunkSet: neither 'file_id' nor 'tag' is set. "
                "This may cause issues with cheatsheet generation.",
                UserWarning,
                stacklevel=2,
            )
        if tag and not file_id:
            warnings.warn(
                "PomaChunkSet: 'tag' is deprecated, use 'file_id' instead.",
                DeprecationWarning,
                stacklevel=2,
            )
        return cls(
            chunkset_index=data["chunkset_index"],
            chunks=data["chunks"],
            to_embed=data.get("to_embed", data.get("contents", "")),
            file_id=file_id or tag,
            tag=tag,
        )


@dataclass
class PomaResult:
    """Result of a POMA ingestion containing chunks, chunksets, and assets."""

    chunks: list[PomaChunk] = field(default_factory=list)
    chunksets: list[PomaChunkSet] = field(default_factory=list)
    #: Images extracted from the ingested document, keyed by image ID
    #: (e.g. ``"image_00001"``). Values are base64 data URI strings
    #: (e.g. ``"data:image/jpeg;base64,..."``).
    #: Empty dict when the document contains no images.
    images: dict[str, str] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "chunks": [c.to_dict() for c in self.chunks],
            "chunksets": [cs.to_dict() for cs in self.chunksets],
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PomaResult:
        images = _resolve_images(
            data.get("image_sources", {}),
            data.get("asset_files", {}),
        )
        return cls(
            chunks=[PomaChunk.from_dict(c) for c in data.get("chunks", [])],
            chunksets=[PomaChunkSet.from_dict(cs) for cs in data.get("chunksets", [])],
            images=images,
        )
