from dash import html, dcc, no_update, set_props
import dash_bootstrap_components as dbc
from dash.dependencies import Input, Output, State

import dataiku
from dataiku import SQLExecutor2
from duckduckgo_search import DDGS
import json

dbc_css = "https://cdn.jsdelivr.net/gh/AnnMarieW/dash-bootstrap-templates/dbc.min.css"
app.config.external_stylesheets = [dbc.themes.SUPERHERO, dbc_css]

# LLM setup
PROJECT = ""  # DSS project key goes here
LLM_ID = ""  # LLM ID for the LLM Mesh connection + model goes here
client = dataiku.api_client()
project = client.get_project(PROJECT)
llm = project.get_llm(LLM_ID)

def create_chat_session():
    """Create a new chat session with tools and prompt"""
    chat = llm.new_completion()
    
    # Import tools from JSON file
    library = project.get_library()
    tools_file = library.get_file("/python/tools.json")
    tools_str = tools_file.read()
    tools_dict = json.loads(tools_str)
    chat.settings["tools"] = tools_dict["tools"]
    
    SYSTEM_PROMPT = """You are a customer information assistant. You can:
    1. Look up customer information in our database
    2. Search for additional information about companies online
    
    When looking up customer info, you will:
    1. First get customer details
    2. Then look up their company
    3. Summarize in a single paragraph with details on both"""
    
    chat.with_message(SYSTEM_PROMPT, role="system")
    return chat

def get_customer_details(customer_id):
    """Get customer information from database"""
    dataset = dataiku.Dataset("pro_customers_sql")
    table_name = dataset.get_location_info().get('info', {}).get('table')
    executor = SQLExecutor2(dataset=dataset)
    customer_id = customer_id.replace("'", "\\'")
    query_reader = executor.query_to_iter(
        f"""SELECT name, job, company FROM "{table_name}" WHERE id = '{customer_id}'""")
    for (name, job, company) in query_reader.iter_tuples():
        return f"The customer's name is \"{name}\", holding the position \"{job}\" at the company named {company}"
    return f"No information can be found about the customer {customer_id}"

def search_company_info(company_name):
    """Search for company information online"""
    with DDGS() as ddgs:
        results = list(ddgs.text(f"{company_name} (company)", max_results=1))
        if results:
            return f"Information found about {company_name}: {results[0]['body']}"
        return f"No information found about {company_name}"

def process_tool_calls(tool_calls):
    """Process tool calls and return results"""
    tool_name = tool_calls[0]["function"]["name"]
    llm_args = json.loads(tool_calls[0]["function"]["arguments"])
    if tool_name == "get_customer_info":
        return get_customer_details(llm_args["customer_id"])
    elif tool_name == "get_company_info":
        return search_company_info(llm_args["company_name"])

# Dash app layout
app.layout = html.Div([
    dbc.Row([html.H2("Using LLM Mesh with an agent in Dash")]),
    dbc.Row(dbc.Label("Please enter the ID of the customer:")),
    dbc.Row([
        dbc.Col(dbc.Input(id="customer_id", placeholder="Customer Id"), width=10),
        dbc.Col(dbc.Button("Search", id="search", color="primary"), width="auto")
    ], justify="between"),
    dbc.Row([dbc.Col(dbc.Textarea(id="result", style={"min-height": "500px"}), width=12)]),
    dbc.Toast(
        [html.P("Searching for information about the customer", className="mb-0"),
         dbc.Spinner(color="primary")],
        id="auto-toast",
        header="Agent working",
        icon="primary",
        is_open=False,
        style={"position": "fixed", "top": "50%", "left": "50%", "transform": "translate(-50%, -50%)"},
    ),
    dcc.Store(id="chat-state"),
    dcc.Store(id="step", data={"current_step": 0}),
], className="container-fluid mt-3")

@app.callback(
    [Output("result", "value"), Output("chat-state", "data")],
    Input("search", "n_clicks"),
    [State("customer_id", "value"), State("chat-state", "data")],
    prevent_initial_call=True,
    running=[(Output("auto-toast", "is_open"), True, False),
             (Output("search", "disabled"), True, False)]
)

def update_output(n_clicks, customer_id, chat_state):
    if not customer_id:
        return no_update, no_update
    
    # Create new chat session
    chat = create_chat_session()
    
    # Start conversation about customer
    content = f"Tell me about the customer with ID {customer_id}"
    chat.with_message(content, role="user")
    
    conversation_log = []
    while True:
        response = chat.execute()
        
        if not response.tool_calls:
            # Final answer received
            chat.with_message(response.text, role="assistant")
            conversation_log.append(f"Final Answer: {response.text}")
            break
            
        # Handle tool calls
        chat.with_tool_calls(response.tool_calls, role="assistant")
        tool_call_result = process_tool_calls(response.tool_calls)
        chat.with_tool_output(tool_call_result, tool_call_id=response.tool_calls[0]["id"])
        
        # Log the step
        tool_name = response.tool_calls[0]["function"]["name"]
        tool_args = response.tool_calls[0]["function"]["arguments"]
        conversation_log.append(f"Tool: {tool_name}\nInput: {tool_args}\nResult: {tool_call_result}\n{'-'*50}")
    
    return "\n".join(conversation_log), {"messages": chat.cq["messages"]}