from __future__ import annotations

import os
import sys
from typing import Annotated, Literal

import gradio as gr
from ._docstrings import autodoc

try:
    import mysql.connector
    from mysql.connector import Error
except ImportError:
    mysql = None
    Error = Exception

# Lấy thông tin kết nối từ biến môi trường
DEFAULT_HOST = os.getenv("MYSQL_HOST", "localhost")
DEFAULT_USER = os.getenv("MYSQL_USER", "root")
DEFAULT_PASSWORD = os.getenv("MYSQL_PASSWORD", "")
DEFAULT_DATABASE = os.getenv("MYSQL_DATABASE", "")
DEFAULT_PORT = int(os.getenv("MYSQL_PORT", "3306"))

# Single source of truth for the LLM-facing tool description
TOOL_SUMMARY = (
    "Execute MySQL queries against a database; returns query results or error messages. "
    "Uses environment variables for connection parameters if not provided."
)


@autodoc(
    summary=TOOL_SUMMARY,
)
def MySQL(
    query: Annotated[str, "SQL query to execute"],
    host: Annotated[str, "MySQL server host address"] = DEFAULT_HOST,
    username: Annotated[str, "MySQL username"] = DEFAULT_USER,
    password: Annotated[str, "MySQL password"] = DEFAULT_PASSWORD,
    database: Annotated[str, "Database name to connect to"] = DEFAULT_DATABASE,
    port: Annotated[int, "MySQL server port"] = DEFAULT_PORT,
    query_type: Annotated[Literal["SELECT", "INSERT", "UPDATE", "DELETE", "DDL"], "Type of SQL query to execute"] = "SELECT"
) -> str:
    if mysql is None:
        return "Error: mysql-connector-python package is not installed. Run: pip install mysql-connector-python"

    if not query or not query.strip():
        return "Error: Query is required."

    if not host or not username:
        return "Error: Host and username are required."

    connection = None
    try:
        # Establish connection to MySQL database
        connection = mysql.connector.connect(
            host=host,
            port=port,
            database=database,
            user=username,
            password=password
        )

        if connection.is_connected():
            cursor = connection.cursor(dictionary=True)
            cursor.execute(query)

            if query_type.upper() in ["SELECT", "SHOW", "DESCRIBE", "EXPLAIN"]:
                # Fetch results for SELECT queries
                records = cursor.fetchall()
                if records:
                    result = "Query Results:\n"
                    for i, row in enumerate(records):
                        result += f"Row {i+1}: {dict(row)}\n"
                    result += f"\nTotal records: {len(records)}"
                else:
                    result = "Query executed successfully. No records found."
            else:
                # For INSERT, UPDATE, DELETE, DDL queries
                connection.commit()
                affected_rows = cursor.rowcount
                result = f"Query executed successfully. Affected rows: {affected_rows}"

            cursor.close()
            return result
    except Error as e:
        error_msg = f"MySQL Error: {str(e)}"
        return error_msg
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        return error_msg
    finally:
        if connection and connection.is_connected():
            connection.close()


def build_interface() -> gr.Interface:
    return gr.Interface(
        fn=MySQL,
        inputs=[
            gr.Code(label="SQL Query", language="sql"),
            gr.Textbox(label="Host", value=DEFAULT_HOST),
            gr.Textbox(label="Username", value=DEFAULT_USER),
            gr.Textbox(label="Password", type="password", value=DEFAULT_PASSWORD),
            gr.Textbox(label="Database", value=DEFAULT_DATABASE),
            gr.Number(label="Port", value=DEFAULT_PORT),
            gr.Radio(
                label="Query Type",
                choices=["SELECT", "INSERT", "UPDATE", "DELETE", "DDL"],
                value="SELECT"
            )
        ],
        outputs=gr.Textbox(label="Result", lines=10, max_lines=20),
        title="MySQL Database Connector",
        description="<div style=\"text-align:center\">Execute MySQL queries against a database. Uses environment variables for defaults.</div>",
        api_description=TOOL_SUMMARY,
        flagging_mode="never",
    )


__all__ = ["MySQL", "build_interface"]