jigaido/apps/telegram-bot/storage.py

"""Per-group JSON file storage for JIGAIDO."""

import json
import os
import tempfile
from pathlib import Path
from typing import Optional

DATA_DIR = Path.home() / ".jigaido"


def _ensure_dirs() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)


def _group_dir(group_id: int) -> Path:
    return DATA_DIR / str(group_id)


def _user_personal_dir(user_id: int) -> Path:
    return DATA_DIR / str(user_id)


def _group_file(group_id: int) -> Path:
    return _group_dir(group_id) / "group.json"


def _user_tracking_file(group_id: int, user_id: int) -> Path:
    return _group_dir(group_id) / f"{user_id}.json"


def _user_personal_file(user_id: int) -> Path:
    return _user_personal_dir(user_id) / "user.json"


def _atomic_write(path: Path, data: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as tmp:
        json.dump(data, tmp, indent=2)
        tmp_path = tmp.name
    os.rename(tmp_path, path)


def load_group_bounties(group_id: int) -> dict:
    _ensure_dirs()
    path = _group_file(group_id)
    if not path.exists():
        return {"bounties": [], "next_id": 1}
    with open(path) as f:
        return json.load(f)


def save_group_bounties(group_id: int, data: dict) -> None:
    _atomic_write(_group_file(group_id), data)


def add_group_bounty(
    group_id: int,
    created_by_user_id: int,
    text: Optional[str],
    link: Optional[str],
    due_date_ts: Optional[int],
) -> dict:
    data = load_group_bounties(group_id)
    bounty = {
        "id": data["next_id"],
        "created_by_user_id": created_by_user_id,
        "text": text,
        "link": link,
        "due_date_ts": due_date_ts,
        "created_at": int(os.path.getctime(_group_file(group_id)))
        if _group_file(group_id).exists()
        else 0,
    }
    data["bounties"].append(bounty)
    data["next_id"] += 1
    save_group_bounties(group_id, data)
    return bounty


def update_group_bounty(
    group_id: int,
    bounty_id: int,
    text: Optional[str],
    link: Optional[str],
    due_date_ts: Optional[int],
) -> bool:
    data = load_group_bounties(group_id)
    for bounty in data["bounties"]:
        if bounty["id"] == bounty_id:
            if text is not None:
                bounty["text"] = text
            if link is not None:
                bounty["link"] = link
            if due_date_ts is not None:
                bounty["due_date_ts"] = due_date_ts
            save_group_bounties(group_id, data)
            return True
    return False


def delete_group_bounty(group_id: int, bounty_id: int) -> bool:
    data = load_group_bounties(group_id)
    for i, bounty in enumerate(data["bounties"]):
        if bounty["id"] == bounty_id:
            data["bounties"].pop(i)
            save_group_bounties(group_id, data)
            return True
    return False


def get_group_bounty(group_id: int, bounty_id: int) -> Optional[dict]:
    data = load_group_bounties(group_id)
    for bounty in data["bounties"]:
        if bounty["id"] == bounty_id:
            return bounty
    return None


def load_user_tracking(group_id: int, user_id: int) -> dict:
    path = _user_tracking_file(group_id, user_id)
    if not path.exists():
        return {"tracked": []}
    with open(path) as f:
        return json.load(f)


def save_user_tracking(group_id: int, user_id: int, data: dict) -> None:
    _atomic_write(_user_tracking_file(group_id, user_id), data)


def track_bounty(group_id: int, user_id: int, bounty_id: int) -> bool:
    data = load_user_tracking(group_id, user_id)
    if any(t["bounty_id"] == bounty_id for t in data["tracked"]):
        return False
    data["tracked"].append({"bounty_id": bounty_id})
    save_user_tracking(group_id, user_id, data)
    return True


def untrack_bounty(group_id: int, user_id: int, bounty_id: int) -> bool:
    data = load_user_tracking(group_id, user_id)
    for i, t in enumerate(data["tracked"]):
        if t["bounty_id"] == bounty_id:
            data["tracked"].pop(i)
            save_user_tracking(group_id, user_id, data)
            return True
    return False


def load_user_personal(user_id: int) -> dict:
    path = _user_personal_file(user_id)
    if not path.exists():
        return {"bounties": [], "next_id": 1}
    with open(path) as f:
        return json.load(f)


def save_user_personal(user_id: int, data: dict) -> None:
    _atomic_write(_user_personal_file(user_id), data)


def add_personal_bounty(
    user_id: int, text: Optional[str], link: Optional[str], due_date_ts: Optional[int]
) -> dict:
    data = load_user_personal(user_id)
    bounty = {
        "id": data["next_id"],
        "text": text,
        "link": link,
        "due_date_ts": due_date_ts,
        "created_at": 0,
    }
    data["bounties"].append(bounty)
    data["next_id"] += 1
    save_user_personal(user_id, data)
    return bounty


def update_personal_bounty(
    user_id: int,
    bounty_id: int,
    text: Optional[str],
    link: Optional[str],
    due_date_ts: Optional[int],
) -> bool:
    data = load_user_personal(user_id)
    for bounty in data["bounties"]:
        if bounty["id"] == bounty_id:
            if text is not None:
                bounty["text"] = text
            if link is not None:
                bounty["link"] = link
            if due_date_ts is not None:
                bounty["due_date_ts"] = due_date_ts
            save_user_personal(user_id, data)
            return True
    return False


def delete_personal_bounty(user_id: int, bounty_id: int) -> bool:
    data = load_user_personal(user_id)
    for i, bounty in enumerate(data["bounties"]):
        if bounty["id"] == bounty_id:
            data["bounties"].pop(i)
            save_user_personal(user_id, data)
            return True
    return False


def get_personal_bounty(user_id: int, bounty_id: int) -> Optional[dict]:
    data = load_user_personal(user_id)
    for bounty in data["bounties"]:
        if bounty["id"] == bounty_id:
            return bounty
    return None