import streamlit as st
import sparknlp

from sparknlp.base import *
from sparknlp.annotator import *
from pyspark.ml import Pipeline

# Page configuration
st.set_page_config(
    layout="wide", 
    initial_sidebar_state="auto"
)

# CSS for styling
st.markdown("""
    <style>
        .main-title {
            font-size: 36px;
            color: #4A90E2;
            font-weight: bold;
            text-align: center;
        }
        .section {
            background-color: #f9f9f9;
            padding: 10px;
            border-radius: 10px;
            margin-top: 10px;
        }
        .section p, .section ul {
            color: #666666;
        }
    </style>
""", unsafe_allow_html=True)

@st.cache_resource
def init_spark():
    return sparknlp.start()

@st.cache_resource
def create_pipeline(model):
    documentAssembler = DocumentAssembler() \
        .setInputCol("text") \
        .setOutputCol("documents")

    t5 = T5Transformer.pretrained(model) \
        .setTask("translate English to SQL:") \
        .setInputCols(["documents"]) \
        .setMaxOutputLength(200) \
        .setOutputCol("sql")
    
    pipeline = Pipeline().setStages([documentAssembler, t5])
    return pipeline

def fit_data(pipeline, data):
    df = spark.createDataFrame([[data]]).toDF("text")
    result = pipeline.fit(df).transform(df)
    return result.select('sql.result').collect()

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["t5_small_wikiSQL"],
    help="For more info about the models visit: https://sparknlp.org/models"
)

# Set up the page layout
title, sub_title = (
    'SQL Query Generation',
    'This demo shows how to generate SQL code from natural language text.'
)

st.markdown(f'<div class="main-title">{title}</div>', unsafe_allow_html=True)
st.write(sub_title)

# Reference notebook link in sidebar
link = """
<a href="https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/streamlit_notebooks/T5_SQL.ipynb">
    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>
</a>
"""
st.sidebar.markdown('Reference notebook:')
st.sidebar.markdown(link, unsafe_allow_html=True)

# Load examples
examples = [
    "How many customers have ordered more than 2 items?",
    "How many players were with the school or club team La Salle?",
    "When the scoring rank was 117, what was the best finish?",
    "When the best finish was T69, how many people came in 2nd?",
    "How many wins were there when the money list rank was 183?",
    "When did the Metrostars have their first Rookie of the Year winner?",
    "What college did the Rookie of the Year from the Columbus Crew attend?"
]

selected_text = st.selectbox("Select an example", examples)
custom_input = st.text_input("Try it with your own Sentence!")

text_to_analyze = custom_input if custom_input else selected_text

st.write('Text to be converted to SQL query:')
HTML_WRAPPER = """<div class="scroll entities" style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem; white-space:pre-wrap">{}</div>"""
st.markdown(HTML_WRAPPER.format(text_to_analyze), unsafe_allow_html=True)

# Initialize Spark and create pipeline
spark = init_spark()
pipeline = create_pipeline(model)
output = fit_data(pipeline, text_to_analyze)

# Display matched sentence
st.write("Generated Output:")

output_text = "".join(output[0][0])
st.markdown(f'<div class="section-content">{output_text}</div>', unsafe_allow_html=True)