src/storage.py

import json
import os
from pathlib import Path

import polars as pl
from huggingface_hub import HfApi, hf_hub_download, upload_file
from huggingface_hub.utils import RepositoryNotFoundError

from src.config import Config
from src.utils import load_json, save_json


class Storage:
    def __init__(self, config: Config, hf_token: str, username: str):
        self.config = config
        self.hf_token = hf_token
        self.username = username
        self.dataset_id = f"{username}/{config.dataset_repo_id}"
        self.api = HfApi(token=hf_token)
        self.config.data_dir.mkdir(parents=True, exist_ok=True)

    def save_local(self, pages_df: pl.DataFrame, chunks_df: pl.DataFrame, facts_df: pl.DataFrame):
        self.config.data_dir.mkdir(parents=True, exist_ok=True)

        if not pages_df.is_empty():
            pages_df.write_parquet(self.config.pages_file)
        if not chunks_df.is_empty():
            chunks_df.write_parquet(self.config.chunks_file)
        if not facts_df.is_empty():
            facts_df.write_parquet(self.config.facts_file)

        save_json(
            {
                "pages": len(pages_df),
                "chunks": len(chunks_df),
                "facts": len(facts_df),
                "has_indexes": (
                    self.config.faiss_file.exists()
                    and self.config.bm25_file.exists()
                ),
            },
            self.config.metadata_file,
        )

    def load_local(self) -> dict:
        result = {"pages": pl.DataFrame(), "chunks": pl.DataFrame(), "facts": pl.DataFrame()}

        if self.config.pages_file.exists():
            result["pages"] = pl.read_parquet(self.config.pages_file)
        if self.config.chunks_file.exists():
            result["chunks"] = pl.read_parquet(self.config.chunks_file)
        if self.config.facts_file.exists():
            result["facts"] = pl.read_parquet(self.config.facts_file)

        return result

    def push_to_hub(self):
        self._ensure_dataset_repo()

        files_to_upload = []

        if self.config.pages_file.exists():
            files_to_upload.append(self.config.pages_file)
        if self.config.chunks_file.exists():
            files_to_upload.append(self.config.chunks_file)
        if self.config.facts_file.exists():
            files_to_upload.append(self.config.facts_file)
        if self.config.faiss_file.exists():
            files_to_upload.append(self.config.faiss_file)
        if self.config.bm25_file.exists():
            files_to_upload.append(self.config.bm25_file)
        if self.config.metadata_file.exists():
            files_to_upload.append(self.config.metadata_file)

        for file_path in files_to_upload:
            try:
                upload_file(
                    path_or_fileobj=str(file_path),
                    path_in_repo=file_path.name,
                    repo_id=self.dataset_id,
                    repo_type="dataset",
                    token=self.hf_token,
                )
            except Exception as e:
                print(f"Failed to upload {file_path.name}: {e}")

        return f"https://huggingface.co/datasets/{self.dataset_id}"

    def pull_from_hub(self) -> dict:
        result = {"pages": pl.DataFrame(), "chunks": pl.DataFrame(), "facts": pl.DataFrame()}

        if not self._repo_exists():
            return result

        self.config.data_dir.mkdir(parents=True, exist_ok=True)

        file_mapping = {
            self.config.pages_path: "pages",
            self.config.chunks_path: "chunks",
            self.config.facts_path: "facts",
            self.config.faiss_index_path: None,
            self.config.bm25_index_path: None,
            self.config.metadata_path: None,
        }

        for file_name, key in file_mapping.items():
            try:
                downloaded = hf_hub_download(
                    repo_id=self.dataset_id,
                    filename=file_name,
                    repo_type="dataset",
                    token=self.hf_token,
                    local_dir=self.config.data_dir,
                    local_dir_use_symlinks=False,
                )
            except Exception:
                continue

        if self.config.pages_file.exists():
            result["pages"] = pl.read_parquet(self.config.pages_file)
        if self.config.chunks_file.exists():
            result["chunks"] = pl.read_parquet(self.config.chunks_file)
        if self.config.facts_file.exists():
            result["facts"] = pl.read_parquet(self.config.facts_file)

        return result

    def _ensure_dataset_repo(self):
        try:
            self.api.repo_info(repo_id=self.dataset_id, repo_type="dataset")
        except RepositoryNotFoundError:
            self.api.create_repo(
                repo_id=self.dataset_id,
                repo_type="dataset",
                private=False,
                exist_ok=True,
            )

    def _repo_exists(self) -> bool:
        try:
            self.api.repo_info(repo_id=self.dataset_id, repo_type="dataset")
            return True
        except RepositoryNotFoundError:
            return False

    def has_local_data(self) -> bool:
        return self.config.pages_file.exists() and self.config.chunks_file.exists()

    def get_stats(self) -> dict:
        stats = {"pages": 0, "chunks": 0, "facts": 0, "indexes": False}
        try:
            meta = load_json(self.config.metadata_file) if self.config.metadata_file.exists() else {}
            return {
                "pages": meta.get("pages", 0),
                "chunks": meta.get("chunks", 0),
                "facts": meta.get("facts", 0),
                "indexes": meta.get("has_indexes", False),
            }
        except Exception:
            return stats