TgMessageHook/app/storage.py

import json
import re
import secrets
import string
import threading
from datetime import UTC, datetime
from pathlib import Path
from typing import List, Optional, Set

from fastapi.concurrency import run_in_threadpool

from .config import get_settings
from .models import HookCreate, HookRead


HOOK_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]{3,64}$")


class HookStore:
    def __init__(self, storage_path: Path, hook_id_length: int) -> None:
        self.storage_path = storage_path
        self.hook_id_length = hook_id_length
        self._lock = threading.Lock()
        self._initialize()

    def _initialize(self) -> None:
        if not self.storage_path.exists():
            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
            self.storage_path.write_text("[]\n", encoding="utf-8")
        else:
            try:
                self._load_raw()
            except (json.JSONDecodeError, UnicodeDecodeError):
                # Corrupted or non-UTF file; back it up and start fresh
                backup = self.storage_path.with_suffix(self.storage_path.suffix + ".bak")
                self.storage_path.replace(backup)
                self.storage_path.write_text("[]\n", encoding="utf-8")

    def _load_raw(self) -> List[dict]:
        text = self.storage_path.read_text(encoding="utf-8")
        data = json.loads(text or "[]")
        if not isinstance(data, list):
            raise json.JSONDecodeError("Hook store must contain a list", text, 0)
        return data

    def _save_raw(self, data: List[dict]) -> None:
        payload = json.dumps(data, indent=2, ensure_ascii=False)
        tmp_path = self.storage_path.with_suffix(self.storage_path.suffix + ".tmp")
        tmp_path.write_text(payload + "\n", encoding="utf-8")
        tmp_path.replace(self.storage_path)

    def _generate_hook_id(self, existing_ids: Set[str]) -> str:
        alphabet = string.ascii_lowercase + string.digits
        while True:
            candidate = "".join(secrets.choice(alphabet) for _ in range(self.hook_id_length))
            if candidate not in existing_ids:
                return candidate

    def list_hooks(self) -> List[HookRead]:
        raw_hooks = self._load_raw()
        hooks = [
            HookRead(
                hook_id=item["hook_id"],
                chat_id=item["chat_id"],
                message=item["message"],
                created_at=datetime.fromisoformat(item["created_at"]),
            )
            for item in raw_hooks
        ]
        hooks.sort(key=lambda h: h.created_at, reverse=True)
        return hooks

    def create_hook(self, payload: HookCreate) -> HookRead:
        created_at = datetime.now(UTC).replace(microsecond=0).isoformat()
        with self._lock:
            raw_hooks = self._load_raw()
            existing_ids = {item["hook_id"] for item in raw_hooks}
            hook_id = self._generate_hook_id(existing_ids)
            raw_hooks.append(
                {
                    "hook_id": hook_id,
                    "chat_id": payload.chat_id,
                    "message": payload.message,
                    "created_at": created_at,
                }
            )
            self._save_raw(raw_hooks)
        return HookRead(
            hook_id=hook_id,
            chat_id=payload.chat_id,
            message=payload.message,
            created_at=datetime.fromisoformat(created_at),
        )

    def get_hook(self, hook_id: str) -> Optional[HookRead]:
        raw_hooks = self._load_raw()
        for item in raw_hooks:
            if item.get("hook_id") == hook_id:
                return HookRead(
                    hook_id=item["hook_id"],
                    chat_id=item["chat_id"],
                    message=item["message"],
                    created_at=datetime.fromisoformat(item["created_at"]),
                )
        return None

    def delete_hook(self, hook_id: str) -> bool:
        with self._lock:
            raw_hooks = self._load_raw()
            new_hooks = [item for item in raw_hooks if item.get("hook_id") != hook_id]
            if len(new_hooks) == len(raw_hooks):
                return False
            self._save_raw(new_hooks)
            return True

    def update_hook_id(self, current_id: str, new_id: str) -> HookRead:
        normalized_new_id = new_id.strip()
        if not normalized_new_id:
            raise ValueError("Hook ID cannot be empty")
        if not HOOK_ID_PATTERN.fullmatch(normalized_new_id):
            raise ValueError("Hook ID must be 3-64 characters of letters, numbers, '_' or '-' only")
        with self._lock:
            raw_hooks = self._load_raw()
            exists = next((item for item in raw_hooks if item.get("hook_id") == current_id), None)
            if not exists:
                raise KeyError("Hook not found")
            if normalized_new_id == current_id:
                return HookRead(
                    hook_id=exists["hook_id"],
                    chat_id=exists["chat_id"],
                    message=exists["message"],
                    created_at=datetime.fromisoformat(exists["created_at"]),
                )
            if any(item.get("hook_id") == normalized_new_id for item in raw_hooks):
                raise ValueError("Hook ID already exists")
            exists["hook_id"] = normalized_new_id
            self._save_raw(raw_hooks)
            return HookRead(
                hook_id=normalized_new_id,
                chat_id=exists["chat_id"],
                message=exists["message"],
                created_at=datetime.fromisoformat(exists["created_at"]),
            )


settings = get_settings()
store = HookStore(settings.database_path, settings.hook_id_length)


async def list_hooks_async() -> List[HookRead]:
    return await run_in_threadpool(store.list_hooks)


async def create_hook_async(payload: HookCreate) -> HookRead:
    return await run_in_threadpool(store.create_hook, payload)


async def get_hook_async(hook_id: str) -> Optional[HookRead]:
    return await run_in_threadpool(store.get_hook, hook_id)


async def delete_hook_async(hook_id: str) -> bool:
    return await run_in_threadpool(store.delete_hook, hook_id)


async def update_hook_id_async(current_hook_id: str, new_hook_id: str) -> HookRead:
    return await run_in_threadpool(store.update_hook_id, current_hook_id, new_hook_id)