"""Set a value in shared memory.

Shared memory is accessible by all agents - use it for coordination,
shared state, or passing data between agents.

Example:
    # Agent A stores results
    await memory.shared_set(key="analysis_results", value={"score": 0.85})

    # Agent B reads them (using shared_get)
    result = await memory.shared_get(key="analysis_results")
"""

import json
import os
import tempfile
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from core.config import get_settings
from core.pathways.context import Context
from core.pathways.models.tool import ToolOutput


async def run(
    key: str,
    value: Any,
    *,
    namespace: str = "default",
    context: Context | None = None,
) -> ToolOutput:
    """Set a value in shared memory.

    Args:
        key: Key to store the value under
        value: Value to store (must be JSON serializable)
        namespace: Optional namespace to organize data (default: "default")
        context: Execution context (injected by framework)

    Returns
    -------
        {"success": True, "key": "...", "namespace": "..."}
    """
    if not key:
        return ToolOutput(success=False, error="key is required")

    try:
        state_store = context.services.state_store if context else None

        if not state_store:
            # Fallback to file-based shared memory
            return await _set_file_based(key, value, namespace)

        # Use StateStore
        await state_store.set_shared(key, value, namespace)

        return ToolOutput(
            success=True,
            data={"key": key, "namespace": namespace},
        )

    except TypeError as e:
        return ToolOutput(success=False, error=f"Value not JSON serializable: {e}")
    except Exception as e:
        return ToolOutput(success=False, error=str(e))


async def _set_file_based(key: str, value: Any, namespace: str) -> ToolOutput:
    """Fallback to file-based shared memory when StateStore not available."""
    # Using hoisted imports: json, os, datetime, timezone, Path
    workspace_root = get_settings().workspace_dir
    shared_dir = Path(workspace_root) / ".albus_store" / "shared"
    shared_dir.mkdir(parents=True, exist_ok=True)

    safe_ns = namespace.replace("/", "_").replace("\\", "_")
    path = shared_dir / f"{safe_ns}.json"

    # Load existing memory
    memory = {}
    if path.exists():
        try:
            memory = json.loads(path.read_text())
        except Exception:
            memory = {}

    # Store the value
    memory[key] = value
    memory["_updated_at"] = datetime.now(UTC).isoformat()

    # Atomic write: temp file + os.replace() prevents corruption on crash.
    import contextlib

    data_bytes = json.dumps(memory, indent=2).encode()
    fd, tmp_path = tempfile.mkstemp(dir=path.parent, suffix=".tmp")
    try:
        os.write(fd, data_bytes)
        os.close(fd)
        os.replace(tmp_path, path)
    except BaseException:
        with contextlib.suppress(OSError):
            os.close(fd)
        with contextlib.suppress(OSError):
            os.unlink(tmp_path)
        raise

    return ToolOutput(
        success=True,
        data={"key": key, "namespace": namespace},
    )
