from __future__ import annotations

import csv
import io
import json
from enum import Enum


class OutputFormat(str, Enum):
    json = "json"
    table = "table"
    csv = "csv"


def _cell_str(v: object) -> str:
    """Format a cell value for table/CSV display."""
    if v is None:
        return ""
    if isinstance(v, (dict, list)):
        return json.dumps(v, ensure_ascii=False)
    return str(v)


def format_output(data: dict | list | str | None, fmt: OutputFormat) -> str:
    """Format data for CLI output. Returns a string ready to print."""
    if fmt == OutputFormat.json:
        return json.dumps(data, ensure_ascii=False, indent=2)

    if fmt == OutputFormat.csv:
        return _format_csv(data)

    if fmt == OutputFormat.table:
        return _format_table(data)

    return json.dumps(data, ensure_ascii=False, indent=2)


def _format_csv(data: dict | list | str | None) -> str:
    if data is None:
        return ""
    if isinstance(data, dict):
        data = [data]
    if isinstance(data, str):
        return data
    if not isinstance(data, list) or not data:
        return ""
    rows = data if isinstance(data[0], dict) else [{"value": r} for r in data]
    # Collect ALL keys across all rows
    all_keys = dict.fromkeys(k for row in rows for k in row.keys())
    buf = io.StringIO()
    writer = csv.DictWriter(buf, fieldnames=all_keys.keys(), extrasaction="ignore")
    writer.writeheader()
    for row in rows:
        writer.writerow({k: _cell_str(v) for k, v in row.items()})
    return buf.getvalue()


def _format_table(data: dict | list | str | None) -> str:
    from rich.console import Console
    from rich.table import Table

    console = Console(file=io.StringIO(), force_terminal=False)

    if isinstance(data, list) and data and isinstance(data[0], dict):
        all_keys = list(dict.fromkeys(k for row in data for k in row.keys()))
        table = Table()
        for key in all_keys:
            table.add_column(key)
        for row in data:
            table.add_row(*[_cell_str(row.get(col)) for col in all_keys])
        console.print(table)
    elif isinstance(data, dict):
        table = Table()
        table.add_column("Field")
        table.add_column("Value")
        for k, v in data.items():
            table.add_row(str(k), _cell_str(v))
        console.print(table)
    else:
        console.print(str(data))

    return console.file.getvalue()
