import os, json, logging
import requests
from transformers import pipeline

from flask import Flask, request, jsonify
app=Flask(__name__)
@app.route("/")
def hello():
    return "welcome!"

def parse_params():
    input=request.args.get('input')
    kargs=request.args.get('kargs')
    try:
        input = json.load(input)
    except:
        pass
    try:
        kargs = json.load(kargs)
    except:
        kargs={}
    return input, kargs
    
from huggingface_hub import HfApi
api=HfApi()
@app.route("/search")
def search():
    args=request.args.to_dict()
    models = api.list_models(**args)
    return jsonify(models)

@app.route("/task_list")
def tasks():
    return [item.strip() for item in 
        '''audio-classification
        automatic-speech-recognition
        conversational
        depth-estimation
        document-question-answering
        feature-extraction
        fill-mask
        image-classification
        image-feature-extraction
        image-segmentation
        image-to-image
        image-to-text
        mask-generation
        object-detection
        question-answering
        summarization
        table-question-answering
        text2text-generation
        text-classification (alias sentiment-analysis available)
        text-generation
        text-to-audio (alias text-to-speech available)
        token-classification (alias ner available)
        translation
        translation_xx_to_yy
        video-classification
        visual-question-answering
        zero-shot-classification
        zero-shot-image-classification
        zero-shot-audio-classification
        zero-shot-object-detection'''.split("\n")
    ]

@app.route("/<task>")
def run_task(task):
    (input, kargs)=parse_params()
    pipe=pipeline(task, **kargs)
    return pipe(input)

@app.route("/<user>/<model>")
def run_model():
    (input, kargs)=parse_params()
    pipe=pipeline(model=f'{user}/{model}', **kargs)
    return pipe(input)

@app.route("/<task>/<user>/<model>")
def run_task_model():
    (input, kargs)=parse_params()
    pipe=pipeline(task, model=f'{user}/{model}', **kargs)
    return pipe(input)

logging.info("xtts is ready")