from typing import Any
import typer
from rich.console import Console
from rich.table import Table
from rich.text import Text

from cli.util import resolve_environment, ENVIRONMENT_OPTION


def execute_statement(cursor: Any, statement: str, console: Console, exit_on_error: bool = True):
    try:
        cursor.execute(statement)
    except Exception as e:
        console.print(Text(f"❌ Error executing statement: {e}", style="bold red"))
        if exit_on_error:
            raise typer.Exit(code=1)
        else:
            return None

    result = cursor.fetchall()
    if isinstance(result, list):
        table = Table(title="SQL Result", show_header=True, header_style="bold cyan")
        col_n = 0
        if cursor.description:
            for col in cursor.description:
                if col.name:
                    table.add_column(col.name, style="green")
                else:
                    table.add_column(f"Column {col_n + 1}", style="green")
                col_n += 1
        else:
            for col_n in range(len(result[0])):
                table.add_column(f"Column {col_n + 1}", style="green")
        for row in result:
            row_values = [str(item) for item in row]
            table.add_row(*row_values)
        console.print(table)
    else:
        console.print(Text(str(result), style="bold green"))


def handle_snowflake_error(e: Exception, console: Console):
    if "Account not found" in str(e):
        console.print(Text("❌ Account not found. Please run `komodo account set` to set the account.", style="bold red"))
    else:
        console.print(Text(f"❌ Error getting Snowflake connection: {e}", style="bold red"))
    raise typer.Exit(code=1)


def get_connection_and_cursor(environment: str, console: Console):
    from komodo.auth.Session import Session
    from komodo import get_snowflake_connection

    session = Session(environment=environment)
    credentials = session.load_credentials()

    if not credentials or not credentials.access_token:
        console.print(Text("❌ No access token found. Please run 'komodo login' first.", style="bold red"))
        raise typer.Exit(code=1)

    try:
        conn = get_snowflake_connection()
    except Exception as e:
        handle_snowflake_error(e, console)

    try:
        cursor = conn.cursor()
    except Exception as e:
        conn.close()
        handle_snowflake_error(e, console)

    return conn, cursor


@resolve_environment
def sql_execute(environment: str = ENVIRONMENT_OPTION, statement: str = typer.Argument(..., help="The SQL statement to execute.")):
    """Execute a SQL query.

    Args:
        environment: The environment to use.
        statement: The SQL statement to execute.
    """
    console = Console()

    conn, cursor = get_connection_and_cursor(environment, console)

    try:
        execute_statement(cursor, statement, console, exit_on_error=True)
    finally:
        cursor.close()
        conn.close()


@resolve_environment
def sql_shell(environment: str = ENVIRONMENT_OPTION):
    """Open an interactive shell to execute Snowflake statements."""

    console = Console()

    conn, cursor = get_connection_and_cursor(environment, console)

    while True:
        query = typer.prompt("Enter a SQL statement to execute (or 'exit' to quit)")
        if query.lower() == "exit":
            break
        execute_statement(cursor, query, console, exit_on_error=False)

    cursor.close()
    conn.close()
