import gradio as gr
import sqlite3
import traceback
import os
import re
import uuid
from agno.tools import tool
import pandas as pd
from matplotlib import pyplot as plt
import seaborn as sns
import plotly.express as px

# --- DB Functions ---
def init_product_db():
    conn = sqlite3.connect("flipkart_mobiles.db")
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS mobiles (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            brand TEXT,
            color TEXT,
            model TEXT,
            memory TEXT,
            storage TEXT,
            rating REAL,
            selling_price REAL,
            original_price REAL
        )
    ''')
    conn.commit()
    conn.close()

def read_products():
    conn = sqlite3.connect("flipkart_mobiles.db")
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM mobiles")
    rows = cursor.fetchall()
    conn.close()
    return rows

# --- Tool Wrappers ---
DB_PATH = "flipkart_mobiles.db"
TABLE_NAME = "mobiles"

@tool(show_result=True, stop_after_tool_call=True)
def get_columns_info_from_database(columns: str = "*"):
    """
    Database Schema: brand, color, model, memory, storage, rating, selling_price, original_price
    Table: mobiles

    Query the 'mobiles' table selecting specified columns dynamically.
    
    Input:
    - columns: a comma-separated string of column names to select, e.g. "brand, model, rating"
      If "*", selects all columns.
    
    Returns:
    - Formatted string of rows with selected columns.
    """
    if columns.strip() != "*":
        if not re.fullmatch(r"[a-zA-Z0-9_,\s]+", columns):
            return "Invalid columns format."

    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    # Build query string dynamically
    query = f"SELECT {columns} FROM {TABLE_NAME}"

    try:
        cursor.execute(query)
        rows = cursor.fetchall()

        # Get column names from cursor description
        col_names = [desc[0] for desc in cursor.description]

        output_lines = []
        for row in rows:
            row_dict = dict(zip(col_names, row))
            formatted_row = ", ".join(f"{col}: {row_dict[col]}" for col in col_names)
            output_lines.append(formatted_row)

        return "\n".join(output_lines) if output_lines else "No rows found."
    except Exception as e:
        return f"Query error: {str(e)}"
    finally:
        conn.close()

@tool(show_result=True, stop_after_tool_call=True)
def generate_python_code(python_code: str) -> str:
    """
    You are a Python data scientist. Take the table and columns information from the chat history or agent memory.

    Your task is to generate a valid Python script from the following response.
    This table and columns information can be in raw English or structured format from the chat history or agent memory like:
    - user: task - description
    - tabular strings
    - JSON-like text
    - general descriptive statistics

    You must:
    1. Convert the data into a pandas DataFrame (use variable name `df`)
    2. Select an appropriate chart (bar chart, pie chart, line chart, etc.) based on the user's query
    3. Use matplotlib, seaborn, or plotly to plot. Any one of it to create the chart or graph or plot
    4. Save the chart using the variable `image_path` to a PNG file
    5. Return only the Python code — no comments, no markdown

    ### Rules:
    - Do not use `plt.show()` or any GUI renderer
    - Use clear axis labels and title
    - Save the figure using `plt.savefig(image_path)`
    - `df` must be used for all data manipulations
    - You must generate the full Python code block
    - execute that Python code and return the path to the saved image folder.
    - Create an image into the "plots" folder.

    Example code:
        ```python
        import pandas as pd
        import matplotlib.pyplot as plt

        data = [
            {"id": 1, "name": "Alice", "task": "NLP"},
            {"id": 2, "name": "Bob", "task": "Vision"},
            {"id": 3, "name": "Alice", "task": "NLP"}
        ]

        df = pd.DataFrame(data)
        task_counts = df["task"].value_counts()

        plt.figure(figsize=(6, 4))
        task_counts.plot(kind="bar", color="skyblue")
        plt.xlabel("Task")
        plt.ylabel("Count")
        plt.title("Task Distribution")
        plt.savefig(image_path)
        ```
    """
    return python_code

@tool(show_result=True, stop_after_tool_call=True)
def visualization_tool(python_code: str) -> str:
    """ This function is for taking the python code as input from chat history or agent memory and cleaning it accordingly so that it can be executed, then executing it and returning the image path.
    """
    try:
        cleaned_code = re.sub(r"^```(?:python)?|```$", "", python_code.strip(), flags=re.MULTILINE)
        image_path = f"plots/{uuid.uuid4().hex}.png"
        os.makedirs("plots", exist_ok=True)
        exec_context = {
            "pd": pd,
            "plt": plt,
            "sns": sns,
            "px": px,
            "image_path": image_path
        }
        exec(cleaned_code, exec_context)
        return image_path
    except Exception:
        return f"Error executing visualization code:\n{traceback.format_exc()}"

# --- Init DB ---
init_product_db()

# --- Define Toolkit ---
toolkit = [
    get_columns_info_from_database,
    generate_python_code,
    visualization_tool
]

# --- Gradio UI ---
tabbed = gr.TabbedInterface(
    interface_list=[
        gr.Interface(
            fn=get_columns_info_from_database.entrypoint,
            inputs=[
                gr.Textbox(label="Columns (comma separated, or * for all)", value="*")
            ],
            outputs=gr.Textbox(label="Query Result"),
            title="Query Products"
        ),
        gr.Interface(
            fn=generate_python_code.entrypoint,
            inputs=[
                gr.Textbox(label="Python code for Visualization", lines=10)
            ],
            outputs=gr.Textbox(label="Python Code for Visualization"),
            title="Python Code Generation"
        ),
        gr.Interface(
            fn=visualization_tool.entrypoint,
            inputs=[
                gr.Textbox(label="Visualization", lines=10)
            ],
            outputs=gr.Textbox(label="Saved Image Path"),
            title="Auto Visualization"
        )
    ],
    tab_names=["Query Products", "Python Code Generation", "Auto Visualization"]
)

# tabbed.launch(mcp_server=True)
tabbed.launch(server_port=7863, mcp_server=True)