import gradio as gr
import requests
import pandas as pd

def api_url(remote):
    return f"https://huynhdoo--mps-api-{remote}.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 retrieve(query='cuisine', origin='Formation', n_results=10):
    queries = query.split(',')

    # Query API
    json = dict(
        query=queries,
        origins=origins[origin],
        n_results=n_results
    )

    try:
        resp = requests.post(url=api_url('retrieve'), json=json)
        data = resp.json()
    except:
        return pd.DataFrame([])

    # Format result
    queries = [pd.DataFrame({'query': [queries[i]]*len(data['distances'][i])}) for i in range(len(queries))]
    distances = [pd.DataFrame({'distance': distances}) for distances in data['distances']]
    metadatas = [pd.DataFrame(metadatas) for metadatas in data['metadatas']]
    documents = [pd.DataFrame({'document': documents}) for documents in data['documents']]
    # ranking = pd.DataFrame({'rang': rank(query, data['documents'])})
    # df = pd.concat([ranking, distances, metadatas, documents], axis=1)
    dfs = [pd.concat([query, dis, metas, docs], axis=1) for query, dis, metas, docs in zip(queries, distances, metadatas, documents)]
    
    for idx in range(len(dfs)):
        dfs[idx]['distance'] = dfs[idx]['distance'].apply(lambda x: round(x, 3))
        dfs[idx]['origin'] = dfs[idx]['origin'].apply(lambda x: x.split('.')[1])
    
    df = pd.concat(dfs)
    return df

def rank(query='cuisine', documents = []):
    # Query API
    json = dict(
        query=query,
        documents=documents
    )

    try:
        resp = requests.post(url=api_url('rank'), json=json)
        data = resp.json()
    except:
        return []

    # Format result
    return data['ranking']

gradio_app = gr.Interface(
    fn=retrieve,
    inputs=[
        gr.Textbox(label="Recherche", info="Votre recherche (mots séparés par des virgules pour des recherches multiples)"),
        gr.Dropdown(choices=list(origins.keys()), value=list(origins.keys())[0], label="Origine", info="Choisir un type de donnée à interroger"),
        gr.Number(value=10, label="Nombre de résultats", info="Nombre de résultats attendus")
    ],
    outputs=[
        gr.DataFrame(label="Résultats", headers=["Query", "Distance", "Key", "Label", "Origin", "Document"])
    ],
    examples=[['militaire', 'Formation'],
              ['cuisine', 'Métier'],
              ['écologie', 'Formation'],
              ['eau', 'Métier'],
              ['math', 'Formation'],
              ['math, histoire, math histoire', 'Formation']
    ],
    cache_examples=False
)

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