import gradio as gr
import json
import os
from huggingface_hub import HfApi, hf_hub_download
import requests
import hashlib

api = HfApi()


def hash_file(path: str) -> str:
    sha256_hash = hashlib.sha256()
    with open(path, "rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()


def hash_repo(repo_name: str):
    converted = []
    blobs = {}
    sub_repo_id = repo_name.replace("/", "-").lower()
    files = api.list_files_info(repo_name)
    for file in files:
        if file.lfs is not None:
            converted.append(
                {
                    "filename": file.rfilename,
                    "revision": "main",
                    "size": file.lfs["size"],
                    "sha256": file.lfs["sha256"],
                }
            )
            blobs[file.lfs["sha256"]] = [
                f"https://modelscope.cn/api/v1/models/hanamizukiai/{sub_repo_id}/repo?Revision=main&FilePath={file.rfilename}",
                # f"https://wisemodel.cn/file-proxy/hanamizuki-ai/{sub_repo_id}/-/raw/main/{file.rfilename}",
                # f"https://gitcode.net/spectator2639/{sub_repo_id}/-/raw/main/{file.rfilename}",
            ]
        else:
            filepath = hf_hub_download(repo_name, file.rfilename, revision="main")
            hash = hash_file(filepath)
            size = os.path.getsize(filepath)
            converted.append(
                {
                    "filename": file.rfilename,
                    "revision": "main",
                    "size": size,
                    "sha256": hash,
                }
            )
            blobs[hash] = [
                f"https://modelscope.cn/api/v1/models/hanamizukiai/{sub_repo_id}/repo?Revision=main&FilePath={file.rfilename}",
                # f"https://wisemodel.cn/file-proxy/hanamizuki-ai/{sub_repo_id}/-/raw/main/{file.rfilename}",
                # f"https://gitcode.net/spectator2639/{sub_repo_id}/-/raw/main/{file.rfilename}",
            ]

    wrapped_meta = {"repo_id": repo_name, "items": converted}
    return wrapped_meta, blobs


iface = gr.Interface(
    fn=hash_repo,
    inputs="text",
    outputs=["json", "json"],
    title="Repo Hasher",
)
iface.launch()