import json
import os

import argilla as rg
from huggingface_hub import HfApi

if __name__ == "__main__":
    rg.init(
        api_url=os.getenv("ARGILLA_API_URL"),
        api_key=os.getenv("ARGILLA_API_KEY"),
        extra_headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
    )

    dataset = rg.FeedbackDataset.from_argilla(
        os.getenv("SOURCE_DATASET"),
        workspace=os.getenv("SOURCE_WORKSPACE"),
    )
    dataset = dataset.filter_by(response_status=["submitted"])  # type: ignore

    output = {}
    for record in dataset.records:
        for response in record.responses:
            if response.user_id not in output:
                output[response.user_id] = 0
            output[response.user_id] += 1

    for key in list(output.keys()):
        output[rg.User.from_id(key).username] = output.pop(key)

    api = HfApi(token=os.getenv("HF_TOKEN"))
    api.upload_file(
        path_or_fileobj=json.dumps(output, indent=4),
        path_in_repo="stats.json",
        repo_id="DIBT/prompt-collective-dashboard",
        repo_type="space",
    )