"""
Multi-Turn Text-to-SQL Agent with Clarification Capabilities
=============================================================

An intelligent SQL assistant that:
- Answers clear database questions with accurate SQL
- Detects ambiguous questions and asks targeted clarifications
- Explains when questions can't be answered with available data
- Self-corrects SQL errors via ReAct reasoning loop
- Maintains multi-turn conversation context

Architecture based on:
- MMSQL (arXiv:2412.17867) — 4-type question classification
- PRACTIQ (arXiv:2410.11076) — clarification dialogue patterns
- SQLFixAgent (arXiv:2406.13408) — self-correcting SQL generation

Built with smolagents CodeAgent + Gradio UI.
"""

import os
import sqlite3
from textwrap import dedent

from smolagents import (
    tool,
    CodeAgent,
    InferenceClientModel,
    GradioUI,
)

# ─────────────────────────────────────────────
# 1. Database Setup — Sample multi-table DB
# ─────────────────────────────────────────────

DB_PATH = "demo_company.db"


def create_demo_database(db_path: str = DB_PATH):
    """Creates a rich demo company database with realistic data and some ambiguous schema elements."""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    for table in ["order_items", "orders", "products", "customers", "employees", "departments"]:
        cursor.execute(f"DROP TABLE IF EXISTS {table}")

    cursor.execute("""
        CREATE TABLE departments (
            dept_id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            location TEXT,
            budget REAL
        )
    """)
    cursor.executemany("INSERT INTO departments VALUES (?, ?, ?, ?)", [
        (1, "Engineering", "San Francisco", 2500000.00),
        (2, "Sales", "New York", 1800000.00),
        (3, "Marketing", "New York", 1200000.00),
        (4, "HR", "Chicago", 800000.00),
        (5, "Finance", "Chicago", 950000.00),
    ])

    cursor.execute("""
        CREATE TABLE employees (
            emp_id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            email TEXT,
            dept_id INTEGER REFERENCES departments(dept_id),
            salary REAL,
            hire_date TEXT,
            manager_id INTEGER REFERENCES employees(emp_id),
            status TEXT DEFAULT 'active'
        )
    """)
    cursor.executemany("INSERT INTO employees VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [
        (1, "Alice Chen", "alice@company.com", 1, 145000, "2019-03-15", None, "active"),
        (2, "Bob Martinez", "bob@company.com", 1, 128000, "2020-06-01", 1, "active"),
        (3, "Carol Smith", "carol@company.com", 2, 95000, "2021-01-10", None, "active"),
        (4, "David Lee", "david@company.com", 2, 88000, "2021-08-20", 3, "active"),
        (5, "Eva Johnson", "eva@company.com", 3, 102000, "2020-11-05", None, "active"),
        (6, "Frank Wilson", "frank@company.com", 1, 135000, "2019-07-22", 1, "active"),
        (7, "Grace Kim", "grace@company.com", 4, 78000, "2022-02-14", None, "active"),
        (8, "Henry Brown", "henry@company.com", 5, 115000, "2020-04-30", None, "active"),
        (9, "Iris Davis", "iris@company.com", 2, 92000, "2022-09-01", 3, "active"),
        (10, "Jack Taylor", "jack@company.com", 1, 140000, "2019-11-18", 1, "inactive"),
        (11, "Karen White", "karen@company.com", 3, 98000, "2021-05-12", 5, "active"),
        (12, "Leo Garcia", "leo@company.com", 5, 105000, "2021-03-28", 8, "active"),
    ])

    cursor.execute("""
        CREATE TABLE customers (
            customer_id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            email TEXT,
            city TEXT,
            state TEXT,
            signup_date TEXT,
            tier TEXT DEFAULT 'standard'
        )
    """)
    cursor.executemany("INSERT INTO customers VALUES (?, ?, ?, ?, ?, ?, ?)", [
        (1, "Acme Corp", "contact@acme.com", "San Francisco", "CA", "2020-01-15", "premium"),
        (2, "Beta Industries", "info@beta.com", "New York", "NY", "2020-03-22", "standard"),
        (3, "Gamma Solutions", "hello@gamma.com", "Chicago", "IL", "2020-06-10", "premium"),
        (4, "Delta Systems", "sales@delta.com", "Austin", "TX", "2021-02-05", "enterprise"),
        (5, "Epsilon LLC", "team@epsilon.com", "Seattle", "WA", "2021-08-18", "standard"),
        (6, "Zeta Partners", "info@zeta.com", "Boston", "MA", "2022-01-30", "premium"),
        (7, "Eta Global", "contact@eta.com", "Denver", "CO", "2022-07-14", "standard"),
        (8, "Theta Inc", "hello@theta.com", "Portland", "OR", "2023-03-01", "enterprise"),
    ])

    cursor.execute("""
        CREATE TABLE products (
            product_id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            category TEXT,
            price REAL,
            cost REAL,
            stock_quantity INTEGER,
            status TEXT DEFAULT 'active'
        )
    """)
    cursor.executemany("INSERT INTO products VALUES (?, ?, ?, ?, ?, ?, ?)", [
        (1, "Widget Pro", "Hardware", 299.99, 150.00, 500, "active"),
        (2, "Widget Basic", "Hardware", 149.99, 75.00, 1200, "active"),
        (3, "DataSync Cloud", "Software", 49.99, 10.00, None, "active"),
        (4, "DataSync Enterprise", "Software", 199.99, 40.00, None, "active"),
        (5, "SecureVault", "Software", 89.99, 20.00, None, "active"),
        (6, "PowerAdapter X", "Hardware", 39.99, 18.00, 3000, "active"),
        (7, "Legacy Suite", "Software", 299.99, 60.00, None, "discontinued"),
        (8, "SmartHub", "Hardware", 449.99, 220.00, 200, "active"),
    ])

    cursor.execute("""
        CREATE TABLE orders (
            order_id INTEGER PRIMARY KEY,
            customer_id INTEGER REFERENCES customers(customer_id),
            employee_id INTEGER REFERENCES employees(emp_id),
            order_date TEXT,
            status TEXT,
            total_amount REAL
        )
    """)
    cursor.executemany("INSERT INTO orders VALUES (?, ?, ?, ?, ?, ?)", [
        (1001, 1, 3, "2024-01-15", "completed", 1499.95),
        (1002, 2, 4, "2024-01-22", "completed", 599.96),
        (1003, 3, 3, "2024-02-10", "completed", 899.97),
        (1004, 1, 9, "2024-02-28", "shipped", 449.99),
        (1005, 4, 4, "2024-03-05", "completed", 2499.90),
        (1006, 5, 3, "2024-03-18", "pending", 149.99),
        (1007, 6, 9, "2024-04-02", "completed", 749.97),
        (1008, 3, 3, "2024-04-15", "completed", 339.98),
        (1009, 7, 4, "2024-05-01", "cancelled", 299.99),
        (1010, 8, 9, "2024-05-20", "shipped", 1349.97),
        (1011, 1, 3, "2024-06-01", "completed", 199.98),
        (1012, 4, 4, "2024-06-15", "completed", 3599.88),
    ])

    cursor.execute("""
        CREATE TABLE order_items (
            item_id INTEGER PRIMARY KEY,
            order_id INTEGER REFERENCES orders(order_id),
            product_id INTEGER REFERENCES products(product_id),
            quantity INTEGER,
            unit_price REAL,
            discount REAL DEFAULT 0.0
        )
    """)
    cursor.executemany("INSERT INTO order_items VALUES (?, ?, ?, ?, ?, ?)", [
        (1, 1001, 1, 5, 299.99, 0.0),
        (2, 1002, 2, 4, 149.99, 0.0),
        (3, 1003, 3, 6, 49.99, 0.0),
        (4, 1003, 5, 3, 89.99, 10.0),
        (5, 1004, 8, 1, 449.99, 0.0),
        (6, 1005, 1, 5, 299.99, 0.0),
        (7, 1005, 4, 5, 199.99, 0.0),
        (8, 1006, 2, 1, 149.99, 0.0),
        (9, 1007, 5, 3, 89.99, 0.0),
        (10, 1007, 3, 9, 49.99, 10.0),
        (11, 1008, 6, 5, 39.99, 0.0),
        (12, 1008, 3, 3, 49.99, 10.0),
        (13, 1009, 1, 1, 299.99, 0.0),
        (14, 1010, 8, 3, 449.99, 0.0),
        (15, 1011, 3, 4, 49.99, 0.0),
        (16, 1012, 4, 12, 199.99, 15.0),
        (17, 1012, 8, 2, 449.99, 10.0),
    ])

    conn.commit()
    conn.close()
    return db_path


# ─────────────────────────────────────────────
# 2. Build Dynamic Schema Description
# ─────────────────────────────────────────────

def get_schema_description(db_path: str = DB_PATH) -> str:
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
    tables = [row[0] for row in cursor.fetchall()]

    schema_parts = []
    for table in tables:
        cursor.execute(f"PRAGMA table_info({table})")
        columns = cursor.fetchall()
        cursor.execute(f"PRAGMA foreign_key_list({table})")
        fks = cursor.fetchall()
        fk_map = {fk[3]: f"→ {fk[2]}({fk[4]})" for fk in fks}
        cursor.execute(f"SELECT COUNT(*) FROM {table}")
        row_count = cursor.fetchone()[0]

        table_desc = f"Table '{table}' ({row_count} rows):\n  Columns:\n"
        for col in columns:
            col_id, col_name, col_type, not_null, default, pk = col
            parts = [f"    - {col_name}: {col_type or 'TEXT'}"]
            if pk: parts.append("PRIMARY KEY")
            if not_null and not pk: parts.append("NOT NULL")
            if default is not None: parts.append(f"DEFAULT {default}")
            if col_name in fk_map: parts.append(f"FK {fk_map[col_name]}")
            table_desc += " ".join(parts) + "\n"

        for col in columns:
            col_name, col_type = col[1], col[2]
            if col_type in ("TEXT", None) and col_name not in ("email",):
                try:
                    cursor.execute(f"SELECT DISTINCT {col_name} FROM {table} WHERE {col_name} IS NOT NULL LIMIT 8")
                    vals = [str(r[0]) for r in cursor.fetchall()]
                    if vals:
                        table_desc += f"  Sample '{col_name}' values: {', '.join(vals)}\n"
                except:
                    pass
        schema_parts.append(table_desc)

    conn.close()
    return "\n".join(schema_parts)


# ─────────────────────────────────────────────
# 3. Define Agent Tools
# ─────────────────────────────────────────────

SCHEMA_DESCRIPTION = ""


@tool
def execute_sql(query: str) -> str:
    """
    Executes a SQL query against the company database and returns the results.
    Use this tool to run SELECT queries to answer user questions about the data.
    
    IMPORTANT RULES:
    - Only use SELECT statements (no INSERT, UPDATE, DELETE, DROP)
    - Always use table and column names exactly as shown in the schema
    - Use JOINs when data spans multiple tables
    - Use LIMIT to avoid overwhelming output (max 50 rows)
    
    DATABASE SCHEMA:
    {schema}

    Args:
        query: A valid SQL SELECT query to execute against the database.
    """
    cleaned = query.strip().upper()
    if not cleaned.startswith("SELECT") and not cleaned.startswith("WITH"):
        return "ERROR: Only SELECT queries are allowed."

    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute(query)
        columns = [desc[0] for desc in cursor.description] if cursor.description else []
        rows = cursor.fetchall()
        conn.close()

        if not rows:
            return f"Query executed successfully.\nColumns: {', '.join(columns)}\nResult: No rows returned."

        result = f"Query executed successfully. {len(rows)} row(s) returned.\n"
        result += "Columns: " + " | ".join(columns) + "\n"
        result += "-" * 60 + "\n"
        for row in rows[:50]:
            result += " | ".join(str(v) for v in row) + "\n"
        if len(rows) > 50:
            result += f"... ({len(rows) - 50} more rows truncated)\n"
        return result

    except Exception as e:
        return f"SQL ERROR: {str(e)}\n\nPlease check your query syntax and column/table names against the schema."


@tool
def inspect_schema(table_name: str = "") -> str:
    """
    Inspect the database schema. If a table_name is provided, shows detailed info
    about that specific table including column types, foreign keys, and sample data.
    If no table_name is given, shows an overview of all tables.
    
    Use this tool BEFORE writing SQL to understand the database structure,
    especially when the user's question is ambiguous about which tables or columns to use.
    
    Args:
        table_name: Name of a specific table to inspect. Leave empty for full schema overview.
    """
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    if not table_name:
        return f"DATABASE SCHEMA OVERVIEW:\n\n{SCHEMA_DESCRIPTION}"

    try:
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = cursor.fetchall()
        if not columns:
            conn.close()
            return f"Table '{table_name}' not found. Use inspect_schema() with no arguments to see all tables."

        result = f"DETAILED INSPECTION OF TABLE '{table_name}':\n\n"
        result += "Columns:\n"
        for col in columns:
            result += f"  {col[1]} ({col[2] or 'TEXT'})"
            if col[5]: result += " [PRIMARY KEY]"
            if col[3]: result += " [NOT NULL]"
            result += "\n"

        cursor.execute(f"PRAGMA foreign_key_list({table_name})")
        fks = cursor.fetchall()
        if fks:
            result += "\nForeign Keys:\n"
            for fk in fks:
                result += f"  {fk[3]} → {fk[2]}({fk[4]})\n"

        cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
        count = cursor.fetchone()[0]
        result += f"\nTotal rows: {count}\n"

        cursor.execute(f"SELECT * FROM {table_name} LIMIT 3")
        sample_rows = cursor.fetchall()
        col_names = [c[1] for c in columns]
        result += f"\nSample rows (first 3):\n"
        result += " | ".join(col_names) + "\n"
        result += "-" * 60 + "\n"
        for row in sample_rows:
            result += " | ".join(str(v) for v in row) + "\n"

        conn.close()
        return result

    except Exception as e:
        conn.close()
        return f"Error inspecting table: {str(e)}"


# ─────────────────────────────────────────────
# 4. Agent System Prompt
# ─────────────────────────────────────────────

SYSTEM_INSTRUCTIONS = dedent("""\
You are an expert SQL assistant that helps users query a company database. You follow a structured multi-turn approach:

## YOUR DECISION PROCESS

For EVERY user question, follow these steps:

### Step 1: Classify the Question
Determine if the question is:
- **ANSWERABLE**: The question is clear and maps directly to the database schema
- **AMBIGUOUS**: The question could have multiple valid SQL interpretations (e.g., "show me the top employees" — top by salary? by sales? by tenure?)
- **UNANSWERABLE**: The question asks for data that doesn't exist in the database

### Step 2: Handle Based on Classification

**If AMBIGUOUS:**
- Identify ALL possible interpretations
- Use `final_answer()` to return a targeted clarification question listing the specific options
- Example: call `final_answer("Your question could mean several things:\\n1. Employees with highest salary\\n2. Employees who handled the most orders\\n3. Employees with the longest tenure\\n\\nWhich interpretation do you mean?")`
- Do NOT generate SQL — return the clarification question immediately using `final_answer()`
- The user will respond in the next turn with their clarification

**If UNANSWERABLE:**
- Use `final_answer()` to explain clearly what data is missing and why the question can't be answered
- Include a suggestion for a related question that CAN be answered with the available data

**If ANSWERABLE:**
- First inspect the schema to confirm the right tables/columns
- Generate and execute the SQL query
- Present results clearly with a natural language summary

### Step 3: Self-Correct
- If your SQL returns an error, analyze the error and fix the query
- If the result seems wrong or empty, verify your joins and filters
- Always sanity-check: does the result make sense given what was asked?

## COMMON AMBIGUITY PATTERNS TO WATCH FOR

1. **Column ambiguity**: "Show employee names" — the 'name' column appears in employees, departments, customers, and products tables
2. **Metric ambiguity**: "Top customers" — by total spending? by number of orders? by most recent activity?
3. **Filter ambiguity**: "Recent orders" — last week? last month? last quarter?
4. **Scope ambiguity**: "Total sales" — all time? this year? by product? by employee?
5. **Status ambiguity**: "List products" — all products? only active ones? including discontinued?
6. **Value ambiguity**: "Expensive products" — what price threshold?

## FORMATTING RULES

- When presenting query results, format them as a clear table
- Always explain what the query does in plain language
- If you make assumptions (e.g., "I'm assuming you mean active employees only"), state them explicitly
- For numerical results, include relevant aggregations (count, sum, average) when helpful
""")


# ─────────────────────────────────────────────
# 5. Main
# ─────────────────────────────────────────────

def create_agent(model_id: str = "Qwen/Qwen2.5-Coder-32B-Instruct"):
    create_demo_database()

    global SCHEMA_DESCRIPTION
    SCHEMA_DESCRIPTION = get_schema_description()
    execute_sql.description = execute_sql.description.replace("{schema}", SCHEMA_DESCRIPTION)

    model = InferenceClientModel(model_id=model_id)

    agent = CodeAgent(
        tools=[execute_sql, inspect_schema],
        model=model,
        instructions=SYSTEM_INSTRUCTIONS,
        max_steps=15,
        additional_authorized_imports=["json", "re"],
    )
    return agent


if __name__ == "__main__":
    agent = create_agent()
    ui = GradioUI(agent, reset_agent_memory=False)
    ui.launch()