from __future__ import annotations

from prompt_toolkit import PromptSession
from rich.console import Console

from awspipe.engine import QueryEngine
from awspipe.render import render_table


def handle_meta_command(
    console: Console,
    engine: QueryEngine,
    line: str,
    profile: str | None,
    region: str | None,
    all_regions: bool,
) -> bool:
    if line == ".tables":
        for provider in engine.available_tables():
            console.print(f"{provider.name}\t{provider.scope}\t{provider.description}")
        return True
    if line.startswith(".inspect "):
        _, table_name = line.split(maxsplit=1)
        provider, columns = engine.inspect_table(table_name)
        console.print(f"{provider.name}: {provider.description}")
        for column in columns:
            console.print(f"  {column.column} ({column.type_name}) - {column.description}")
        return True
    if line.startswith(".schema "):
        _, table_name = line.split(maxsplit=1)
        _, columns = engine.inspect_table(table_name)
        for column in columns:
            console.print(f"{column.column} {column.type_name}")
        return True
    if line == ".connections":
        console.print("AWS CLI credentials are resolved from --profile, AWS_PROFILE, or the default AWS SDK chain.")
        return True
    if line.startswith(".refresh "):
        _, table_name = line.split(maxsplit=1)
        engine.query(f"select * from {table_name}", profile=profile, region=region, refresh=True, all_regions=all_regions)
        console.print(f"Refreshed {table_name}.")
        return True
    return False


def run_shell(console: Console, engine: QueryEngine, profile: str | None, region: str | None, all_regions: bool) -> None:
    session = PromptSession("awspipe> ")
    while True:
        try:
            line = session.prompt().strip()
        except (EOFError, KeyboardInterrupt):
            console.print()
            return
        if not line:
            continue
        if line == ".quit":
            return
        if line.startswith("."):
            if handle_meta_command(console, engine, line, profile=profile, region=region, all_regions=all_regions):
                continue
        result = engine.query(line, profile=profile, region=region, all_regions=all_regions)
        render_table(console, result)
