import json
import os
import re
import time

import gradio as gr
from huggingface_hub import InferenceClient

from sqlagent import database
from sqlagent.seed import seed_database

DB_PATH = "/tmp/shop.db"
MODEL = "Qwen/Qwen2.5-7B-Instruct"
MAX_QUESTION_CHARS = 400

client = InferenceClient()


def ensure_db():
    if not os.path.exists(DB_PATH):
        connection = database.connect(DB_PATH)
        seed_database(connection)
        connection.close()


ensure_db()


def schema_text():
    with database.connect(DB_PATH, read_only=True) as connection:
        lines = []
        for table in database.list_tables(connection):
            cols = database.table_schema(connection, table)
            joined = ", ".join(f"{c['column']} {c['type']}" for c in cols)
            lines.append(f"{table}({joined})")
        return "\n".join(lines)


def ask(messages, max_tokens=350):
    last_error = None
    for attempt in range(3):
        try:
            response = client.chat_completion(messages=messages, model=MODEL, max_tokens=max_tokens)
            return response.choices[0].message.content.strip()
        except Exception as error:
            last_error = error
            time.sleep(1.5 * (attempt + 1))
    raise last_error


def extract_sql(text):
    fence = re.search(r"```(?:sql)?\s*(.+?)```", text, re.S)
    body = fence.group(1) if fence else text
    match = re.search(r"(?is)\b(select|with)\b.*", body)
    sql = match.group(0) if match else body
    sql = sql.split(";")[0]
    sql = sql.split("\n\n")[0]
    return sql.strip()


def answer(question):
    question = (question or "").strip()
    if not question:
        return "Ask a question about the shop database."
    if len(question) > MAX_QUESTION_CHARS:
        return f"Please keep the question under {MAX_QUESTION_CHARS} characters."
    try:
        schema = schema_text()
        raw_sql = ask([
            {"role": "system", "content": "You write SQLite SQL. Given the schema, write ONE read-only SELECT query that answers the question. Return only the SQL query: no explanation, no markdown, no comments."},
            {"role": "user", "content": f"Schema:\n{schema}\n\nQuestion: {question}"},
        ], max_tokens=200)
        sql = extract_sql(raw_sql)
        with database.connect(DB_PATH, read_only=True) as connection:
            rows = database.run_select(connection, sql)
        results = json.dumps(rows, default=str)[:2000]
        final = ask([
            {"role": "system", "content": "Answer the question concisely using the query results. State the figures clearly."},
            {"role": "user", "content": f"Question: {question}\nResults: {results}"},
        ], max_tokens=300)
        return f"{final}\n\n---\nSQL used:\n{sql}"
    except Exception:
        return "The service is busy or unavailable right now. Please try again in a moment."


demo = gr.Interface(
    fn=answer,
    inputs=gr.Textbox(lines=2, label="Ask a question about the shop database",
                      placeholder="Which product category brings the most revenue?"),
    outputs=gr.Textbox(lines=10, label="Answer"),
    title="SQL Question Answering Agent",
    description=(
        "Ask a question in plain English about a small sample shop database "
        "(customers, products, orders). The system writes a read-only SQL query, "
        "runs it and answers from the results, showing the SQL it used."
    ),
    article="Code: https://github.com/delcenjo/llm-sql-agent",
    cache_examples=False,
    examples=[
        ["Which product category brings the most revenue?"],
        ["How many customers are from Spain?"],
        ["List the top 3 products by total quantity sold."],
    ],
)

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