from __future__ import annotations

import copy
import json
import os
import time
from dataclasses import dataclass, field
from typing import Any

from maxim.utils.atomic_io import atomic_write_json


@dataclass
class RuntimeState:
    max_steps: int = 100
    confirmed: bool = False
    steps_taken: int = 0
    done: bool = False
    last_error: str | None = None
    data: dict[str, Any] = field(default_factory=dict)

    def update(self, observation: dict[str, Any] | None) -> None:
        if not observation:
            return
        if isinstance(observation, dict):
            self.data.update(observation)

    def snapshot(self) -> dict[str, Any]:
        return {
            "steps_taken": int(self.steps_taken),
            "max_steps": int(self.max_steps),
            "confirmed": bool(self.confirmed),
            "done": bool(self.done),
            "last_error": self.last_error,
            "data": copy.deepcopy(self.data),  # Deep copy to prevent shared mutable state
        }

    def save_json(self, path: str, *, meta: dict[str, Any] | None = None) -> None:
        from maxim.utils.format_version import with_format_version

        abs_path = os.path.abspath(path)
        payload: dict[str, Any] = {"saved_at": time.time(), **self.snapshot()}
        if meta:
            payload.update(meta)
        atomic_write_json(abs_path, with_format_version(payload))

    @classmethod
    def load_json(cls, path: str) -> "RuntimeState":
        import logging

        from maxim.utils.format_version import check_format_version

        with open(path, "r", encoding="utf-8") as fp:
            payload = json.load(fp)
        check_format_version(payload, "runtime_state", log=logging.getLogger(__name__))

        data = payload.get("data") if isinstance(payload, dict) else None
        if not isinstance(data, dict):
            data = {}

        raw_max_steps = payload.get("max_steps", 100)
        try:
            max_steps = int(raw_max_steps)
        except Exception:
            max_steps = 100

        return cls(
            max_steps=max_steps,
            confirmed=bool(payload.get("confirmed", False)),
            steps_taken=int(payload.get("steps_taken", 0) or 0),
            done=bool(payload.get("done", False)),
            last_error=payload.get("last_error"),
            data=data,
        )

    def mark_failure(self, error: str | None) -> None:
        self.last_error = str(error) if error is not None else None

    def is_done(self) -> bool:
        return self.done
