from h2o_wave import main, app, Q, ui, data
from gradio_client import Client
import ast


async def init_ui(q: Q) -> None:
    q.page['meta'] = ui.meta_card(
        box='',
        layouts=[
            ui.layout(breakpoint='xs', min_height='100vh', zones=[
                ui.zone('main', size='1', direction=ui.ZoneDirection.ROW, zones=[
                    ui.zone('sidebar', size='250px'),
                    ui.zone('body', direction=ui.ZoneDirection.COLUMN, zones=[
                        ui.zone('title', size='55px'),
                        ui.zone('content', size='1'),
                        ui.zone('footer'),
                    ]),
                ])
            ])
        ],
        title='H2O GPT',
    )
    q.page['sidebar'] = ui.nav_card(
        box='sidebar', color='primary', title='GPT Wave app', subtitle='Powered by H2O GPT',
        value=f"#{q.args['#']}' if q.args['#'] else '#page1",
        image='https://wave.h2o.ai/img/h2o-logo.svg', items=[
            ui.nav_group('', items=[
                ui.nav_item(name='wave-docs', label='Wave docs', path='https://wave.h2o.ai/'),
                ui.nav_item(name='h2o-gpt', label='H2O GPT', path='https://github.com/h2oai/h2ogpt'),
                ui.nav_item(name='fine-tune', label='LLM Studio', path='https://github.com/h2oai/h2o-llmstudio'),
                ui.nav_item(name='more-models', label='More models', path='https://huggingface.co/h2oai'),
            ]),
        ],
        secondary_items=[
            ui.toggle(name='dark_mode', label='Dark mode', trigger=True),
            ui.text('<center>Made with H2O Wave.</center>')
        ]
    )

    q.page['chatbot'] = ui.chatbot_card(
        box=ui.box('content'),
        data=data('content from_user', t='list'),
        name='chatbot'
    )
    q.page['title'] = ui.section_card(
        box='title',
        title='',
        subtitle='',
        items=[
            ui.dropdown(name='model', trigger=True, label='', value='gpt', choices=[
                ui.choice(name='gpt', label='H2O GPT'),
                ui.choice(name='falcon', label='h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v3'),
                ui.choice(name='llama', label='h2oai/h2ogpt-research-oasst1-llama-65b'),
                ui.choice(name='mpt', label='mosaicml/mpt-30b-instruct'),
            ]),
            ui.button(name='clear', label='Clear', icon='Delete'),
        ],
    )


@app('/')
async def serve(q: Q):
    if not q.client.initialized:
        await init_ui(q)
        q.client.model_client = Client('https://gpt.h2o.ai/')
        q.client.initialized = True

    # A new message arrived.
    if q.args.chatbot:
        # Append user message.
        q.page['chatbot'].data += [q.args.chatbot, True]
        # Append bot response.
        kwargs = dict(instruction_nochat=q.args.chatbot)
        try:
            res = q.client.model_client.predict(str(dict(kwargs)), api_name='/submit_nochat_api')
            bot_res = ast.literal_eval(res)['response']
            q.page['chatbot'].data += [bot_res, False]
        except:
            q.page['meta'] = ui.meta_card(box='', notification_bar=ui.notification_bar(
                text='An error occurred during prediction. Please try later or a different model.',
                type='error',
            ))
    elif q.args.clear:
        # Recreate the card.
        q.page['chatbot'] = ui.chatbot_card(
            box=ui.box('content'),
            data=data('content from_user', t='list'),
            name='chatbot'
        )
    elif q.args.dark_mode is not None:
        q.page['meta'].theme = 'h2o-dark' if q.args.dark_mode else 'light' 
        q.page['sidebar'].color = 'card' if q.args.dark_mode else 'primary'
    elif q.args.model:
        try:
            q.client.model_client = Client(f'https://{q.args.model}.h2o.ai/')
            q.page['meta'] = ui.meta_card(box='', notification_bar=ui.notification_bar(
                text='Model changed successfully.',
                type='success',
            ))
        except:
            q.page['meta'] = ui.meta_card(box='', notification_bar=ui.notification_bar(
                text='An error occurred while changing the model. Please try a different one.',
                type='error',
            ))

    await q.page.save()