logic/data_utils.py

import os
import json
import time
import glob
from huggingface_hub import HfApi, create_repo, CommitScheduler
import bcrypt
import shutil
import uuid
import gradio as gr
from PIL import Image
import numpy as np
from data.words_map import words_mapping
# from logic.supabase_client import auth_handler

def load_concepts(path="data/concepts.json"):
    with open(path, encoding='utf-8') as f:
        data = json.load(f)
    sorted_data = dict()
    duplicate_entries = []
    for country in sorted(data):
        sorted_data[country] = dict()
        for lang in sorted(data[country]):
            sorted_data[country][lang] = dict()
            for cat, concepts in data[country][lang].items():
                sorted_data[country][lang][cat] = sorted(set(concepts))
            all_concepts = [v for l in sorted_data[country][lang].values() for v in l]
            unq_concepts = set(all_concepts)
            if len(unq_concepts) < len(all_concepts):
                duplicate_entries.append(f"{country}/{lang}")
    if duplicate_entries:
        # raise ValueError(f"Duplicate concepts in: {duplicate_entries}")
        print(f"Duplicate concepts in: {duplicate_entries}")  # FIXME raise error above

    return sorted_data


def _normalize_concepts_list(val):
    """Ensure category_N_concepts is a list of strings, never None or nested None."""
    if val is None:
        return [""]
    if not isinstance(val, list):
        return [str(val)] if val else [""]
    return [str(x).strip() if x is not None else "" for x in val] or [""]


_VLM_USAGE_KEYS = (
    "prompt_tokens",
    "completion_tokens",
    "total_tokens",
    "reasoning_tokens",
)


def _normalize_vlm_usage(value):
    """Keep a stable Arrow-compatible schema across old and new samples."""
    value = value if isinstance(value, dict) else {}
    normalized = {}
    for key in _VLM_USAGE_KEYS:
        try:
            normalized[key] = int(value.get(key) or 0)
        except (TypeError, ValueError):
            normalized[key] = 0
    return normalized


def _extract_timestamp_from_id(raw_id):
    if not raw_id or not isinstance(raw_id, str):
        return None
    suffix = raw_id.rsplit("_", 1)[-1].strip()
    return int(suffix) if suffix.isdigit() else None


def mark_sample_excluded(example_id, country, language, username, local_ds_folder):
    """
    Directly update JSON file(s) to set excluded=True. Bypasses full save flow.
    Returns True if at least one file was updated.
    """
    if not example_id or not country or not language or not username:
        return False
    username = username.strip()
    if not username:
        return False
    user_root = os.path.join(
        local_ds_folder,
        "logged_in_users",
        country,
        language,
        username,
    )
    ts = _extract_timestamp_from_id(example_id)
    updated = False
    try:
        for fp in glob.glob(os.path.join(user_root, "**", "*.json"), recursive=True):
            try:
                with open(fp, "r", encoding="utf-8") as f:
                    data = json.load(f)
                file_id = data.get("id")
                if file_id != example_id and (ts is None or _extract_timestamp_from_id(file_id) != ts):
                    continue
                data["excluded"] = True
                with open(fp, "w", encoding="utf-8") as f:
                    json.dump(data, f, indent=2)
                updated = True
            except Exception:
                continue
    except Exception:
        pass
    return updated


def _normalize_text(value):
    return value.strip() if isinstance(value, str) else value


def _canonicalize_category(values_dic):
    """
    Keep saved categories in the canonical taxonomy language.
    The UI may display localized labels, but JSON should store concept keys.
    """
    category = _normalize_text(values_dic.get("category"))
    country = _normalize_text(values_dic.get("country"))
    language = _normalize_text(values_dic.get("language"))
    if not category or not country or not language:
        return

    concepts_dict = values_dic.get("concepts_dict") or {}
    country_lang_map = values_dic.get("country_lang_map") or {}
    language_lookup = country_lang_map.get(language, language)
    categories = concepts_dict.get(country, {}).get(language_lookup, {})
    if not isinstance(categories, dict):
        return

    if category in categories:
        values_dic["category"] = category
        return

    reverse_translations = {
        display: canonical
        for canonical, display in (words_mapping.get(language) or {}).items()
    }
    canonical = reverse_translations.get(category)
    if canonical in categories:
        values_dic["category"] = canonical
        return

    values_dic["category"] = None


def _infer_category_from_concept(values_dic):
    """
    Best-effort server-side category recovery.
    Useful when UI state loses dropdown value but concept is still present.
    """
    if values_dic.get("category"):
        return

    concept = _normalize_text(values_dic.get("concept"))
    country = _normalize_text(values_dic.get("country"))
    language = _normalize_text(values_dic.get("language"))
    if not concept or not country or not language:
        return

    concepts_dict = values_dic.get("concepts_dict") or {}
    country_lang_map = values_dic.get("country_lang_map") or {}
    language_lookup = country_lang_map.get(language, language)

    categories = concepts_dict.get(country, {}).get(language_lookup, {})
    if not isinstance(categories, dict):
        return

    matches = [cat for cat, concepts in categories.items() if concept in (concepts or [])]
    if len(matches) == 1:
        values_dic["category"] = matches[0]


def load_metadata(path="data/metadata.json"):
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    sorted_data = dict()
    for country in sorted(data):
        sorted_data[country] = dict()
        for lang in sorted(data[country]):
            sorted_data[country][lang] = data[country][lang]
    return sorted_data


class CustomHFDatasetSaver:
    def __init__(self, api_token, dataset_name, private=False):
        self.api_token = api_token
        self.dataset_name = dataset_name
        self.private = private
        self.api = HfApi()
        
    def setup(self, data_outputs, local_ds_folder, metadata_dict=None):
        # create repo is not exist
        self.dataset_name = create_repo(
            repo_id=self.dataset_name,
            token=self.api_token,
            private=self.private,
            repo_type="dataset",
            exist_ok=True,
        ).repo_id

        # Create the local data folder if not exist
        self.local_ds_folder = local_ds_folder
        os.makedirs(self.local_ds_folder, exist_ok=True)

        # Migrate any existing JSON files to include new VLM fields
        self._migrate_existing()

        self.data_outputs = data_outputs  # list of components to read values from
        self.metadata_dict = metadata_dict or {}  # localized UI strings

        # create scheduler to commit the data to the hub every x minutes
        self.scheduler = CommitScheduler(
            repo_id=self.dataset_name,
            repo_type="dataset",
            folder_path=self.local_ds_folder,
            every=1,
            token=self.api_token,
        )

    def _migrate_existing(self):
        """Normalize the schema of existing sample JSON files."""
        for root, _, files in os.walk(self.local_ds_folder):
            for fname in files:
                if not fname.endswith(".json"):
                    continue
                fpath = os.path.join(root, fname)
                with open(fpath, "r+", encoding="utf-8") as f:
                    data = json.load(f)
                    if fname == "profiles.json" and isinstance(data, dict):
                        # Repair keys added by the previous over-broad migration.
                        migration_keys = {
                            "vlm_caption",
                            "vlm_feedback",
                            "vlm_model",
                            "hf_username",
                            "vlm_usage",
                        }
                        cleaned = {key: value for key, value in data.items() if key not in migration_keys}
                        if cleaned != data:
                            f.seek(0)
                            json.dump(cleaned, f, indent=2)
                            f.truncate()
                        continue
                    if not isinstance(data, dict) or "id" not in data or "image" not in data:
                        continue
                    updated = False
                    for key in ["vlm_caption", "vlm_feedback", "vlm_model", "hf_username"]:
                        if key not in data:
                            data[key] = ""
                            updated = True
                    normalized_usage = _normalize_vlm_usage(data.get("vlm_usage"))
                    if data.get("vlm_usage") != normalized_usage:
                        data["vlm_usage"] = normalized_usage
                        updated = True
                    if updated:
                        f.seek(0)
                        json.dump(data, f, indent=2)
                        f.truncate()
    
    def validate_data(self, values_dic):
        """
        Validates the data before saving to ensure no required fields are empty.
        Returns (bool, str) tuple where first value indicates if validation passed
        and second value contains error message if validation failed.
        """
        # Remove 'image' from required fields since we handle it separately
        required_fields = ['country', 'language', 'category', 'concept']
        if not values_dic.get('quick_submit_mode'):
            required_fields.append('caption')
        
        # Check if image is provided (either uploaded or via URL)
        image = values_dic.get('image')
        image_url = values_dic.get('image_url')
        
        # Check if image exists and is not None
        has_image = image is not None and (
            isinstance(image, dict)
            or isinstance(image, Image.Image)
            or (hasattr(image, "shape") and image.shape[0] > 0)
            or (isinstance(image, str) and image.strip() != "" and os.path.exists(image))
        )
        has_url = image_url is not None and image_url.strip() != ""
        
        if not has_image and not has_url:
            return False, "Either an image or image URL must be provided"
            
        # Check required fields
        for field in required_fields:
            value = values_dic.get(field)
            if value is None or (isinstance(value, str) and value.strip() == ""):
                return False, f"Required field '{field}' cannot be empty"

        # VLM feedback is required when a VLM caption was generated
        vlm_caption = values_dic.get('vlm_caption')
        vlm_feedback = values_dic.get('vlm_feedback')
        if vlm_caption and str(vlm_caption).strip():
            if not vlm_feedback or not str(vlm_feedback).strip():
                country = values_dic.get('country', '')
                language = values_dic.get('language', '')
                meta = self.metadata_dict.get(country, {}).get(language, {}) if country and language else {}
                msg = meta.get(
                    'VLM_Feedback_Required_msg',
                    'Please select Yes or No for the VLM description feedback before submitting.',
                )
                return False, msg
                
        # Check if image file exists if image path is provided
        if has_image and isinstance(image, dict):
            if not os.path.exists(image.get('path', '')):
                return False, "Image file not found"
                
        return True, ""

    #TODO: add a function to check if the user is logged in
    def is_logged_in(self):
        pass

    #TODO: check if the user is logged in (add a decorator to the save function)
    def save(self, *values):
        # 'values' are the outputs from your data collection components,
        # you can map these to field names as needed
        values_dic = dict(zip(self.data_outputs, values))
        # Normalize text fields and recover category if the UI dropped it.
        for key in ("country", "language", "category", "concept", "caption", "image_url"):
            values_dic[key] = _normalize_text(values_dic.get(key))
        _canonicalize_category(values_dic)
        _infer_category_from_concept(values_dic)
        
        # print(f"Values received: {values_dic}")
        
        # Validate data before proceeding
        is_valid, error_msg = self.validate_data(values_dic)
        if not is_valid:
            raise gr.Error(error_msg)
            # raise ValueError(error_msg)
            
        # Password field is legacy; keep for backward compatibility if present.
        if "password" in values_dic and values_dic["password"]:
            values_dic["password"] = self.hash_password(values_dic["password"])

        # # Process main category and concept
        # main_category = values_dic.get('category', '')
        # main_concept = values_dic.get('concept', '')
        
        # # Process category-specific concept dropdowns
        # additional_concepts_by_category = {}
        # # Extract predefined categories and their corresponding dropdowns from values_dic
        # predefined_categories = sorted(list(values_dic.get('concepts_dict', {})
        #                             .get(values_dic.get('country', 'USA'), {})
        #                             .get(values_dic.get('language', 'English'), {}).keys()))[:5]
        # # Process each category dropdown
        # for i, category in enumerate(predefined_categories):
        #     dropdown_key = f'category{i+1}_concepts'
        #     if dropdown_key in values_dic and values_dic[dropdown_key]:
        #         # Only add non-empty concept selections
        #         if values_dic[dropdown_key]:
        #             additional_concepts_by_category[category] = values_dic[dropdown_key]

        ### TODO: fix saving additional concepts if not displayed in English
        # # Process category-specific concept dropdowns
        # additional_concepts_by_category = {}

        # # Extract the country and language
        # country = values_dic.get('country', 'USA')
        # language = values_dic.get('language', 'English')
        # concepts_dict = values_dic.get('concepts_dict', {})
        # lang2eng_mapping = values_dic.get('country_lang_map', {})

        # # Get the English version of the language for dictionary lookup
        # eng_lang = lang2eng_mapping.get(language, language)

        # # Get the predefined categories in English
        # predefined_categories = sorted(list(concepts_dict.get(country, {}).get(eng_lang, {}).keys()))[:5]

        # # Process each category dropdown
        # for i, category in enumerate(predefined_categories):
        #     dropdown_key = f'category_{i+1}_concepts'
        #     if dropdown_key in values_dic and values_dic[dropdown_key]:
        #         # Only add non-empty concept selections
        #         additional_concepts_by_category[category] = values_dic[dropdown_key]

        raw_id = values_dic.get("id")
        parsed_timestamp = _extract_timestamp_from_id(raw_id)
        current_timestamp = parsed_timestamp or int(time.time() * 1000)
        country = values_dic.get("country") or ""
        language = values_dic.get("language") or ""
        category = values_dic.get("category") or ""
        concept = values_dic.get("concept") or ""
        values_dic["id"] = f"{country}_{language}_{category}_{concept}_{current_timestamp}"

        # Prepare the main directory of the sample.
        email = (values_dic.get("username") or "").strip()
        if email:
            sample_dir = os.path.join(
                "logged_in_users",
                values_dic["country"],
                values_dic["language"],
                email,
                str(current_timestamp),
            )
            print(f"Sample directory for logged in user: {sample_dir}")
        else:
            sample_dir = os.path.join(
                "anonymous_users",
                values_dic["country"],
                values_dic["language"],
                str(uuid.uuid4()),
                str(current_timestamp),
            )
        print(f"Sample directory: {sample_dir}")

        os.makedirs(os.path.join(self.local_ds_folder, sample_dir), exist_ok=True)

        # Destination path 
        dest_image_path = os.path.join(sample_dir, "image.png")
        # Source path (to be used for copying the file in the with lock block)
        # This is the path of the image file that was uploaded by the user
        # I want to save the values_dic['image'] in the dest_image_path
        # Convert numpy array to PIL Image and save it
        # ===
        # uploaded_image_path = os.path.join(self.local_ds_folder, dest_image_path)
        # img = Image.fromarray(values_dic['image'])
        # img.save(uploaded_image_path)
        full_dest_path = os.path.join(self.local_ds_folder, dest_image_path)
        
        # Handle different image types
        image_data = values_dic['image']
        
        if isinstance(image_data, dict) and 'path' in image_data:
            # New upload case - copy from the uploaded path
            uploaded_image_path = image_data['path']
            with self.scheduler.lock:
                shutil.copy(uploaded_image_path, full_dest_path)
        elif isinstance(image_data, np.ndarray):  # not values_dic.get('excluded', False) and 
            # Exclude case with numpy array - save the array as an image
            with self.scheduler.lock:
                # Convert numpy array to PIL image and save
                img = Image.fromarray(image_data)
                img.save(full_dest_path)
        elif isinstance(image_data, Image.Image):
            # PIL image case
            with self.scheduler.lock:
                image_data.save(full_dest_path)
        
        values_dic['image'] = dest_image_path
        image_file_path_on_hub =  f"https://huggingface.co/datasets/{self.dataset_name}/resolve/main/{dest_image_path}"

        # print(f"Saving sample: {values}")
        # Build the metadata dictionary.
        data_dict = {
            # in case using windows
            "image": values_dic['image'].replace("\\", "/"),
            "image_file": image_file_path_on_hub.replace("\\", "/"),
            # "image": values_dic['image'],
            # "image_file": image_file_path_on_hub,
            "image_url": values_dic['image_url'] or "",
            "caption": values_dic['caption'] or "",
            "vlm_caption": values_dic['vlm_caption'] or "",
            "vlm_feedback": values_dic['vlm_feedback'] or "" if values_dic['vlm_caption'] else "",
            "vlm_model": values_dic['vlm_model'] or "" if values_dic['vlm_caption'] else "",
            "vlm_usage": _normalize_vlm_usage(values_dic.get("vlm_usage")),
            "country": values_dic['country'] or "",
            "language": values_dic['language'] or "",
            "category": values_dic['category'] or "",
            "concept": values_dic['concept'] or "",
            "category_1_concepts": _normalize_concepts_list(values_dic.get('category_1_concepts')),
            "category_2_concepts": _normalize_concepts_list(values_dic.get('category_2_concepts')),
            "category_3_concepts": _normalize_concepts_list(values_dic.get('category_3_concepts')),
            "category_4_concepts": _normalize_concepts_list(values_dic.get('category_4_concepts')),
            "category_5_concepts": _normalize_concepts_list(values_dic.get('category_5_concepts')),
            "timestamp": current_timestamp,
            "username": email,
            "hf_username": values_dic.get("hf_username") or "",
            "password": values_dic.get('password') or "",
            "id": str(values_dic["id"]),
            "excluded": False if values_dic.get('excluded') is None else bool(values_dic.get('excluded')),
            # "is_blurred": str(values_dic.get('is_blurred'))
        }

        print(f"Data dictionary: {data_dict}")
        # Define a unique filename for the JSON metadata file (stored in self.folder).
        json_filename = f"sample_{current_timestamp}.json"
        json_file_path = os.path.join(self.local_ds_folder, sample_dir, json_filename)

        with self.scheduler.lock:
            # Save the metadata to the sample file in the local dataset folder
            with open(json_file_path, "w", encoding="utf-8") as f:
                json.dump(data_dict, f, indent=2)

        user_root = os.path.join(
            self.local_ds_folder,
            "logged_in_users",
            country,
            language,
            email,
        )
        try:
            for fp in glob.glob(os.path.join(user_root, "**", "*.json"), recursive=True):
                try:
                    with open(fp, "r", encoding="utf-8") as f:
                        data = json.load(f)
                    file_id = data.get("id")
                    if _extract_timestamp_from_id(file_id) != current_timestamp:
                        continue
                    if file_id == data_dict["id"]:
                        continue
                    data["id"] = data_dict["id"]
                    data["category"] = data_dict["category"]
                    data["concept"] = data_dict["concept"]
                    data["category_1_concepts"] = data_dict.get("category_1_concepts", [""])
                    data["category_2_concepts"] = data_dict.get("category_2_concepts", [""])
                    data["category_3_concepts"] = data_dict.get("category_3_concepts", [""])
                    data["category_4_concepts"] = data_dict.get("category_4_concepts", [""])
                    data["category_5_concepts"] = data_dict.get("category_5_concepts", [""])
                    data["excluded"] = data_dict.get("excluded", False)
                    with open(fp, "w", encoding="utf-8") as f:
                        json.dump(data, f, indent=2)
                except Exception:
                    continue
        except Exception:
            pass

        print("Data saved successfully")

    def hash_password(self, raw_password):
        """
        Hashes a raw password using bcrypt and returns the hashed password.
            raw_password (str): The plain text password to be hashed.
            str: The hashed password as a string.
        """
        hashed_password = bcrypt.hashpw(raw_password.encode(), bcrypt.gensalt()).decode()
        return hashed_password