from typing import Optional

import gradio as gr
import pandas as pd

from utils.similarity import batch_cos_sim


def read_data(filepath: str) -> Optional[pd.DataFrame]:
    if not filepath:
        return None
    if filepath.endswith('.xlsx'):
        df = pd.read_csv(filepath)
    elif filepath.endswith('.csv'):
        df = pd.read_csv(filepath)
    else:
        raise Exception('File type not supported')
    return df


def process(model_name: str,
            prompt: str,
            file=None,
            ):
    df = read_data(file.name)
    df = batch_cos_sim(df, model_name)
    path = 'output.csv'
    df.to_csv(path, index=False, encoding='utf-8-sig')
    return df.to_markdown(), path


model_name_input = gr.components.Textbox(
    value='paraphrase-multilingual-MiniLM-L12-v2',
    lines=1,
    type="text"
)

prompt_input = gr.components.Textbox(
    value='prompt,response',
    lines=10,
    type="text"
)

file_output = gr.components.File(label="Output File",
                                 file_count="single",
                                 file_types=["", ".", ".csv", ".xls", ".xlsx"])

app = gr.Interface(
    fn=process,
    inputs=[model_name_input, prompt_input, "file" ],
    outputs=["text", file_output]
)
app.launch()