import sqlparse
import typer
from transformers import pipeline
from groq import Groq
import json
app = typer.Typer()
from dotenv import load_dotenv
import os
load_dotenv()
import streamlit as st
import gradio as gr
client = Groq(api_key=os.getenv("GROQ_API_KEY"))
# print("API Key Loaded:", os.getenv("GROQ_API_KEY"))
# st.title("🚀 SQL Copilot ")


#client = Groq(api_key=GROQ_API_KEY)

from datetime import datetime

# ---------------------------
# Tools (Agents)
# ---------------------------
def format_agent(query: str = typer.Option(..., "--query", "-q", help="The SQL query to format")):
    """Formatter agent 📝"""
    return sqlparse.format(query, reindent=True, keyword_case="upper")

def explainer_agent(query: str = typer.Option(..., "--query", "-q", help="The SQL query to format")):
    """AI-powered SQL explanation"""
    prompt = f"""
    You are an expert SQL analyst. 
    Analyze this SQL query and return:
    1. Business logic in plain English point wise in client or business person who is not technical understandable format.
    2. Tables used
    3. Joins used
    4. Filters/conditions
     
    

    SQL Query:
    {query}
    """
    completion = client.chat.completions.create(
        model="llama-3.3-70b-versatile",  
        messages=[{"role": "user", "content": prompt}],
    )
    
    return (completion.choices[0].message.content)

AGENTS = {
    "Formatter 📝": format_agent,
    "Explainer 🤖": explainer_agent
  
}
AGENT_COLORS = {
    "Formatter 📝": "#ADD8E6",  # light blue
    "Explainer 🤖": "#90EE90" # light green

}
# ---------------------------
# Process Queries
# ---------------------------
def process_queries(query_text, uploaded_file):
    queries = []

    if uploaded_file is not None:
        with open(uploaded_file.name, "r", encoding="utf-8") as f:
            content = f.read()
        #content = uploaded_file.read().decode("utf-8")
        queries = [q.strip() for q in sqlparse.split(content) if q.strip()]
    elif query_text and query_text.strip():
        queries = [query_text.strip()]

    if not queries:
        return [], None, None

    results = []
    formatted_list = []
    rules_list = []

    for q in queries:
        formatted_q = format_agent(q)
        rules_q = explainer_agent(q)

        results.append(["Formatter 📝", formatted_q])
        results.append(["Explainer 🤖", rules_q])

        formatted_list.append(formatted_q)
        rules_list.append(f"Query:\n{q}\n\nBusiness Rules:\n{rules_q}\n{'-'*50}\n")

    # Save files for download
    formatted_file = "formatted_queries.sql"
    rules_file = "business_rules.txt"

    with open(formatted_file, "w", encoding="utf-8") as f:
        f.write("\n\n".join(formatted_list))

    with open(rules_file, "w", encoding="utf-8") as f:
        f.write("\n".join(rules_list))

    return results, formatted_file, rules_file

# ---------------------------
# Gradio UI
# ---------------------------
with gr.Blocks(title="🚀 SQL Copilot Agent") as demo:
    gr.Markdown("## 🚀 SQL Copilot Agent")

    with gr.Tab("Single Query / File"):
        query_box = gr.Textbox(label="Paste your SQL query here", lines=6)
        file_box = gr.File(label="Or upload SQL file", file_types=[".sql"])
        process_btn = gr.Button("Process")

        output_display = gr.Dataframe(headers=["Agent", "Output"], interactive=False)
        download_formatted = gr.File(label="Download Formatted SQL")
        download_rules = gr.File(label="Download Business Rules")

        process_btn.click(
            fn=process_queries,
            inputs=[query_box, file_box],
            outputs=[output_display, download_formatted, download_rules],
        )

demo.launch()