import json
import os
import urllib.parse

import gradio as gr
import requests
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from huggingface_hub import InferenceClient

example = HuggingfaceHubSearch().example_value()

client = InferenceClient(
    "meta-llama/Meta-Llama-3.1-70B-Instruct",
    token=os.environ["HF_TOKEN"],
)


def get_iframe(hub_repo_id, sql_query=None):
    if not hub_repo_id:
        raise ValueError("Hub repo id is required")
    if sql_query:
        sql_query = urllib.parse.quote(sql_query)
        url = f"https://huggingface.co./datasets/{hub_repo_id}/embed/viewer?sql_console=true&sql={sql_query}"
    else:
        url = f"https://huggingface.co./datasets/{hub_repo_id}/embed/viewer"
    iframe = f"""
    <iframe
  src="{url}"
  frameborder="0"
  width="100%"
  height="800px"
></iframe>
"""
    return iframe


def get_column_info(hub_repo_id):
    url: str = f"https://datasets-server.huggingface.co/info?dataset={hub_repo_id}"
    response = requests.get(url)
    try:
        data = response.json()
        data = data.get("dataset_info")
        key = list(data.keys())[0]
        features: str = json.dumps(data.get(key).get("features"))
    except Exception as e:
        gr.Error(f"Error getting column info: {e}")
    return features


def query_dataset(hub_repo_id, features, query):
    messages = [
        {
            "role": "system",
            "content": "You are a SQL query expert assistant that returns a DuckDB SQL queries based on the user's natural language query and dataset features. You might need to use DuckDB functions for lists and aggregations, given the features. Only return the SQL query, no other text.",
        },
        {
            "role": "user",
            "content": f"""table train
# Features
{features}

# Query
{query}
""",
        },
    ]
    response = client.chat_completion(
        messages=messages,
        max_tokens=1000,
        stream=False,
    )
    query = response.choices[0].message.content
    return query, get_iframe(hub_repo_id, query)


with gr.Blocks() as demo:
    gr.Markdown("""# 🐥 🦙 🤗 Text To SQL Hub Datasets 🤗 🦙 🐥

                This is a basic text to SQL tool that allows you to query datasets on Huggingface Hub.
                It is built with [DuckDB](https://duckdb.org/), [Huggingface's Inference API](https://huggingface.co./docs/api-inference/index), and [LLama 3.1 70B](https://huggingface.co./meta-llama/Meta-Llama-3.1-70B-Instruct).
                Also, it uses the [dataset-server API](https://redocly.github.io/redoc/?url=https://datasets-server.huggingface.co/openapi.json#operation/isValidDataset).
                """)
    with gr.Row():
        with gr.Column():
            search_in = HuggingfaceHubSearch(
                label="Search Huggingface Hub",
                placeholder="Search for models on Huggingface",
                search_type="dataset",
                sumbit_on_select=True,
            )
            query = gr.Textbox(
                label="Natural Language Query",
                placeholder="Enter a natural language query to generate SQL",
            )
            sql_out = gr.Code(
                label="SQL Query",
                interactive=True,
                language="sql",
                lines=1,
                visible=False,
            )
    with gr.Row():
        with gr.Column():
            btn = gr.Button("Show Dataset")
        with gr.Column():
            btn2 = gr.Button("Query Dataset")
    with gr.Row():
        search_out = gr.HTML(label="Search Results")
    with gr.Row():
        features = gr.Code(label="Features", language="json", visible=False)
    gr.on(
        [btn.click, search_in.submit],
        fn=get_iframe,
        inputs=[search_in],
        outputs=[search_out],
    ).then(
        fn=get_column_info,
        inputs=[search_in],
        outputs=[features],
    )
    gr.on(
        [btn2.click, query.submit],
        fn=query_dataset,
        inputs=[search_in, features, query],
        outputs=[sql_out, search_out],
    )

if __name__ == "__main__":
    demo.launch()