from __future__ import annotations

from uuid import uuid4

from app.modules.agent.engine.orchestrator.models import ArtifactItem, ArtifactType


class ArtifactStore:
    def __init__(self) -> None:
        self._by_id: dict[str, ArtifactItem] = {}
        self._by_key: dict[str, ArtifactItem] = {}

    def put(self, *, key: str, artifact_type: ArtifactType, content=None, meta: dict | None = None) -> ArtifactItem:
        item_meta = dict(meta or {})
        if content is not None and not isinstance(content, str):
            item_meta.setdefault("value", content)
        item = ArtifactItem(
            artifact_id=f"artifact_{uuid4().hex}",
            key=key,
            type=artifact_type,
            content=self._as_content(content),
            meta=item_meta,
        )
        self._by_id[item.artifact_id] = item
        self._by_key[key] = item
        return item

    def get(self, key: str) -> ArtifactItem | None:
        return self._by_key.get(key)

    def get_content(self, key: str, default=None):
        item = self.get(key)
        if item is None:
            return default
        if item.content is not None:
            return item.content
        return item.meta.get("value", default)

    def has(self, key: str) -> bool:
        return key in self._by_key

    def all_items(self) -> list[ArtifactItem]:
        return list(self._by_id.values())

    def _as_content(self, value):
        if value is None:
            return None
        if isinstance(value, str):
            return value
        return None
