"""Slot-scoped opaque strategy-instance state."""
from __future__ import annotations

import os
import threading
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Optional, Protocol
from uuid import UUID

_STATE_CAP_BYTES = int(os.environ.get("TREX_STRATEGY_STATE_CAP_BYTES", "4096"))

REASON_VERSION_CONFLICT = "version_conflict"
REASON_SLOT_INACTIVE = "slot_inactive"
REASON_STATE_TOO_LARGE = "state_too_large"


@dataclass(frozen=True)
class StrategyStateEntry:
    """One opaque entry in a slot's strategy_instance_id map."""

    version: int
    data: str
    updated_at: datetime


@dataclass(frozen=True)
class SetResult:
    ok: bool
    version: Optional[int] = None
    reason: Optional[str] = None


def _byte_len(data: str) -> int:
    return len(data.encode("utf-8"))


def _cas(
    existing: Optional[StrategyStateEntry], *, expected_version: int, data: str, slot_active: bool
) -> SetResult:
    if not slot_active:
        return SetResult(ok=False, reason=REASON_SLOT_INACTIVE)
    if _byte_len(data) > _STATE_CAP_BYTES:
        return SetResult(ok=False, reason=REASON_STATE_TOO_LARGE)
    stored_version = existing.version if existing is not None else 0
    if expected_version != stored_version:
        return SetResult(ok=False, reason=REASON_VERSION_CONFLICT, version=stored_version)
    return SetResult(ok=True, version=stored_version + 1)


class StrategyStateStore(Protocol):
    def set(
        self,
        slot_id: UUID,
        strategy_instance_id: str,
        *,
        expected_version: int,
        data: str,
        slot_active: bool,
    ) -> SetResult: ...

    def get_for_slots(self, slot_ids: list[UUID]) -> dict[UUID, dict[str, StrategyStateEntry]]: ...


class InMemoryStrategyStateStore:
    """Process-local slot-scoped state-map store."""

    def __init__(self) -> None:
        self._store: dict[UUID, dict[str, StrategyStateEntry]] = {}
        self._lock = threading.Lock()

    def set(
        self,
        slot_id: UUID,
        strategy_instance_id: str,
        *,
        expected_version: int,
        data: str,
        slot_active: bool,
    ) -> SetResult:
        with self._lock:
            entries = self._store.setdefault(slot_id, {})
            existing = entries.get(strategy_instance_id)
            result = _cas(
                existing, expected_version=expected_version, data=data, slot_active=slot_active
            )
            if not result.ok:
                if not entries:
                    self._store.pop(slot_id, None)
                return result
            entries[strategy_instance_id] = StrategyStateEntry(
                version=result.version,  # type: ignore[arg-type]
                data=data,
                updated_at=datetime.now(tz=timezone.utc),
            )
            return result

    def get_for_slots(self, slot_ids: list[UUID]) -> dict[UUID, dict[str, StrategyStateEntry]]:
        wanted = set(slot_ids)
        with self._lock:
            return {sid: dict(entries) for sid, entries in self._store.items() if sid in wanted}