database.py

import json
import os
import re
import shutil
import asyncio
from huggingface_hub import HfApi, hf_hub_download
import config

import tempfile

DATA_FILE_NAME = "database.json"
LOCAL_PATH = os.path.join(tempfile.gettempdir(), DATA_FILE_NAME)

# In-memory database state
db_state = {
    "admins": [],       # List of admin user IDs (ints)
    "users": [],        # List of authorized user IDs (ints)
    "urls": [],         # List of dicts: {"url": str, "added_by": int}
    "user_metadata": {} # Dict: "user_id_str" -> {"username": str, "name": str}
}

# Helper to extract repository ID from a HF URL
def extract_repo_id(url: str) -> str:
    if not url:
        return None
    url = url.strip()
    if "/" in url and not url.startswith("http"):
        return url
    # Match standard dataset url: huggingface.co/datasets/username/dataset-name
    match = re.search(r"huggingface\.co/datasets/([^/]+)/([^/]+)", url, re.IGNORECASE)
    if match:
        return f"{match.group(1)}/{match.group(2)}"
    return None

repo_id = extract_repo_id(config.HF_DATASET_URL)

def get_api():
    if not config.HF_TOKEN:
        return None
    return HfApi(token=config.HF_TOKEN)

def init_db():
    global db_state
    if not repo_id or not config.HF_TOKEN:
        print("HF_DATASET_URL or HF_TOKEN not configured. Working locally.")
        if os.path.exists(LOCAL_PATH):
            try:
                with open(LOCAL_PATH, "r", encoding="utf-8") as f:
                    loaded = json.load(f)
                    # Migrate or update local state
                    for key in db_state:
                        if key in loaded:
                            db_state[key] = loaded[key]
                print("Loaded database from local file.")
            except Exception as e:
                print(f"Error loading local database: {e}")
        return

    # Sync from Hugging Face Dataset
    try:
        api = get_api()
        # Ensure repository exists
        api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
        
        # Download file
        print(f"Downloading database from HF Dataset: {repo_id}")
        downloaded_path = hf_hub_download(
            repo_id=repo_id,
            filename=DATA_FILE_NAME,
            repo_type="dataset",
            token=config.HF_TOKEN
        )
        
        # Copy to local path
        shutil.copy(downloaded_path, LOCAL_PATH)
        
        # Load JSON
        with open(LOCAL_PATH, "r", encoding="utf-8") as f:
            loaded = json.load(f)
            for key in db_state:
                if key in loaded:
                    db_state[key] = loaded[key]
                    
        print("Successfully synced database from Hugging Face.")
    except Exception as e:
        err_msg = str(e)
        if "404" in err_msg or "Entry Not Found" in err_msg:
            print("Database file not found in HF dataset. Initializing a new one.")
        else:
            print(f"Error loading database from HF: {err_msg}. Initializing a new one.")
        if os.path.exists(LOCAL_PATH):
            try:
                with open(LOCAL_PATH, "r", encoding="utf-8") as f:
                    loaded = json.load(f)
                    for key in db_state:
                        if key in loaded:
                            db_state[key] = loaded[key]
            except Exception:
                pass
        save_db()

def save_db():
    global db_state
    # Ensure types and clean duplicates
    db_state["admins"] = list(sorted(list(set(int(uid) for uid in db_state.get("admins", [])))))
    db_state["users"] = list(sorted(list(set(int(uid) for uid in db_state.get("users", [])))))
    
    unique_urls = []
    seen = set()
    for item in db_state.get("urls", []):
        url = item.get("url")
        if url not in seen:
            seen.add(url)
            unique_urls.append({
                "url": url,
                "added_by": int(item.get("added_by"))
            })
    db_state["urls"] = unique_urls

    # Save locally
    try:
        with open(LOCAL_PATH, "w", encoding="utf-8") as f:
            json.dump(db_state, f, indent=2, ensure_ascii=False)
        print("Database saved locally.")
    except Exception as e:
        print(f"Error saving database locally: {e}")
        return

    # Upload to HF
    if not repo_id or not config.HF_TOKEN:
        return

    try:
        api = get_api()
        api.upload_file(
            path_or_fileobj=LOCAL_PATH,
            path_in_repo=DATA_FILE_NAME,
            repo_id=repo_id,
            repo_type="dataset",
            token=config.HF_TOKEN
        )
        print("Database uploaded and synced to Hugging Face dataset.")
    except Exception as e:
        print(f"Error uploading database to Hugging Face: {e}")

async def save_db_async():
    await asyncio.to_thread(save_db)

# Normalization of Hugging Face Space URLs
def normalize_hf_url(url: str) -> str:
    url = url.strip()
    
    # Remove any query parameters or trailing slashes
    url = url.split("?")[0].rstrip("/")
    
    # Case 1: Match standard Hugging Face space URLs
    # e.g., https://huggingface.co/spaces/username/spacename
    # or huggingface.co/spaces/username/spacename/tree/main
    # We want to extract username and spacename.
    hf_co_pattern = r"(?:https?://)?(?:www\.)?huggingface\.co/spaces/([^/]+)/([^/]+)"
    match = re.match(hf_co_pattern, url, re.IGNORECASE)
    if match:
        username = match.group(1).lower().replace("_", "-").replace(".", "-")
        spacename = match.group(2).lower().replace("_", "-").replace(".", "-")
        return f"https://{username}-{spacename}.hf.space"
    
    # Case 2: Match direct hf.space URLs
    # e.g., https://username-spacename.hf.space
    # or username-spacename.hf.space/subpath
    hf_space_pattern = r"(?:https?://)?([^/.]+)\.hf\.space"
    match = re.match(hf_space_pattern, url, re.IGNORECASE)
    if match:
        subdomain = match.group(1).lower().replace("_", "-").replace(".", "-")
        return f"https://{subdomain}.hf.space"
        
    # Case 3: Match simple username/spacename form
    # e.g., username/spacename
    simple_pattern = r"^([^/]+)/([^/]+)$"
    match = re.match(simple_pattern, url)
    if match:
        username = match.group(1).lower().replace("_", "-").replace(".", "-")
        spacename = match.group(2).lower().replace("_", "-").replace(".", "-")
        return f"https://{username}-{spacename}.hf.space"
        
    raise ValueError("Invalid Hugging Face Space URL. Make sure it contains 'huggingface.co/spaces/...' or '<subdomain>.hf.space'")

# Authorization helper functions
def is_super_admin(user_id: int) -> bool:
    return config.SUPER_ADMIN_ID is not None and user_id == config.SUPER_ADMIN_ID

def is_admin(user_id: int) -> bool:
    if is_super_admin(user_id):
        return True
    return user_id in db_state.get("admins", [])

def is_authorized(user_id: int) -> bool:
    if is_admin(user_id):
        return True
    return user_id in db_state.get("users", [])

# User metadata updates
def update_user_metadata(user_id: int, username: str = None, first_name: str = None, last_name: str = None):
    if "user_metadata" not in db_state:
        db_state["user_metadata"] = {}
    
    name = first_name or ""
    if last_name:
        name += f" {last_name}"
    name = name.strip()
    
    user_id_str = str(user_id)
    # Check if there is any change to avoid redundant saves
    existing = db_state["user_metadata"].get(user_id_str)
    if not existing or existing.get("username") != username or existing.get("name") != name:
        db_state["user_metadata"][user_id_str] = {
            "username": username,
            "name": name
        }
        save_db()

def get_user_display(user_id: int) -> str:
    user_id_str = str(user_id)
    metadata = db_state.get("user_metadata", {}).get(user_id_str)
    if metadata:
        username = metadata.get("username")
        name = metadata.get("name")
        if username:
            return f"{name} (@{username}) [`{user_id}`]"
        elif name:
            return f"{name} [`{user_id}`]"
    return f"`{user_id}`"

# User administration functions
def add_admin(user_id: int) -> bool:
    if user_id in db_state["admins"]:
        return False
    db_state["admins"].append(user_id)
    save_db()
    return True

def remove_admin(user_id: int) -> bool:
    if user_id not in db_state["admins"]:
        return False
    db_state["admins"].remove(user_id)
    save_db()
    return True

def add_user(user_id: int) -> bool:
    if user_id in db_state["users"]:
        return False
    db_state["users"].append(user_id)
    save_db()
    return True

def remove_user(user_id: int) -> bool:
    if user_id not in db_state["users"]:
        return False
    db_state["users"].remove(user_id)
    save_db()
    return True

def list_admins() -> list:
    return db_state["admins"]

def list_users() -> list:
    return db_state["users"]

# URL management functions
def add_url(url: str, added_by: int) -> tuple[bool, str]:
    normalized = normalize_hf_url(url)
    # Check if already exists
    for item in db_state["urls"]:
        if item["url"] == normalized:
            return False, normalized
    db_state["urls"].append({
        "url": normalized,
        "added_by": added_by
    })
    save_db()
    return True, normalized

def remove_url(url: str, user_id: int) -> tuple[bool, str]:
    try:
        normalized = normalize_hf_url(url)
    except ValueError:
        normalized = url.strip()

    # Find the URL
    target_idx = -1
    for idx, item in enumerate(db_state["urls"]):
        if item["url"] == normalized or item["url"].rstrip("/") == normalized.rstrip("/"):
            target_idx = idx
            break
            
    if target_idx == -1:
        return False, normalized
        
    item = db_state["urls"][target_idx]
    # Check authorization
    # Super admins and admins can remove any URL.
    # Regular users can only remove their own.
    if not is_admin(user_id) and item["added_by"] != user_id:
        raise PermissionError("You can only remove URLs that you added.")
        
    db_state["urls"].pop(target_idx)
    save_db()
    return True, normalized

def get_urls(user_id: int) -> list:
    # If admin or super admin, return all
    if is_admin(user_id):
        return db_state["urls"]
    # Else, return only owned by this user
    return [item for item in db_state["urls"] if item["added_by"] == user_id]

def get_all_urls() -> list:
    return [item["url"] for item in db_state["urls"]]