import json
import os
from pathlib import Path
from typing import Dict, Iterator, List, Tuple

import huggingface_hub


class HPSPrompts:
    def __init__(self):
        super().__init__()
        self.hps_prompt_files = ['anime.json', 'concept-art.json', 'paintings.json', 'photo.json']
        self._download_benchmark_prompts()
        self.prompts: Dict[str, str] = {}
        self._size = 0
        for file in self.hps_prompt_files:
            category = file.replace('.json', '')
            with open(os.path.join('downloads/hps', file), 'r') as f:
                prompts = json.load(f)
                for i, prompt in enumerate(prompts):
                    if i == 100:
                        break
                    filename = f"{category}_{i:03d}.png"
                    self.prompts[filename] = prompt
                    self._size += 1
    
    def __iter__(self) -> Iterator[Tuple[str, Path]]:
        for filename, prompt in self.prompts.items():
            yield prompt, Path(filename)
    
    @property
    def name(self) -> str:
        return "hps"
    
    @property
    def size(self) -> int:
        return self._size

    def _download_benchmark_prompts(self) -> None:
        folder_name = Path('downloads/hps')
        folder_name.mkdir(parents=True, exist_ok=True)
        for file in self.hps_prompt_files:
            file_name = huggingface_hub.hf_hub_download("zhwang/HPDv2", file, subfolder="benchmark", repo_type="dataset")
            if not os.path.exists(os.path.join(folder_name, file)):
                os.symlink(file_name, os.path.join(folder_name, file))

    @property
    def metrics(self) -> List[str]:
        return ["hps"]