"""
store.py — a Cosmos DB-shaped conversation store (Knowledge Module 07 §6-7).

This is an in-memory stand-in that models the THREE things that decide success
or failure with Cosmos DB, so you can reason about them without a cloud account:

  1. PARTITION KEY — items are bucketed by a partition key you choose. A POINT
     READ (you supply the partition key + id) is O(1) and ~1 RU. A CROSS-PARTITION
     query (no partition key) fans out across buckets and costs far more RU.
  2. REQUEST UNITS (RU) — every operation costs RUs; we tally them so you can SEE
     why point reads beat queries.
  3. PARTITION KEY CHOICE — we partition conversations by `conversationId`, the
     dominant access pattern, so reading a conversation is a single-partition
     point read and load spreads evenly (no hot partition).

PRODUCTION: replace with azure-cosmos; `read_item(item, partition_key)` and
`query_items(..., partition_key=...)` map 1:1 onto these methods.
"""
from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class CosmosLikeStore:
    # partition_key_value -> {item_id -> document}
    _partitions: dict[str, dict[str, dict]] = field(default_factory=dict)
    ru_consumed: float = 0.0
    partition_key: str = "conversationId"

    def create_item(self, doc: dict) -> dict:
        pk = doc[self.partition_key]
        self._partitions.setdefault(pk, {})[doc["id"]] = doc
        self.ru_consumed += 5.0                       # a write ~5 RU
        return doc

    def point_read(self, item_id: str, partition_key_value: str) -> dict | None:
        """Cheapest op: you provide the partition key, so it's a single-partition
        O(1) lookup (~1 RU). This is why we partition by conversationId."""
        self.ru_consumed += 1.0
        return self._partitions.get(partition_key_value, {}).get(item_id)

    def read_conversation(self, conversation_id: str) -> list[dict]:
        """Single-partition read of all messages in a conversation — cheap because
        they share the partition key."""
        items = list(self._partitions.get(conversation_id, {}).values())
        self.ru_consumed += max(1.0, len(items) * 0.5)
        return sorted(items, key=lambda d: d["seq"])

    def cross_partition_query(self, predicate) -> list[dict]:
        """Expensive: no partition key, so it fans out across ALL partitions. RU
        cost scales with the data scanned — avoid on hot paths."""
        scanned = sum(len(p) for p in self._partitions.values())
        self.ru_consumed += max(2.5, scanned * 2.5)   # much costlier per item
        return [d for p in self._partitions.values() for d in p.values() if predicate(d)]

    def stats(self) -> dict:
        return {"partitions": len(self._partitions),
                "items": sum(len(p) for p in self._partitions.values()),
                "ru_consumed": round(self.ru_consumed, 1)}