import ast
import gradio as gr
from functions import example_question_generator, chatbot_func
from data_sources import connect_sql_db
from utils import message_dict

with gr.Blocks() as demo:
    with gr.Accordion("ℹ️  About the SQL Connector", open=False):
        gr.HTML("""
            <div class="max-w-4xl mx-auto text-center">
                <div class="bg-blue-50 border border-blue-200 rounded-lg max-w-2xl mx-auto p-4">
                    <p>Connect to a PostgreSQL database and query it using natural language.</p>
                    <p style="font-weight:bold;">
                        No credentials are retained — they are passed as session variables and disappear when you leave or refresh.
                        Queries run through Pandas <code>read_sql_query</code>, which cannot delete, drop, or insert rows.
                        Use caution connecting production databases to third-party tools.
                    </p>
                    <p>Contact me if you'd like this built for your organization with proper infrastructure and security controls.</p>
                </div>
            </div>
        """)

    gr.HTML("""
        <div style="max-width:560px;margin:8px auto 4px;padding:8px 14px;background:#f0f9ff;
                    border:1px solid #bae6fd;border-radius:8px;text-align:center;">
            <p style="margin:0;font-size:13px;color:#0369a1;">
                <i class="fas fa-flask" style="margin-right:6px;"></i>
                <strong>Demo credentials pre-filled.</strong>
                &nbsp;Replace with your own database details to analyze your own data.
            </p>
        </div>
    """)

    sql_url = gr.Textbox(label="URL", value="virtual-data-analyst-pg.cyetm2yjzppu.us-west-1.rds.amazonaws.com")
    with gr.Row():
        sql_port = gr.Textbox(label="Port", value="5432")
        sql_user = gr.Textbox(label="Username", value="postgres")
        sql_pass = gr.Textbox(label="Password", value="Vda-1988", type="password")
        sql_db_name = gr.Textbox(label="Database Name", value="dvdrental")

    submit = gr.Button(value="Connect", variant="primary")

    @gr.render(inputs=[sql_url, sql_port, sql_user, sql_pass, sql_db_name], triggers=[submit.click])
    def sql_chat(request: gr.Request, url=sql_url.value, sql_port=sql_port.value, sql_user=sql_user.value, sql_pass=sql_pass.value, sql_db_name=sql_db_name.value):
        if request.session_hash not in message_dict:
            message_dict[request.session_hash] = {}
        message_dict[request.session_hash]['sql'] = None
        if url:
            print("SQL APP")
            process_message = process_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, request.session_hash)
            gr.HTML(value=process_message[1], padding=False)
            if process_message[0] == "success":
                if "virtual-data-analyst-pg.cyetm2yjzppu.us-west-1.rds.amazonaws.com" in url:
                    example_questions = [
                                            ["Describe the dataset"],
                                            ["What is the total revenue generated by each store?"],
                                            ["Can you generate and display a bar chart of film category to number of films in that category?"],
                                            ["Can you generate a pie chart showing the top 10 most rented films by revenue?"],
                                            ["Can you generate a line chart of rental revenue over time?"],
                                            ["What is the relationship between film length and rental frequency?"]
                                        ]
                else:
                    try:
                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'sql', sql_db_name, process_message[2], ""))
                        example_questions = [["Describe the dataset"]]
                        for example in generated_examples:
                            example_questions.append([example])
                    except Exception as e:
                        print("SQL QUESTION GENERATION ERROR")
                        print(e)
                        example_questions = [
                                            ["Describe the dataset"],
                                            ["List the tables in the database"],
                                            ["What could this data be used for?"],
                                        ]
                session_hash = gr.Textbox(visible=False, value=request.session_hash)
                db_url = gr.Textbox(visible=False, value=url)
                db_port = gr.Textbox(visible=False, value=sql_port)
                db_user = gr.Textbox(visible=False, value=sql_user)
                db_pass = gr.Textbox(visible=False, value=sql_pass)
                db_name = gr.Textbox(visible=False, value=sql_db_name)
                titles = gr.Textbox(value=process_message[2], interactive=False, label="SQL Tables")
                data_source = gr.Textbox(visible=False, value='sql')
                schema = gr.Textbox(visible=False, value='')
                bot = gr.Chatbot(type='messages', label="SQL DB Chat Window", render_markdown=True, sanitize_html=False, show_label=True, render=False, visible=True, elem_classes="chatbot")
                chat = gr.ChatInterface(
                                    fn=chatbot_func,
                                    type='messages',
                                    chatbot=bot,
                                    title="Chat with your Database",
                                    examples=example_questions,
                                    concurrency_limit=None,
                                    additional_inputs=[session_hash, data_source, titles, schema, db_url, db_port, db_user, db_pass, db_name]
                                    )

    def process_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, session_hash):
        if url:
            process_message = connect_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, session_hash)
        return process_message

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