from flask import Flask, request
import json
from huggingface_hub import HfApi, ModelFilter, DatasetFilter, ModelSearchArguments
from pprint import pprint

app = Flask(__name__)


@app.route("/hello")
def hello():
    return "<h1 style='color:blue'>Hello There!</h1>"


@app.route("/hfapi/search", methods=["POST"])
def hf_api():
    request_data = request.get_json()
    query = request_data.get("query")
    filters = json.loads(request_data.get("filters"))
    limit = request_data.get("limit", 5)
    print("query", query)
    print("filters", filters)
    print("limit", limit)

    api = HfApi()
    filt = ModelFilter(
        task=filters["task"],
        library=filters["library"],
    )
    models = api.list_models(search=query, filter=filt, limit=limit, full=True)
    res = []
    for model in models:
        model = model.__dict__
        res.append(
            {
                "modelId": model.get("modelId"),
                "tags": model.get("tags"),
                "downloads": model.get("downloads"),
                "likes": model.get("likes"),
            }
        )
    count = len(res)
    if len(res) > limit:
        res = res[:limit]
    pprint(res)
    return json.dumps({"value": res, "count": count})


@app.route("/custom/search", methods=["POST"])
def main():
    request_data = request.get_json()
    query = request_data.get("query")
    filters = json.loads(request_data.get("filters"))
    limit = request_data.get("limit", 5)
    print("query", query)
    print("filters", filters)
    print("limit", limit)

    # records, count_filtered = search_query(query=request_data["query"], filters=filters, top=request_data["top"])
    # assert len(set([record["id"] for record in records])) == len(records), "ids of results are not unique"
    # res = {"value": records, "count": count_filtered}
    # return json.dumps(res, indent=2)


if __name__ == "__main__":
    app.run(host="localhost", port=5000)