from __future__ import annotations

import json
import os
from pathlib import Path
from threading import RLock
from typing import 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()

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

    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)
        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)