import gradio as gr
import models

with gr.Blocks() as demo:
    Models: list[models.BaseTCOModel] = [models.OpenAIModel, models.OpenSourceLlama2Model]
    model_names = [Model().get_name() for Model in Models]
    with gr.Row():
        with gr.Column():
            page1 = models.ModelPage(Models)
            dropdown = gr.Dropdown(model_names, interactive=True)
            page1.render()

        with gr.Column():
            page2 = models.ModelPage(Models)
            dropdown2 = gr.Dropdown(model_names, interactive=True)
            page2.render()
            
    dropdown.change(page1.make_model_visible, inputs=dropdown, outputs=page1.get_all_components())
    dropdown2.change(page2.make_model_visible, inputs=dropdown2, outputs=page2.get_all_components())

    compute_tco_btn = gr.Button("Compute TCO")
    tco_output = gr.Text("Output: ")
    compute_tco_btn.click(page1.compute_cost_per_token, inputs=page1.get_all_components_for_cost_computing() + [dropdown], outputs=tco_output)

demo.launch(debug=True)