from __future__ import annotations

import csv
import json
from io import StringIO
from typing import Iterable

from rich.console import Console
from rich.table import Table

from awspipe.backend import QueryResult


def render_table(console: Console, result: QueryResult) -> None:
    table = Table(show_header=True, header_style="bold cyan")
    for column in result.columns:
        table.add_column(column)
    for row in result.rows:
        table.add_row(*["" if value is None else str(value) for value in row])
    console.print(table)


def render_json(console: Console, result: QueryResult) -> None:
    payload = [dict(zip(result.columns, row, strict=False)) for row in result.rows]
    console.print_json(json.dumps(payload, default=str))


def render_csv(console: Console, result: QueryResult) -> None:
    output = StringIO()
    writer = csv.writer(output)
    writer.writerow(result.columns)
    writer.writerows(result.rows)
    console.print(output.getvalue().strip())


def render_rows(console: Console, rows: Iterable[dict[str, object]]) -> None:
    materialized = list(rows)
    if not materialized:
        console.print("[yellow]No rows returned.[/yellow]")
        return
    result = QueryResult(columns=list(materialized[0].keys()), rows=[tuple(row.values()) for row in materialized])
    render_table(console, result)
