from __future__ import annotations

import json
import os
from pathlib import Path
from threading import RLock
from typing import cast, TypeVar

from pydantic import BaseModel

from app.core.models import BackupBundle, ControlState, LearningState, RuntimeState

T = TypeVar("T", bound=BaseModel)


class JsonDocumentStore:
    def __init__(self, root: str | Path) -> None:
        self.root = Path(root).resolve()
        self.root.mkdir(parents=True, exist_ok=True)
        self._lock = RLock()
        self._cache: dict[str, BaseModel] = {}

    def load(self, name: str, model: type[T], default: T) -> T:
        path = self.root / name
        with self._lock:
            cached = self._cache.get(name)
            if cached is not None:
                return cast(T, cached)
            if not path.exists():
                self._cache[name] = default
                return default
            value = model.model_validate_json(path.read_text(encoding="utf-8"))
            self._cache[name] = value
            return value

    def save(self, name: str, value: T) -> T:
        path = self.root / name
        temporary = path.with_suffix(path.suffix + ".tmp")
        with self._lock:
            temporary.write_text(
                json.dumps(value.model_dump(mode="json"), ensure_ascii=True, sort_keys=True) + "\n",
                encoding="utf-8",
            )
            os.replace(temporary, path)
            self._cache[name] = value
        return value


class FutureStores:
    def __init__(self, root: str | Path = ".future_store") -> None:
        self._docs = JsonDocumentStore(root)

    def runtime(self) -> RuntimeState:
        return self._docs.load("runtime.json", RuntimeState, RuntimeState())

    def save_runtime(self, state: RuntimeState) -> RuntimeState:
        return self._docs.save("runtime.json", state)

    def learning(self) -> LearningState:
        return self._docs.load("learning.json", LearningState, LearningState())

    def save_learning(self, state: LearningState) -> LearningState:
        return self._docs.save("learning.json", state)

    def control(self) -> ControlState:
        return self._docs.load("control.json", ControlState, ControlState())

    def save_control(self, state: ControlState) -> ControlState:
        return self._docs.save("control.json", state)

    def export_backup(self) -> BackupBundle:
        return BackupBundle(
            runtime=self.runtime(),
            learning=self.learning(),
            control=self.control(),
        )

    def restore_backup(self, bundle: BackupBundle) -> None:
        self.save_runtime(bundle.runtime)
        self.save_learning(bundle.learning)
        self.save_control(bundle.control)