from __future__ import annotations

import asyncio
from typing import Any
import json

import aiofiles

class SharedState:
    def __init__(self, initial: dict[str, Any] | None = None) -> None:
        self._state: dict[str, Any] = dict(initial or {})
        self._lock = asyncio.Lock()

    async def read(self, keys: list[str] | None = None) -> dict[str, Any]:
        async with self._lock:
            snapshot = dict(self._state)
        if keys is None:
            return snapshot
        return {k: snapshot.get(k) for k in keys}

    async def patch(self, patch: dict[str, Any]) -> dict[str, Any]:
        def _merge(existing: Any, incoming: Any) -> Any:
            # list + list => append new items to keep history instead of overwriting
            if isinstance(existing, list) and isinstance(incoming, list):
                return list(existing) + list(incoming)
            # dict + dict => shallow recursive merge to preserve missing keys
            if isinstance(existing, dict) and isinstance(incoming, dict):
                merged = dict(existing)
                for k, v in incoming.items():
                    if k in merged:
                        merged[k] = _merge(merged[k], v)
                    else:
                        merged[k] = v
                return merged
            # fallback: replace (for scalars or mismatched types)
            return incoming

        async with self._lock:
            for k, v in patch.items():
                if k in self._state:
                    self._state[k] = _merge(self._state[k], v)
                else:
                    self._state[k] = v
            return dict(self._state)

    def get_nowait(self, key: str, default: Any = None) -> Any:
        # Best-effort read without locking; useful for formatting small templates.
        return self._state.get(key, default)

    async def save_state_locally(self, path: str) -> dict[str, Any]:
        """Append current state as one JSONL line to the given path."""
        dumpable_state = {k: v for k, v in self._state.items() if v is not None}
        line = json.dumps(dumpable_state, ensure_ascii=False) + "\n"
        async with aiofiles.open(path, "a", encoding="utf-8") as f:
            await f.write(line)
        return {"status": "success", "message": f"State saved to {path}"}
