import json
import textwrap
from dataclasses import dataclass
from datetime import datetime
from typing import Any, TypeGuard, cast

from alloy_runtime_types.dtos.executions import (
    AssistantTextActivity,
    ExecutionActivityItem,
    GetExecutionResponse,
    ReasoningActivity,
    ToolCallActivity,
    ToolErrorActivity,
    ToolOutputActivity,
)
from rich.markup import escape


@dataclass(frozen=True)
class KindDescriptor:
    label: str
    style: str


KIND_DESCRIPTORS: dict[str, KindDescriptor] = {
    "reasoning": KindDescriptor("REASONING", "dim cyan"),
    "assistant_text": KindDescriptor("ASSISTANT", "bold"),
    "tool_call": KindDescriptor("TOOL CALL", "magenta"),
    "tool_output": KindDescriptor("TOOL OUTPUT", "green"),
    "tool_error": KindDescriptor("TOOL ERROR", "bold red"),
}

RULE = "─" * 80


def render_execution_header(response: GetExecutionResponse, *, width: int = 80) -> str:
    execution_id = f"{str(response.execution_id)[:8]}…"
    duration = _format_duration(response.total_duration_ms)
    rounds = (
        _plural(response.round_count, "round")
        if response.round_count is not None
        else "0 rounds"
    )
    activity = _plural(len(response.activity), "activity item")
    lines = [
        "Execution  "
        f"{escape(execution_id)}  ·  "
        f"{_styled(escape(response.status), _status_style(response.status))}  ·  "
        f"{escape(duration)}  ·  {escape(rounds)}  ·  {escape(activity)}"
    ]

    provider = _provider_line(response)
    if provider:
        lines.append(provider)

    times = _times_line(response)
    if times:
        lines.append(times)

    if response.error_message:
        lines.append(f"Error      {_styled(escape(response.error_message), 'red')}")

    lines.append(_styled(escape(RULE[:width]), "dim"))
    return "\n".join(lines)


def render_activity_blocks(
    activities: list[ExecutionActivityItem],
    *,
    start_index: int = 0,
    full_output: bool,
    width: int = 80,
) -> str:
    blocks: list[str] = []
    activity_slice = activities[start_index:]
    grouped_results = _group_tool_results(activity_slice)
    rendered_result_indexes: set[int] = set()

    for activity in activity_slice:
        if isinstance(activity, ToolCallActivity):
            results = grouped_results.get(activity.tool_call_id or "", [])
            rendered_result_indexes.update(result.index for result in results)
            blocks.append(
                _render_tool_group(
                    activity,
                    results,
                    full_output=full_output,
                    width=width,
                )
            )
            continue
        if _is_tool_result(activity) and activity.index in rendered_result_indexes:
            continue
        blocks.append(_render_activity(activity, full_output=full_output, width=width))
    return "\n\n".join(block for block in blocks if block)


def render_execution_snapshot(
    response: GetExecutionResponse,
    *,
    include_metadata: bool = True,
    include_appendices: bool = True,
    start_index: int = 0,
    full_output: bool,
) -> str:
    blocks: list[str] = []
    if include_metadata:
        blocks.append(render_execution_header(response))

    activity_blocks = render_activity_blocks(
        response.activity, start_index=start_index, full_output=full_output
    )
    if activity_blocks:
        blocks.append(activity_blocks)
    elif include_metadata and response.activity_state == "unavailable":
        blocks.append("No live activity is available yet.")

    if include_appendices:
        appendices = render_appendices(response)
        if appendices:
            blocks.append(appendices)

    return "\n\n".join(blocks).rstrip()


def render_appendices(response: GetExecutionResponse) -> str:
    appendices: list[str] = []
    if response.output_text:
        appendices.append(f"Final Output\n{escape(response.output_text)}")
    return "\n\n".join(appendices)


def _render_activity(
    activity: ExecutionActivityItem, *, full_output: bool, width: int
) -> str:
    body = ""
    inline_meta: str | None = None
    tool_name: str | None = None

    if isinstance(activity, ReasoningActivity):
        body = _wrap_text(activity.text, width=width, indent="  ", style="dim")
    elif isinstance(activity, AssistantTextActivity):
        body = _wrap_text(activity.text, width=width, indent="  ")
    elif isinstance(activity, ToolCallActivity):
        tool_name = activity.tool_name
        body = _format_tool_input(activity.input, width=width)
    elif isinstance(activity, ToolOutputActivity):
        tool_name = activity.tool_name
        body, inline_meta = _dispatch_tool_output(
            activity.tool_name, activity.output, full=full_output, width=width
        )
    else:
        tool_name = activity.tool_name
        body = _wrap_text(activity.error, width=width, indent="  ", style="red")

    header = _activity_header(activity, tool_name=tool_name, inline_meta=inline_meta)
    if not body:
        return header
    return f"{header}\n{body}"


def _render_tool_group(
    call: ToolCallActivity,
    results: list[ToolOutputActivity | ToolErrorActivity],
    *,
    full_output: bool,
    width: int,
) -> str:
    blocks = [_render_activity(call, full_output=full_output, width=width)]
    for result in results:
        blocks.append(
            _indent_block(
                _render_activity(result, full_output=full_output, width=width)
            )
        )
    return "\n\n".join(blocks)


def _group_tool_results(
    activities: list[ExecutionActivityItem],
) -> dict[str, list[ToolOutputActivity | ToolErrorActivity]]:
    results: dict[str, list[ToolOutputActivity | ToolErrorActivity]] = {}
    for activity in activities:
        if not _is_tool_result(activity) or activity.tool_call_id is None:
            continue
        results.setdefault(activity.tool_call_id, []).append(activity)
    return results


def _is_tool_result(
    activity: ExecutionActivityItem,
) -> TypeGuard[ToolOutputActivity | ToolErrorActivity]:
    return isinstance(activity, ToolOutputActivity | ToolErrorActivity)


def _indent_block(block: str) -> str:
    return "\n".join(f"  {line}" if line else line for line in block.splitlines())


def _activity_header(
    activity: ExecutionActivityItem, *, tool_name: str | None, inline_meta: str | None
) -> str:
    descriptor = KIND_DESCRIPTORS[activity.kind]
    pieces = [
        _styled(escape(_format_time(activity.timestamp)), "dim"),
        _styled(descriptor.label, descriptor.style),
    ]
    if tool_name:
        pieces.append(_styled(escape(tool_name), "bold"))
    if inline_meta:
        pieces.append(_styled(escape(inline_meta), "dim"))
    return "  ".join(pieces)


def _format_tool_input(payload: dict[str, Any], *, width: int) -> str:
    if not payload:
        return ""
    key_width = min(max(len(str(key)) for key in payload), 24)
    lines: list[str] = []
    available = max(20, width - 2 - key_width - 2)
    for key, value in payload.items():
        key_text = str(key)
        rendered = _format_input_value(value)
        prefix = "  " + _styled(escape(key_text.ljust(key_width)), "cyan") + "  "
        if "\n" not in rendered and len(rendered) <= available:
            lines.append(prefix + escape(rendered))
            continue
        lines.append(prefix.rstrip())
        for line in rendered.splitlines():
            lines.append("  " + " " * (key_width + 2) + escape(line))
    return "\n".join(lines)


def _format_input_value(value: Any) -> str:
    if isinstance(value, str):
        return value
    compact = json.dumps(value, ensure_ascii=False, separators=(",", ":"))
    if len(compact) <= 80:
        return compact
    return json.dumps(value, ensure_ascii=False, indent=2)


def _dispatch_tool_output(
    tool_name: str, payload: dict[str, Any] | list[Any] | str, *, full: bool, width: int
) -> tuple[str, str | None]:
    if tool_name == "rag_search" and isinstance(payload, dict):
        return _format_rag_search(payload, full=full, width=width)
    return _format_generic(payload, full=full, width=width)


def _format_rag_search(
    payload: dict[str, Any], *, full: bool, width: int
) -> tuple[str, str | None]:
    persisted_output = payload.get("output")
    raw_metadata = payload.get("metadata")
    metadata = (
        cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
    )
    if isinstance(persisted_output, str):
        hits = _int_value(metadata.get("results_count"), default=None)
        if hits is None:
            hits = _int_value(metadata.get("total_candidates"), default=0)
        timing_payload = metadata.get("timing_ms")
        timing = (
            _int_value(
                cast(dict[str, Any], timing_payload).get("total_ms"), default=None
            )
            if isinstance(timing_payload, dict)
            else _int_value(metadata.get("timing_ms"), default=None)
        )
        inline_meta = f"{hits} hits" + (f" · {timing}ms" if timing is not None else "")
        body, _ = _format_generic(persisted_output, full=full, width=width)
        return body, inline_meta

    titles = _list_of_strings(payload.get("top_titles"))
    hits = _int_value(payload.get("hits"), default=len(titles))
    timing = _int_value(payload.get("timing_ms"), default=None)
    inline_meta = f"{hits} hits" + (f" · {timing}ms" if timing is not None else "")
    scores = _list(payload.get("scores"))
    source_refs = _list(payload.get("source_refs"))
    chunks = _list(payload.get("chunks"))

    lines: list[str] = []
    for index, title in enumerate(titles):
        if not full:
            lines.append(f"  • {escape(title)}")
            continue
        meta_parts: list[str] = []
        if index < len(scores):
            meta_parts.append(f"score {_score(scores[index])}")
        if index < len(source_refs):
            meta_parts.append(str(source_refs[index]))
        meta = (
            f"  {_styled(escape(' · '.join(meta_parts)), 'dim')}" if meta_parts else ""
        )
        lines.append(f"  • {escape(title)}{meta}")
        snippet = _chunk_text(chunks[index]) if index < len(chunks) else None
        if snippet:
            lines.extend(
                escape(line)
                for line in _wrap_plain(
                    snippet, width=width, indent="    ", max_lines=2
                )
            )

    if not full:
        lines.append("")
        lines.append(
            "  "
            + _styled(
                "Use --full-output to see scores, source refs, and content snippets.",
                "dim",
            )
        )
    return "\n".join(lines), inline_meta


def _format_generic(
    payload: dict[str, Any] | list[Any] | str, *, full: bool, width: int
) -> tuple[str, str | None]:
    raw = (
        payload
        if isinstance(payload, str)
        else json.dumps(payload, indent=2, ensure_ascii=False)
    )
    wrapped: list[str] = []
    for line in raw.splitlines() or [""]:
        wrapped.extend(_wrap_plain(line, width=width, indent="  "))
    total = len(wrapped)
    visible = wrapped if full or total <= 12 else wrapped[:12]
    escaped_lines = [escape(line) for line in visible]
    if not full and total > 12:
        escaped_lines.append(
            "  "
            + _styled(
                f"+ {total - 12} more lines — use --full-output to expand",
                "dim",
            )
        )
    return "\n".join(escaped_lines), _plural(total, "line")


def _wrap_text(text: str, *, width: int, indent: str, style: str | None = None) -> str:
    lines = [escape(line) for line in _wrap_plain(text, width=width, indent=indent)]
    if style:
        return "\n".join(_styled(line, style) for line in lines)
    return "\n".join(lines)


def _wrap_plain(
    text: str, *, width: int, indent: str, max_lines: int | None = None
) -> list[str]:
    lines: list[str] = []
    body_width = max(20, width - len(indent))
    for source_line in text.splitlines() or [""]:
        wrapped = textwrap.wrap(
            source_line,
            width=body_width,
            replace_whitespace=False,
            drop_whitespace=False,
        ) or [""]
        for line in wrapped:
            lines.append(indent + line)
            if max_lines is not None and len(lines) >= max_lines:
                return lines
    return lines


def _provider_line(response: GetExecutionResponse) -> str | None:
    if (
        not response.provider_key
        and not response.provider_model_name
        and not response.origin_label
    ):
        return None
    provider = (
        f"{response.provider_key} / {response.provider_model_name}"
        if response.provider_key or response.provider_model_name
        else ""
    )
    return f"Provider   {escape(provider)}  ·  origin: {escape(response.origin_label)}"


def _times_line(response: GetExecutionResponse) -> str | None:
    parts: list[str] = []
    include_completed_date = (
        isinstance(response.completed_at, datetime)
        and response.queued_at.date() != response.completed_at.date()
    )
    if response.queued_at:
        parts.append(f"queued {_format_time(response.queued_at)}")
    if response.started_at:
        parts.append(f"started {_format_time(response.started_at)}")
    if response.completed_at:
        parts.append(
            "completed "
            + (
                response.completed_at.strftime("%Y-%m-%d %H:%M:%S")
                if include_completed_date
                else _format_time(response.completed_at)
            )
        )
    if not parts:
        return None
    return "Times      " + escape("  ·  ".join(parts))


def _format_duration(total_duration_ms: int | None) -> str:
    if total_duration_ms is None:
        return "(running)"
    seconds = total_duration_ms / 1000
    return f"{seconds:.1f}s" if seconds < 60 else f"{seconds / 60:.1f}m"


def _format_time(value: Any) -> str:
    if not isinstance(value, datetime):
        return "--:--:--"
    return value.strftime("%H:%M:%S")


def _status_style(status: str) -> str:
    if status in {"failed", "cancelled", "timeout"}:
        return "bold red"
    if status in {"completed", "succeeded"}:
        return "green"
    return "yellow"


def _plural(count: int, noun: str) -> str:
    return f"{count} {noun}" if count == 1 else f"{count} {noun}s"


def _list(value: Any) -> list[Any]:
    return cast(list[Any], value) if isinstance(value, list) else []


def _list_of_strings(value: Any) -> list[str]:
    items = cast(list[object], value) if isinstance(value, list) else []
    return [str(item) for item in items]


def _int_value(value: Any, *, default: int | None) -> int | None:
    if isinstance(value, bool):
        return default
    if isinstance(value, int):
        return value
    return default


def _score(value: Any) -> str:
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return f"{float(value):.2f}"
    return str(value)


def _chunk_text(value: Any) -> str | None:
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        mapping = cast(dict[str, object], value)
        text = mapping.get("text")
        if isinstance(text, str):
            return text
    return None


def _styled(content: str, style: str) -> str:
    return _open_tag(style) + content + _close_tag(style)


def _open_tag(style: str) -> str:
    return "[" + style + "]"


def _close_tag(style: str) -> str:
    return "[/" + style + "]"
