assistant/storage.py

import os
import json

# Base paths relative to the backend root
BASE_DIR = os.path.join(os.path.dirname(__file__), '..', 'data')
DOCUMENTS_DIR = os.path.join(BASE_DIR, 'documents')
METADATA_FILE = os.path.join(BASE_DIR, 'metadata.json')

# Ensure directories exist on import
os.makedirs(DOCUMENTS_DIR, exist_ok=True)


def _load_metadata() -> dict:
    """Load metadata from disk. Returns empty dict if file doesn't exist."""
    if not os.path.exists(METADATA_FILE):
        return {}
    with open(METADATA_FILE, 'r', encoding='utf-8') as f:
        return json.load(f)


def _write_metadata(metadata: dict):
    """Persist metadata dict to disk."""
    with open(METADATA_FILE, 'w', encoding='utf-8') as f:
        json.dump(metadata, f, indent=2)


def generate_doc_id() -> str:
    """Auto-increment doc IDs: doc1, doc2, doc3 ..."""
    metadata = _load_metadata()
    next_index = len(metadata) + 1
    return f"doc{next_index}"


def save_document(doc_id: str, text: str):
    """Save extracted document text to data/documents/{doc_id}.txt"""
    filepath = os.path.join(DOCUMENTS_DIR, f"{doc_id}.txt")
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(text)


def load_document(doc_id: str) -> str | None:
    """Load document text from disk. Returns None if not found."""
    filepath = os.path.join(DOCUMENTS_DIR, f"{doc_id}.txt")
    if not os.path.exists(filepath):
        return None
    with open(filepath, 'r', encoding='utf-8') as f:
        return f.read()


def save_metadata(doc_id: str, filename: str):
    """Append a new document entry to metadata.json."""
    metadata = _load_metadata()
    metadata[doc_id] = {"filename": filename}
    _write_metadata(metadata)


def get_all_docs() -> list:
    """Return list of all uploaded document IDs and their filenames."""
    metadata = _load_metadata()
    return [{"doc_id": doc_id, "filename": info.get("filename", "")}
            for doc_id, info in metadata.items()]