import gradio as gr
import requests
import pandas as pd

api_url = 'https://huynhdoo--mps-api-query.modal.run'

origins = {
    'Formation' : ['formation.presentation', 'formation.summary'],
    'Métier' : ['metier.presentation', 'metier.competences',
                'metier.condition_travail', 'metier.nature_travail',
                'metier.acces_metier', 'metier.vie_professionnelle',
                'metier.accroche_metier', 'metier.format_court1',
                'metier.format_court2']
}

def API(origin='Formation', query='cuisine'):
    # Query API
    json = dict(
        query=query,
        origins=origins[origin]
    )

    resp = requests.post(url=api_url, json=json)
    data = resp.json()

    # Format result
    distances = pd.DataFrame({'distance': data['distances']}) 
    metadatas = pd.DataFrame(data['metadatas'])
    documents = pd.DataFrame({'document': data['documents']})
    df = pd.concat([distances, metadatas, documents], axis=1)
    df['distance'] = df['distance'].apply(lambda x: round(x, 3))
    df['origin'] = df['origin'].apply(lambda x: x.split('.')[1])
    return df

gradio_app = gr.Interface(
    fn=API,
    inputs=[
        gr.Dropdown(list(origins.keys()), label="Origine", info="Choisir un type de donnée à interroger"),
        gr.Textbox(label="Recherche", info="Votre recherche")
    ],
    outputs=[
        gr.DataFrame(label="Résultats", headers=["Distance", "Key", "Label", "Origin", "Document"])
    ],
    examples=[['Formation', 'militaire'],
              ['Métier', 'cuisine'],
              ['Formation', 'écologie'],
              ['Métier', 'eau'],
              ['Formation', 'math'],
    ],
    cache_examples=False
)

if __name__ == "__main__":
    gradio_app.launch(auth=("mps", "sup"), share=True)