import os
import gradio as gr
import pandas as pd
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, insert, text, inspect
from smolagents import tool, CodeAgent, InferenceClientModel

# --- Setup SQLite database ---
engine = create_engine("sqlite:///data.db")
metadata_obj = MetaData()

# --- Create students table ---
students = Table(
    "students",
    metadata_obj,
    Column("student_id", Integer, primary_key=True),
    Column("student_name", String(32), nullable=False),
)
metadata_obj.create_all(engine)

# Insert sample student data
student_rows = [
    {"student_id": 1, "student_name": "Alice Johnson"},
    {"student_id": 2, "student_name": "Bob Smith"},
    {"student_id": 3, "student_name": "Charlie Brown"},
    {"student_id": 4, "student_name": "Diana Prince"},
]
for row in student_rows:
    stmt = insert(students).values(**row)
    with engine.begin() as conn:
        conn.execute(stmt)

# --- Create subjects table ---
subjects = Table(
    "subjects",
    metadata_obj,
    Column("student_id", Integer, primary_key=True),
    Column("subject_name", String(32), primary_key=True),
    Column("marks", Float, nullable=False),
    Column("teacher", String(32), nullable=False),
)
metadata_obj.create_all(engine)

# Insert sample subject/marks data
subject_rows = [
    {"student_id": 1, "subject_name": "Math", "marks": 95.0, "teacher": "Mr. Adams"},
    {"student_id": 1, "subject_name": "English", "marks": 88.5, "teacher": "Ms. Baker"},
    {"student_id": 1, "subject_name": "Science", "marks": 92.0, "teacher": "Dr. Carter"},
    {"student_id": 2, "subject_name": "Math", "marks": 72.0, "teacher": "Mr. Adams"},
    {"student_id": 2, "subject_name": "Science", "marks": 81.0, "teacher": "Dr. Carter"},
    {"student_id": 3, "subject_name": "Math", "marks": 85.0, "teacher": "Mr. Adams"},
    {"student_id": 3, "subject_name": "History", "marks": 90.0, "teacher": "Mrs. Davis"},
    {"student_id": 4, "subject_name": "English", "marks": 78.0, "teacher": "Ms. Baker"},
    {"student_id": 4, "subject_name": "Science", "marks": 82.0, "teacher": "Dr. Carter"},
]
for row in subject_rows:
    stmt = insert(subjects).values(**row)
    with engine.begin() as conn:
        conn.execute(stmt)

# --- Convert tables to DataFrames for front-end display ---
def fetch_table(table_name: str) -> pd.DataFrame:
    with engine.connect() as con:
        rows = con.execute(text(f"SELECT * FROM {table_name}"))
        df = pd.DataFrame(rows.fetchall(), columns=rows.keys())
    return df

students_df = fetch_table("students")
subjects_df = fetch_table("subjects")

# --- Define SQL tool for the agent ---
@tool
def sql_engine(query: str) -> list:
    """
    Executes SQL queries on the available tables: students and subjects.

    Args:
        query (str): The SQL query string to execute.

    Returns:
        list: List of tuples containing the query results.
    """
    try:
        print(f"🧩 Executing query: {query}")
        with engine.connect() as con:
            rows = con.execute(text(query))
            results = [tuple(row) for row in rows]
        return results or []
    except Exception as e:
        return [f" SQL Error: {str(e)}"]

# Dynamically describe tables for the agent
updated_description = "You can run SQL queries on these tables:\n"
inspector = inspect(engine)
for table in ["students", "subjects"]:
    columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
    table_description = f"\n\nTable '{table}':\nColumns:\n" + "\n".join(
        [f" - {name}: {col_type}" for name, col_type in columns_info]
    )
    updated_description += table_description

sql_engine.description = updated_description

# --- Create the agent ---
agent = CodeAgent(
    tools=[sql_engine],
    model=InferenceClientModel(
        "meta-llama/Meta-Llama-3-8B-Instruct",
        api_key=os.environ.get("HF_TOKEN")
    ),
)

# --- Sample prompts to display ---
sample_prompts = [
    "Which student scored highest in Math?",
    "List all subjects and marks for Alice Johnson.",
    "Who is the teacher of Science for Diana Prince?",
    "Show students who scored more than 80 in Science.",
    "Average marks per subject."
]

# --- Define Gradio interface ---
def ask_agent(question: str) -> str:
    """Ask the AI agent a question and return its answer."""
    try:
        result = agent.run(question)
        if isinstance(result, list):
            result_str = "\n".join(str(r) for r in result)
            return result_str or "No results."
        return str(result)
    except Exception as e:
        return f" Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("##  Text to SQL Agent")
    gr.Markdown("### Students Table")
    gr.DataFrame(value=students_df, interactive=False)
    gr.Markdown("### Subjects Table")
    gr.DataFrame(value=subjects_df, interactive=False)
    gr.Markdown("### Ask questions about students, subjects, marks, and teachers")
    question_input = gr.Textbox(label="Your question", placeholder="e.g., Which student scored highest in Math?")
    answer_output = gr.Textbox(label="Agent response")
    gr.Markdown("### Sample Prompts")
    for prompt in sample_prompts:
        gr.Markdown(f"- {prompt}")
    question_input.submit(fn=ask_agent, inputs=question_input, outputs=answer_output)

demo.launch()