import io
import json
from dataclasses import dataclass
from typing import Any, List

from huggingface_hub import HfApi, create_repo
from loguru import logger


@dataclass
class DreamboothPreprocessor:
    concept_images: List[Any]
    concept_name: str
    username: str
    project_name: str
    token: str

    def __post_init__(self):
        self.repo_name = f"{self.username}/autotrain-data-{self.project_name}"
        try:
            create_repo(
                repo_id=self.repo_name,
                repo_type="dataset",
                token=self.token,
                private=True,
                exist_ok=False,
            )
        except Exception:
            logger.error("Error creating repo")
            raise ValueError("Error creating repo")

    def _upload_concept_images(self, file, api):
        logger.info(f"Uploading {file} to concept1")
        api.upload_file(
            path_or_fileobj=file.name,
            path_in_repo=f"concept1/{file.name.split('/')[-1]}",
            repo_id=self.repo_name,
            repo_type="dataset",
            token=self.token,
        )

    def _upload_concept_prompts(self, api):
        _prompts = {}
        _prompts["concept1"] = self.concept_name

        prompts = json.dumps(_prompts)
        prompts = prompts.encode("utf-8")
        prompts = io.BytesIO(prompts)
        api.upload_file(
            path_or_fileobj=prompts,
            path_in_repo="prompts.json",
            repo_id=self.repo_name,
            repo_type="dataset",
            token=self.token,
        )

    def prepare(self):
        api = HfApi()
        for _file in self.concept_images:
            self._upload_concept_images(_file, api)

        self._upload_concept_prompts(api)