import dataiku
from dataiku.langchain.dku_llm import DKUChatLLM
import json
import sqlparse
import re
import pandas as pd

import dash
from dash import dcc, html, dash_table, ctx, ALL
from dash.dependencies import Input, Output, State
import dash_bootstrap_components as dbc

LLM_ID = "agent:MSIGB581:v1"
llm_provided = len(dataiku.get_custom_variables()["LLM_id"]) > 0
if llm_provided:
    llm = dataiku.api_client().get_default_project().get_llm(LLM_ID)

DATASETS = ["match_results", "match_goalscorers"]

######################################
# Webapp configs
######################################
# App Title
app.title = "Text2SQL web app"
app.config.external_stylesheets = [
    dbc.themes.ZEPHYR,
    dbc.icons.BOOTSTRAP
]

######################################
# Layout of the webapp
######################################
send_icon = html.Span(html.I(className="bi bi-send"))
# Creating an input group using Dash Bootstrap Components
question_bar = dbc.InputGroup(
    [
        # Adding a textarea input for entering a query
        dbc.Textarea(id='query', 
                     placeholder="Enter a query you want answered on available data", 
                     minLength=0,
                    style={"border-color": "lightgrey"}),
        
        # Adding a button with the send icon
        dbc.Button(send_icon, id='send-btn', title='Get an answer')
    ],
    style={"margin-top": "20px", 'display': 'flex'}
)

table_list = ", ".join([f"`{dataset}`" for dataset in DATASETS])

sample_questions = [
    "Who are the top 3 goal scorers?",
    "In which game did Eden Hazard last scored in? And what tournament was it in?",
    "How many unique scorers are there?"
]


# Define App Layout
app.layout = html.Div(
    [
        # Title at the Top of Webapp
        html.H4(
            "Text2SQL web app",
            style={"margin-top": "20px", "text-align": "center"}
        ),
        
        # Describe the Webapp
        html.Div(dcc.Markdown(f"Interact with the following datasets: {table_list}"),
            style={"margin-top": "20px"}  # Add margin at the top
        ),
        
        dbc.Alert(
            [
                html.I(className="bi bi-info-circle-fill me-2"),
                "Kindly note that this web app uses an LLM to generate the responses to the questions. Its answers may be inaccurate.",
            ],
            color="primary",
            className="d-flex align-items-center",
            style={"font-size": 14}
        ),
                
        # Text bar to enter question to be answered
        question_bar,   
                        
        # Components to store answers and show answers/feedback from LLM, including output of SQL code, SQL code and time taken for query & saving as recipe
        html.Div(
            [
                dbc.Spinner([
                        dcc.Markdown(
                            id='answer',
                            link_target="_blank"
                        ),
                    
                    ],
                    color="primary",
                    spinner_style={"display": "none"}
                )
            ],
            style={"margin-top": "20px","align-items": "flex-start","display": "flex","height": "auto"}
        ),
        
        dbc.Spinner(
            [
                dcc.Markdown(
                    id='intermediate_steps',
                    link_target="_blank",
                    style={"white-space": "pre-line"}
                ),
                # DataTable to show results of the SQL Query
                dash_table.DataTable(
                    data=[{}],
                    columns=[{}],
                    style_table={'overflowX': 'auto', "margin-top": "5px", "margin-bottom": "20px", "display": "none", "font-size": 12},  # Set table style
                    id='query_table',
                    export_format=""

                ),
                dcc.Markdown(
                    id='execution_time',
                    link_target="_blank",
                    style={"white-space": "pre-line"}
                )
            ],
            color="primary"
        ),
        html.Div(
                [
                    dbc.Button(
                        [html.I(className="bi bi-chat-quote"), f" {sample_question}"], 
                        id={"index": i, "type": "sample-question"},
                        outline=True,
                        color="primary",
                        style={"margin": "20px 5px"}
                    )
                    for (i, sample_question) in enumerate(sample_questions)
                ],
                style={'display': 'flex', 'alignItems': 'center'},
        )
    ],
    style={
        "margin": "auto",
        "text-align": "left",
        "max-width": "800px"
    }
)

######################################
# Callbacks
######################################

# There are two callbacks :
# - One for updating the textfield with the question when a sample question is selected
# - One for the execution of text2SQL pipeline once the button is clicked

# Callback to ask LLM to generate SQL to answer question about data selected and execute the SQL statement to return result
@app.callback(
    [Output('answer', 'children'),
    Output('query_table', 'data'),
    Output('query_table', 'columns'),
    Output('query_table', 'style_table'),
    Output('intermediate_steps', 'children'),
    Output('execution_time', 'children')],
    [Input('send-btn', 'n_clicks'),
    Input('query', 'n_submit')],
    [State('query', 'value'),
    State('query_table', 'style_table')],
    prevent_initial_call=True
)
def answer_question(n_clicks, n_submit, query, query_table_style):
    """
    Display the answer
    """
    
    # Don't do anything if no query
    if len(query) == 0:
        query_table_style['display'] = 'none'
        return "Please enter your question in the text field above", [{}], [{}], query_table_style, "", ""

    # Return message to set up llm agent using openai key in user secrets
    if not llm_provided:
        query_table_style['display'] = 'none'
        return "You need an LLM connection (cf. requirements in the wiki of this project).", [{}], [{}], query_table_style, "", ""

    # invoking the agent chain
    completion = llm.new_completion()
    completion.with_message(
        json.dumps(
            {
                "request": query,
            }
        )
    )
    resp = completion.execute()
    result = json.loads(resp.text)
    
    reply = result["reply"]
    sql_query = result["sql_query"]
    answer_details = result["answer_details"]
    processing_time = result["processing_time"]

    m = re.match("[0-9]\. execute_sql.*--> (\{.*\})", answer_details.split("\n")[-1])
    if m:
        answer_data = eval(json.loads(m.group(1))["result"])
    else:
        answer_data = [{}]

    answer = f'''** Answer**: {reply}'''
    execution_time = f"**Execution time**: {processing_time} seconds"
    
    if len(sql_query) > 0:
        answer_details = f'''
**SQL query generated by the LLM**:
```sql
{sqlparse.format(sql_query, reindent=True, keyword_case='upper')}
```
** Result of the SQL query**:
'''         
        answer_data_columns = [{"name": i, "id": i} for i in pd.DataFrame(answer_data).columns]
        query_table_style['display'] = 'flex' 
    else:
        execution_time = f"**Execution time**: {(time.time()-start):.1f} seconds"
        answer_data_columns = [{}]
        answer_details = ""
        query_table_style['display'] = 'none'
        
    return answer, answer_data, answer_data_columns, query_table_style, answer_details, execution_time
    
# Callback to enter the sample question into the text area
@app.callback(
    Output('query', 'value'),
    Input({'type': 'sample-question', 'index': ALL}, 'n_clicks_timestamp')
)
def select_sample_question(values):
    """
    Update text area field to use sample question selected
    """
    max_value, argmax = -1, None
    for i in range(len(values)):
        if values[i] is not None:
            if values[i] > max_value:
                max_value, argmax = values[i], i
    if argmax is None:
        return ""
    idx = int(dash.callback_context.inputs_list[-1][argmax]['id']['index'])
    return sample_questions[idx]