import marimo

__generated_with = "0.11.17"
app = marimo.App(width="full")


@app.cell
def _():
    import marimo as mo
    import pandas as pd
    import os
    import io
    import json
    return io, json, mo, os, pd


@app.cell
def _(mo):
    ### Credentials for the watsonx.ai SDK clien
    # Endpoints
    wx_platform_url = "https://api.dataplatform.cloud.ibm.com"
    regions = {
        "US": "https://us-south.ml.cloud.ibm.com",
        "EU": "https://eu-de.ml.cloud.ibm.com",
        "GB": "https://eu-gb.ml.cloud.ibm.com",
        "JP": "https://jp-tok.ml.cloud.ibm.com",
        "AU": "https://au-syd.ml.cloud.ibm.com",
        "CA": "https://ca-tor.ml.cloud.ibm.com"
    }

    # Create a form with multiple elements
    client_instantiation_form = (
        mo.md('''
        ###**watsonx.ai credentials:**

        {wx_region}

        {wx_api_key}

        {project_id}

    ''')
        .batch(
            wx_region = mo.ui.dropdown(regions, label="Select your watsonx.ai region:", value="US", searchable=True),
            wx_api_key = mo.ui.text(placeholder="Add your IBM Cloud api-key...", label="IBM Cloud Api-key:", kind="password"),
            project_id = mo.ui.text(placeholder="Add your watsonx.ai project_id...", label="Project_ID:", kind="text"),
            space_id = mo.ui.text(placeholder="Add your watsonx.ai space_id...", label="Space_ID:", kind="text")
        ,)
        .form(show_clear_button=True, bordered=False)
    )

    return client_instantiation_form, regions, wx_platform_url


@app.cell
def _(client_instantiation_form, mo):
    from ibm_watsonx_ai import APIClient, Credentials
    import requests

    def get_iam_token(api_key):
        return requests.post(
            'https://iam.cloud.ibm.com/identity/token',
            headers={'Content-Type': 'application/x-www-form-urlencoded'},
            data={'grant_type': 'urn:ibm:params:oauth:grant-type:apikey', 'apikey': api_key}
        ).json()['access_token']

    def setup_task_credentials(project_client):
        # Get existing task credentials
        existing_credentials = project_client.task_credentials.get_details()

        # Delete existing credentials if any
        if "resources" in existing_credentials and existing_credentials["resources"]:
            for cred in existing_credentials["resources"]:
                cred_id = project_client.task_credentials.get_id(cred)
                project_client.task_credentials.delete(cred_id)

        # Store new credentials
        return project_client.task_credentials.store()

    if client_instantiation_form.value:
    ### Instantiate the watsonx.ai client
        wx_api_key = client_instantiation_form.value["wx_api_key"]

        wx_credentials = Credentials(
            url=client_instantiation_form.value["wx_region"],
            api_key=wx_api_key
        )

        project_client = APIClient(credentials=wx_credentials, project_id=client_instantiation_form.value["project_id"])
        # deployment_client = APIClient(credentials=wx_credentials, space_id=client_instantiation_form.value["space_id"])

        token = get_iam_token(wx_api_key)
        task_credentials_details = setup_task_credentials(project_client)

    else:
        project_client = None
        # deployment_client = None
        task_credentials_details = None
        wx_api_key = None
        token = None

    client_status = mo.md("### Client Instantiation Status will turn Green When Ready")

    if project_client is not None:
        client_callout_kind = "success"
    else:
        client_callout_kind = "neutral"
    return (
        APIClient,
        Credentials,
        client_callout_kind,
        client_status,
        get_iam_token,
        project_client,
        requests,
        setup_task_credentials,
        task_credentials_details,
        token,
        wx_api_key,
        wx_credentials,
    )


@app.cell
def _(client_callout_kind, client_instantiation_form, client_status, mo):
    client_callout = mo.callout(client_status, kind=client_callout_kind)

    client_stack = mo.hstack([client_instantiation_form, client_callout], align="center", justify="space-around")

    client_stack
    return client_callout, client_stack


@app.cell
def _(mo, project_client):
    if project_client:
        # model_specs = project_client.foundation_models.get_chat_model_specs()  ### if you want models that support chat_completions via the RestAPI or Python SDK.
        model_specs = project_client.foundation_models.get_model_specs()
        resources = model_specs["resources"]
        model_id_list = []
        for resource in resources:
            model_id_list.append(resource["model_id"])
        model_selection = mo.ui.table(
            model_id_list,
            selection="single",  # Only allow selecting one row
            label="Select a model to use.",
            page_size=30,
            initial_selection=[16]
        )
    else:
        model_specs = []
        resources = []
        model_id_list = []
        model_selection = mo.md("**Instantiate a watsonx.ai client with a valid project_id**")

    return model_id_list, model_selection, model_specs, resource, resources


@app.cell
def _(mo, model_selection):
    from ibm_watsonx_ai.foundation_models import ModelInference
    from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
    # Create a form with multiple elements
    llm_parameters = (
        mo.md('''
        ###**LLM parameters:**

        {decoding_method}

        {repetition_penalty}

        {min_tokens}

        {max_tokens}

        {stop_sequences}
    ''')
        .batch(
            ### Temporary version with preset credentials
            decoding_method = mo.ui.dropdown(options=["greedy", "sample"], value="greedy",label="Decoding Method:"),
            min_tokens = mo.ui.number(start=1, stop=1, label="Minimum Output Tokens:"),
            max_tokens = mo.ui.number(start=1, stop=8096, value=500, label="Maximum Output Tokens:"),
            repetition_penalty = mo.ui.number(start=1.0, stop=2.0, step=0.01, label="Repetition Penalty:"),
            stop_sequences = mo.ui.text(label="Stopping Sequences:", value="['<|end_of_text|>','</s>']", placeholder="List of Strings, e.g. ['<|end_of_text|>','</s>']", full_width=False)
        ).form(show_clear_button=True, bordered=True)
    )

    # llm_parameters
    llm_setup = mo.hstack([model_selection, llm_parameters], align="center", justify="space-around")

    llm_setup
    return GenParams, ModelInference, llm_parameters, llm_setup


@app.cell
def _(
    GenParams,
    ModelInference,
    llm_parameters,
    model_selection,
    project_client,
):
    import ast

    if llm_parameters.value:
        params = {
            GenParams.DECODING_METHOD: llm_parameters.value['decoding_method'],
            GenParams.MAX_NEW_TOKENS: llm_parameters.value['max_tokens'],
            GenParams.MIN_NEW_TOKENS: llm_parameters.value['min_tokens'],
            GenParams.REPETITION_PENALTY: llm_parameters.value['repetition_penalty'],
            GenParams.STOP_SEQUENCES: ast.literal_eval(llm_parameters.value['stop_sequences']),
            GenParams.RETURN_OPTIONS: {
                'input_text': False,
                'generated_tokens': False,
                'input_tokens': True,
                'token_logprobs': False
            }
        }

        inf_model = ModelInference(api_client=project_client, model_id=model_selection.value[0], params=params)
    else:
        params = {}
        inf_model = None
    return ast, inf_model, params


@app.cell
def _(mo):
    prompt_template_mistral = """[INST] write your prompt here [/INST]"""

    prompt_template_llama = """<|start_header_id|>user<|end_header_id|>\n\n write your prompt here <|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"""


    templates = {
        "mistral": prompt_template_mistral,
        "llama": prompt_template_llama
    }

    template_selector = mo.ui.dropdown(templates, value="mistral", label="Select Prompt Template with Syntax:")
    return (
        prompt_template_llama,
        prompt_template_mistral,
        template_selector,
        templates,
    )


@app.cell
def _(mo, template_selector):
    prompt_editor = (
        mo.md('''
        #### **Create your prompt here by editing the template:**

        {editor}

        ''')
        .batch(
            editor = mo.ui.code_editor(value=template_selector.value, language="python", min_height=50)
        )
        .form(show_clear_button=True, bordered=False)
    )
    return (prompt_editor,)


@app.cell
def _(mo, prompt_editor, template_selector):
    prompt_setup = mo.vstack([template_selector, prompt_editor])
    prompt_setup
    return (prompt_setup,)


@app.cell
def _(mo, prompt_editor):
    _ = prompt_editor.value
    
    if prompt_editor.value is not None:
        prompt_printout = mo.md(f"**Current prompt template:**\n\n {prompt_editor.value['editor']}")
    else:
        prompt_printout = mo.md(f"**Submit a prompt template.**")

    prompt_printout
    return (prompt_printout,)


@app.cell
def _(inf_model, params, process_with_llm, prompt_editor):
    _ = prompt_editor.value
    
    if prompt_editor.value is not None:
        prompt_response = process_with_llm(inf_model, prompt_template=prompt_editor.value['editor'], params=params, return_full_json_response=False)
        prompt_response_full = process_with_llm(inf_model, prompt_template=prompt_editor.value['editor'], params=params, return_full_json_response=True)
    else:
        prompt_response = []
        prompt_response_full = []
    return prompt_response, prompt_response_full


@app.cell
def _(mo, prompt_response, prompt_response_full):
    mo.vstack([mo.md(f"{prompt_response}"),prompt_response_full], align="center",justify="space-around")
    return


@app.cell
def _():
    def process_with_llm(inf_model, prompt_template, params, return_full_json_response=False):
        """
        Process a prompt with an LLM model.
    
        Returns full JSON response or just text based on return_full_json_response parameter.
        """
        # Check for required model
        if not inf_model:
            print("Missing required inference model")
            return None
    
        # Extract prompt value if it's a dict
        if hasattr(prompt_template, 'get') and prompt_template.get('value'):
            prompt_template = prompt_template['value']
    
        try:
            # Call appropriate method based on return type preference
            if return_full_json_response:
                return inf_model.generate(prompt=prompt_template, params=params)
            else:
                return inf_model.generate_text(prompt=prompt_template, params=params)
            
        except Exception as e:
            print(f"Error during inference: {str(e)}")
            return None
    return (process_with_llm,)


if __name__ == "__main__":
    app.run()