from typing import Any, cast

from pydantic import BaseModel

from cli.infrastructure.output import OutputFormat, OutputService
from cli.infrastructure.renderers.entity_renderer import FieldConfig
from cli.infrastructure.renderers.list_renderer import ColumnConfig
from alloy_runtime_types.dtos.generation import (
    GenerateTextResponse,
    GenerationFullTrace,
)


class GeneratedTextOutput(BaseModel):
    output_text: str


def _format_optional_text(value: str | None) -> Any:
    if value is None:
        return cast(str, None)
    return value


def _format_tags(value: list[str]) -> Any:
    if not value:
        return cast(str, None)
    return ", ".join(value)


def _format_total_tokens(response: GenerateTextResponse) -> str:
    return str(response.usage.input_tokens + response.usage.output_tokens)


def _format_reasoning_tokens(usage: Any) -> Any:
    reasoning_tokens = getattr(usage, "reasoning_tokens", None)
    if reasoning_tokens is None:
        return cast(str, None)
    return str(reasoning_tokens)


def _format_content_parts(parts: list[Any]) -> str:
    if not parts:
        return "-"
    first_part: Any = parts[0]
    return getattr(first_part, "content_text", None) or "-"


class GenerateTextPresenter:
    def __init__(self, output: OutputService):
        self.output = output

    def present_result(
        self,
        response: GenerateTextResponse,
        show_metadata: bool = True,
        *,
        metadata_title: str = "Generation Result",
        output_title: str = "Generated Text",
    ) -> None:
        if getattr(self.output, "format", OutputFormat.COMPACT) == OutputFormat.JSON:
            self.output.raw(response)
            return

        if show_metadata:
            self.output.entity(
                response,
                metadata_title,
                [
                    FieldConfig("execution_id", "Execution ID", str),
                    FieldConfig("chat_session_id", "Session ID", str),
                    FieldConfig("assistant_message_id", "Assistant Message ID", str),
                    FieldConfig(
                        "primary_content_part_id", "Primary Content Part ID", str
                    ),
                    FieldConfig("provider_key", "Provider"),
                    FieldConfig("provider_model_name", "Model"),
                    FieldConfig("finish_reason", "Finish Reason"),
                    FieldConfig("output_incomplete", "Output Incomplete"),
                    FieldConfig(
                        "usage", "Input Tokens", lambda usage: str(usage.input_tokens)
                    ),
                    FieldConfig(
                        "usage", "Output Tokens", lambda usage: str(usage.output_tokens)
                    ),
                    FieldConfig(
                        "usage",
                        "Total Tokens",
                        lambda _: _format_total_tokens(response),
                    ),
                    FieldConfig("usage", "Reasoning Tokens", _format_reasoning_tokens),
                    FieldConfig(
                        "cost", "Input Cost", lambda cost: f"${cost.input_cost_usd:.6f}"
                    ),
                    FieldConfig(
                        "cost",
                        "Output Cost",
                        lambda cost: f"${cost.output_cost_usd:.6f}",
                    ),
                    FieldConfig(
                        "cost", "Total Cost", lambda cost: f"${cost.total_cost_usd:.6f}"
                    ),
                    FieldConfig("tags", "Tags", _format_tags),
                    FieldConfig("external_id", "External ID", _format_optional_text),
                    FieldConfig("used_fallback", "Used Fallback"),
                    FieldConfig(
                        "actual_provider_key", "Actual Provider", _format_optional_text
                    ),
                    FieldConfig(
                        "actual_provider_model_name",
                        "Actual Model",
                        _format_optional_text,
                    ),
                ],
                raw_payload=response,
            )

        self.present_text_output(response.output_text, title=output_title)
        if response.full_trace is not None:
            self.present_full_trace(response.full_trace)

    def present_multiple_results(
        self, responses: list[GenerateTextResponse], show_metadata: bool = True
    ) -> None:
        if getattr(self.output, "format", OutputFormat.COMPACT) == OutputFormat.JSON:
            self.output.raw(responses)
            return

        total = len(responses)
        for index, response in enumerate(responses, start=1):
            suffix = f"{index}/{total}"
            self.present_result(
                response,
                show_metadata=show_metadata,
                metadata_title=f"Generation Result {suffix}",
                output_title=f"Generated Text {suffix}",
            )

    def present_text_output(
        self, content: str, *, title: str = "Generated Text"
    ) -> None:
        output = GeneratedTextOutput(output_text=content)
        self.output.entity(
            output,
            title,
            [FieldConfig("output_text", "Output")],
            raw_payload=output,
        )

    def present_full_trace(self, trace: GenerationFullTrace) -> None:
        self.output.table(
            items=cast(list[BaseModel], trace.messages),
            title=f"Full Trace ({trace.finish_reason})",
            columns=[
                ColumnConfig(
                    source_field="sequence_number",
                    header_label="#",
                    align="right",
                    compact_line=1,
                    compact_style="cyan",
                ),
                ColumnConfig(
                    source_field="role",
                    header_label="Role",
                    compact_line=1,
                    compact_style="yellow",
                ),
                ColumnConfig(
                    source_field="created_at",
                    header_label="Time",
                    formatter=lambda value: (
                        value.strftime("%H:%M:%S") if value else "-"
                    ),
                    compact_line=1,
                    compact_style="blue",
                ),
                ColumnConfig(
                    source_field="content_parts",
                    header_label="Content",
                    formatter=_format_content_parts,
                    compact_line=2,
                    compact_style="dim",
                ),
            ],
            raw_payload=trace,
            total_count=len(trace.messages),
        )
