"""Formatting utilities for chain execution results."""

from __future__ import annotations

import csv
import io
import json
from datetime import datetime


def to_markdown(chain_result: dict) -> str:
    """Format a chain result dict as Markdown.

    Produces a ``# Chain Result`` header followed by one
    ``## Step N: <model>`` section per step with the output as a quote
    block and timing in italics.

    Args:
        chain_result: Dict as returned by the ``/chain`` endpoint with
            keys ``steps`` (list of step dicts), optionally ``success``
            and ``total_duration_ms``.

    Returns:
        A Markdown-formatted string.
    """
    lines: list[str] = ["# Chain Result", ""]

    steps = chain_result.get("steps", [])

    for i, step in enumerate(steps, 1):
        model = step.get("model", "unknown")
        output = step.get("content", step.get("output", ""))
        duration = step.get("duration_ms")

        lines.append(f"## Step {i}: {model}")
        lines.append("")

        for output_line in str(output).splitlines():
            lines.append(f"> {output_line}")
        lines.append("")

        if duration is not None:
            lines.append(f"*{duration}ms*")
            lines.append("")

    return "\n".join(lines)


def format_result(chain_result: dict, fmt: str) -> str:
    """Format a chain result based on the requested format.

    Dispatches to the appropriate formatter function. For formats that
    expect a list (e.g. csv, table), wraps the single result in a list.

    Args:
        chain_result: A single chain execution result dict.
        fmt: The requested format (markdown, html, json_pretty, plain_text,
            csv, table).

    Returns:
        The formatted string.

    Raises:
        ValueError: If the format is unknown.
    """
    f = fmt.lower()
    if f == "markdown":
        return to_markdown(chain_result)
    elif f == "html":
        return to_html(chain_result)
    elif f == "json_pretty" or f == "json":
        # Handle datetime serialization if needed, though default json
        # doesn't handle it out of the box without a default kwarg.
        import datetime
        return json.dumps(
            chain_result,
            indent=2,
            sort_keys=True,
            default=lambda o: o.isoformat() if isinstance(o, datetime.datetime) else str(o)
        )
    elif f == "plain_text" or f == "text":
        return to_plain_text(chain_result)
    elif f == "csv":
        return to_csv([chain_result])
    elif f == "table":
        return to_table([chain_result])
    else:
        raise ValueError(f"Unknown format: {fmt}")



def _escape_html(text: str) -> str:
    """Minimal HTML escaping for user-provided content."""
    return (
        text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


# Palette of badge colours cycled per step.
_BADGE_COLORS = [
    "#6366f1",  # indigo
    "#059669",  # emerald
    "#d97706",  # amber
    "#dc2626",  # red
    "#7c3aed",  # violet
    "#0284c7",  # sky
]


def to_html(chain_result: dict) -> str:
    """Format a chain result dict as a complete HTML document with inline CSS.

    Each step is rendered as a styled card with the model name in a coloured
    badge, the output in a ``<pre>`` block, and latency in a small footer.

    Args:
        chain_result: Dict as returned by the ``/chain`` endpoint with
            keys ``steps`` (list of step dicts), optionally ``success``
            and ``total_duration_ms``.

    Returns:
        A self-contained HTML document string.
    """
    steps = chain_result.get("steps", [])

    cards: list[str] = []
    for i, step in enumerate(steps, 1):
        model = _escape_html(step.get("model", "unknown"))
        output = _escape_html(str(step.get("content", step.get("output", ""))))
        duration = step.get("duration_ms")
        color = _BADGE_COLORS[(i - 1) % len(_BADGE_COLORS)]

        footer = ""
        if duration is not None:
            footer = (
                f'<div style="margin-top:8px;font-size:12px;color:#888;">'
                f"{_escape_html(str(duration))}ms</div>"
            )

        cards.append(
            f'<div style="background:#fff;border:1px solid #e5e7eb;border-radius:8px;'
            f'padding:16px;margin-bottom:16px;box-shadow:0 1px 3px rgba(0,0,0,0.06);">'
            f'<span style="display:inline-block;background:{color};color:#fff;'
            f'font-size:13px;font-weight:600;padding:2px 10px;border-radius:12px;'
            f'margin-bottom:10px;">Step {i}: {model}</span>'
            f'<pre style="background:#f9fafb;border:1px solid #e5e7eb;border-radius:4px;'
            f'padding:12px;margin:0;white-space:pre-wrap;word-wrap:break-word;'
            f'font-size:14px;line-height:1.5;">{output}</pre>'
            f"{footer}"
            f"</div>"
        )

    body = "\n".join(cards)

    return (
        "<!DOCTYPE html>\n"
        '<html lang="en">\n'
        "<head>\n"
        '<meta charset="utf-8">\n'
        '<meta name="viewport" content="width=device-width,initial-scale=1">\n'
        "<title>Chain Result</title>\n"
        "</head>\n"
        '<body style="margin:0;padding:32px;font-family:-apple-system,BlinkMacSystemFont,'
        '\'Segoe UI\',Roboto,sans-serif;background:#f3f4f6;color:#111827;">\n'
        '<div style="max-width:720px;margin:0 auto;">\n'
        '<h1 style="font-size:22px;margin:0 0 20px;">Chain Result</h1>\n'
        f"{body}\n"
        "</div>\n"
        "</body>\n"
        "</html>"
    )


_CSV_COLUMNS = [
    "timestamp",
    "strategy",
    "models",
    "input_preview",
    "output_preview",
    "tokens",
    "latency_ms",
]


def to_csv(chain_results: list[dict]) -> str:
    """Format a list of chain result dicts as CSV text.

    Args:
        chain_results: List of dicts as returned by the ``/chain`` or
            ``/v1/chain`` endpoint, each with keys like ``steps``,
            ``total_duration_ms`` / ``total_latency_ms``, ``strategy``,
            ``timestamp``, ``input``, etc.

    Returns:
        A CSV string with columns: timestamp, strategy, models,
        input_preview (first 100 chars), output_preview (first 200 chars),
        tokens, latency_ms.
    """
    buf = io.StringIO()
    writer = csv.writer(buf)
    writer.writerow(_CSV_COLUMNS)

    for result in chain_results:
        timestamp = result.get("timestamp", result.get("ts", ""))

        strategy = result.get("strategy", result.get("type", ""))

        models_raw = result.get("models", result.get("models_used", []))
        if isinstance(models_raw, list):
            models = ", ".join(models_raw)
        else:
            models = str(models_raw)

        # Input preview: top-level input/prompt, first 100 chars
        raw_input = result.get("input", result.get("prompt", ""))
        input_preview = str(raw_input)[:100]

        # Output preview: last step content, first 200 chars
        steps = result.get("steps", [])
        output = ""
        for step in reversed(steps):
            content = step.get("content", step.get("output", ""))
            if content:
                output = str(content)
                break
        output_preview = output[:200]

        # Tokens: sum across steps, or top-level value
        tokens = result.get("tokens", 0)
        if not tokens and steps:
            tokens = sum(
                s.get("token_count", s.get("tokens", 0)) for s in steps
            )

        latency_ms = result.get(
            "total_latency_ms", result.get("total_duration_ms", 0)
        )

        writer.writerow([
            timestamp,
            strategy,
            models,
            input_preview,
            output_preview,
            tokens,
            latency_ms,
        ])

    return buf.getvalue()


def _json_default(obj: object) -> str:
    """Default handler for json.dumps that converts datetime to ISO format."""
    if isinstance(obj, datetime):
        return obj.isoformat()
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


def to_json_pretty(chain_result: dict) -> str:
    """Format a chain result dict as pretty-printed JSON.

    Args:
        chain_result: Dict as returned by the ``/chain`` endpoint.

    Returns:
        A JSON string with 2-space indentation and sorted keys.
        Datetime objects are converted to ISO format strings.
    """
    return json.dumps(chain_result, indent=2, sort_keys=True, default=_json_default)


def to_table(chain_results: list[dict]) -> str:
    """Format a list of chain result dicts as a fixed-width ASCII table.

    Args:
        chain_results: List of dicts as returned by the ``/chain`` or
            ``/v1/chain`` endpoint, each with keys like ``steps``,
            ``total_duration_ms`` / ``total_latency_ms``, ``strategy``,
            ``models``, etc.

    Returns:
        A fixed-width ASCII table string with columns:
        #, Strategy, Models, Tokens, Latency.
    """
    headers = ("#", "Strategy", "Models", "Tokens", "Latency")

    rows: list[tuple[str, str, str, str, str]] = []
    for i, result in enumerate(chain_results, 1):
        strategy = result.get("strategy", result.get("type", ""))

        models_raw = result.get("models", result.get("models_used", []))
        if isinstance(models_raw, list):
            models = ", ".join(models_raw)
        else:
            models = str(models_raw)

        steps = result.get("steps", [])
        tokens = result.get("tokens", 0)
        if not tokens and steps:
            tokens = sum(
                s.get("token_count", s.get("tokens", 0)) for s in steps
            )

        latency = result.get(
            "total_latency_ms", result.get("total_duration_ms", 0)
        )

        rows.append((
            str(i),
            str(strategy),
            str(models),
            str(tokens),
            f"{latency}ms",
        ))

    # Compute column widths from headers and data.
    widths = [len(h) for h in headers]
    for row in rows:
        for j, cell in enumerate(row):
            widths[j] = max(widths[j], len(cell))

    def _fmt_row(cells: tuple[str, ...] | tuple[str, str, str, str, str]) -> str:
        parts = [cell.ljust(widths[j]) for j, cell in enumerate(cells)]
        return "| " + " | ".join(parts) + " |"

    separator = "+-" + "-+-".join("-" * w for w in widths) + "-+"

    lines = [separator, _fmt_row(headers), separator]
    for row in rows:
        lines.append(_fmt_row(row))
    lines.append(separator)

    return "\n".join(lines)


def format_result(chain_result: dict, fmt: str) -> str:
    """Format a chain result based on the requested format.

    Dispatches to the appropriate formatter function. For formats that
    expect a list (e.g. csv, table), wraps the single result in a list.

    Args:
        chain_result: A single chain execution result dict.
        fmt: The requested format (markdown, html, json_pretty, plain_text,
            csv, table).

    Returns:
        The formatted string.

    Raises:
        ValueError: If the format is unknown.
    """
    f = fmt.lower()
    if f == "markdown":
        return to_markdown(chain_result)
    elif f == "html":
        return to_html(chain_result)
    elif f == "json_pretty" or f == "json":
        # Handle datetime serialization if needed, though default json
        # doesn't handle it out of the box without a default kwarg.
        import datetime
        return json.dumps(
            chain_result,
            indent=2,
            sort_keys=True,
            default=lambda o: o.isoformat() if isinstance(o, datetime.datetime) else str(o)
        )
    elif f == "plain_text" or f == "text":
        return to_plain_text(chain_result)
    elif f == "csv":
        return to_csv([chain_result])
    elif f == "table":
        return to_table([chain_result])
    else:
        raise ValueError(f"Unknown format: {fmt}")



def to_plain_text(chain_result: dict) -> str:
    """Format a chain result dict as a clean plain-text summary.

    Each step is rendered as a single line with the model name and the first
    line of its output, followed by a total latency footer.

    Args:
        chain_result: Dict as returned by the ``/chain`` endpoint with
            keys ``steps`` (list of step dicts), optionally
            ``total_duration_ms`` / ``total_latency_ms``.

    Returns:
        A plain-text summary string.
    """
    lines: list[str] = []

    for step in chain_result.get("steps", []):
        model = step.get("model", "unknown")
        output = str(step.get("content", step.get("output", "")))
        first_line = output.split("\n", 1)[0].strip()
        if first_line:
            lines.append(f"Model: {model} | {first_line}...")
        else:
            lines.append(f"Model: {model} | (no output)")

    total = chain_result.get("total_duration_ms", chain_result.get("total_latency_ms", 0))
    lines.append(f"Total: {total}ms")

    return "\n".join(lines)


def format_result(chain_result: dict, fmt: str) -> str:
    """Format a chain result based on the requested format.

    Dispatches to the appropriate formatter function. For formats that
    expect a list (e.g. csv, table), wraps the single result in a list.

    Args:
        chain_result: A single chain execution result dict.
        fmt: The requested format (markdown, html, json_pretty, plain_text,
            csv, table).

    Returns:
        The formatted string.

    Raises:
        ValueError: If the format is unknown.
    """
    f = fmt.lower()
    if f == "markdown":
        return to_markdown(chain_result)
    elif f == "html":
        return to_html(chain_result)
    elif f == "json_pretty" or f == "json":
        # Handle datetime serialization if needed, though default json
        # doesn't handle it out of the box without a default kwarg.
        import datetime
        return json.dumps(
            chain_result,
            indent=2,
            sort_keys=True,
            default=lambda o: o.isoformat() if isinstance(o, datetime.datetime) else str(o)
        )
    elif f == "plain_text" or f == "text":
        return to_plain_text(chain_result)
    elif f == "csv":
        return to_csv([chain_result])
    elif f == "table":
        return to_table([chain_result])
    else:
        raise ValueError(f"Unknown format: {fmt}")

