"""
Development utility for subscribing to async task messages with streaming support.

This module provides utilities to read existing messages from a task and subscribe
to new streaming messages, handling mid-stream connections gracefully.

Uses TextStreamPart format (v2) for streaming.
"""

from __future__ import annotations

import json
from collections.abc import Iterator
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, List, Optional

from rich.console import Console
from rich.markdown import Markdown
from rich.markup import escape
from rich.panel import Panel
from yaspin import yaspin  # type: ignore[import-untyped]
from yaspin.core import Yaspin  # type: ignore[import-untyped]

from terminaluse.lib.client import TerminalUse
from terminaluse.lib.utils.output.defaults import CLI_TOOL_RESPONSE_FILTERS
from terminaluse.lib.utils.output.extractors import extract_text
from terminaluse.lib.utils.output.formatters import clean_lines
from terminaluse.lib.utils.output.renderers import ToolOutputData, ToolRequestExtractor, get_renderers
from terminaluse.types import Task, TaskResponse, UiMessage
from terminaluse.types.text_stream_part_wrapper import (
    TextStreamPartWrapper_Error,
    TextStreamPartWrapper_Finish,
    TextStreamPartWrapper_FinishStep,
    TextStreamPartWrapper_ReasoningDelta,
    TextStreamPartWrapper_ReasoningEnd,
    TextStreamPartWrapper_ReasoningStart,
    TextStreamPartWrapper_Start,
    TextStreamPartWrapper_StartStep,
    TextStreamPartWrapper_TextDelta,
    TextStreamPartWrapper_TextEnd,
    TextStreamPartWrapper_TextStart,
    TextStreamPartWrapper_ToolCall,
    TextStreamPartWrapper_ToolInputDelta,
    TextStreamPartWrapper_ToolInputEnd,
    TextStreamPartWrapper_ToolInputStart,
    TextStreamPartWrapper_ToolResult,
)
from terminaluse.types.ui_message_part_wrapper import (
    UiMessagePartWrapper_Error,
    UiMessagePartWrapper_Reasoning,
    UiMessagePartWrapper_Text,
    UiMessagePartWrapper_Tool,
)


def print_ui_message(
    message: UiMessage,
    print_messages: bool = True,
    rich_print: bool = True,
) -> None:
    """
    Print a UIMessage (v2 format) in a formatted way.

    Args:
        message: The UiMessage to print
        print_messages: Whether to actually print the message (for debugging)
        rich_print: Whether to use rich to print the message
    """
    if not print_messages:
        return

    # Skip messages with no parts
    if not message.parts:
        return

    timestamp = message.created_at or "N/A"

    console = None
    if rich_print:
        console = Console(width=80)  # Fit better in Jupyter cells

    # Process each part
    for part in message.parts:
        content = ""
        content_type = "other"

        if isinstance(part, UiMessagePartWrapper_Text):
            content = part.text or ""
            content_type = "text"
            if not content.strip():
                continue

        elif isinstance(part, UiMessagePartWrapper_Tool):
            tool_name = part.tool_name or "unknown"
            state = getattr(part, "state", "input-available")

            if state in ("input-streaming", "input-available"):
                # Tool request
                tool_input = getattr(part, "input", {})
                try:
                    formatted_args = json.dumps(tool_input, indent=2)
                    content = f"**Tool Request: {tool_name}**\n\n**Arguments:**\n```json\n{formatted_args}\n```"
                except (json.JSONDecodeError, TypeError):
                    content = f"**Tool Request: {tool_name}**\n\n**Arguments:**\n{tool_input}"
                content_type = "tool_request"
            else:
                # Tool response
                tool_output = getattr(part, "output", None)
                try:
                    if tool_output is not None:
                        formatted_output = json.dumps(tool_output, indent=2)
                        content = f"**Tool Response: {tool_name}**\n\n**Response:**\n```json\n{formatted_output}\n```"
                    else:
                        content = f"**Tool Response: {tool_name}**\n\n(no output)"
                except (json.JSONDecodeError, TypeError):
                    content = f"**Tool Response: {tool_name}**\n\n{tool_output}"
                content_type = "tool_response"

        elif isinstance(part, UiMessagePartWrapper_Reasoning):
            reasoning_text = part.text or ""
            if not reasoning_text.strip():
                continue
            content = f"**Reasoning**\n\n{reasoning_text}"
            content_type = "reasoning"

        elif isinstance(part, UiMessagePartWrapper_Error):
            error_text = part.error or "Unknown error"
            content = f"**Error**\n\n{error_text}"
            content_type = "error"

        else:
            content = f"{type(part).__name__}"
            content_type = "other"

        if not content:
            continue

        if rich_print and console:
            author_color = "bright_cyan" if message.role == "user" else "green"

            # Use different border styles and colors for different content types
            if content_type == "tool_request":
                border_style = "yellow"
            elif content_type == "tool_response":
                border_style = "bright_green"
            elif content_type == "reasoning":
                border_style = "bright_magenta"
                author_color = "bright_magenta"
            elif content_type == "error":
                border_style = "red"
            else:
                border_style = author_color

            title = f"[bold {author_color}]{str(message.role).upper()}[/bold {author_color}] [{timestamp}]"
            panel = Panel(Markdown(content), title=title, border_style=border_style, width=80)
            console.print(panel)
        else:
            title = f"{str(message.role).upper()} [{timestamp}]"
            print(f"{title}\n{content}\n")


# Alias for backward compatibility
def print_task_message(
    message: UiMessage,
    print_messages: bool = True,
    rich_print: bool = True,
) -> None:
    """Alias for print_ui_message for backward compatibility."""
    print_ui_message(message, print_messages, rich_print)


def subscribe_to_async_task_messages(
    client: TerminalUse,
    task: Task,
    only_after_timestamp: Optional[datetime] = None,
    print_messages: bool = True,
    rich_print: bool = True,
    timeout: int = 10,
) -> List[UiMessage]:
    """
    Subscribe to async task messages and collect completed messages.

    This function:
    1. Reads all existing messages from the task
    2. Optionally filters messages after a timestamp
    3. Shows a loading message while listening
    4. Subscribes to task message events
    5. Fetches and displays complete messages when they finish
    6. Returns all messages collected during the session

    Features:
    - Uses Rich library for beautiful formatting in Jupyter notebooks
    - Agent messages are formatted as Markdown
    - User and agent messages are displayed in colored panels with fixed width
    - Optimized for Jupyter notebook display

    Args:
        client: The TerminalUse client instance
        task: The task to subscribe to
        print_messages: Whether to print messages as they arrive
        only_after_timestamp: Only include messages created after this timestamp. If None, all messages will be included.
        rich_print: Whether to use rich to print the message
        timeout: The timeout in seconds for the streaming connection. If the connection times out, the function will return with any messages collected so far.
    Returns:
        List of UiMessage objects collected during the session

    Raises:
        ValueError: If the task doesn't have a name (required for streaming)
    """

    messages_to_return: List[UiMessage] = []

    # Read existing messages using v2 API
    messages: List[UiMessage] = []
    try:
        response = client.messages_v2.list(task_id=task.id)
        messages = list(response.data) if response.data else []
    except Exception as e:
        print(f"Error reading existing messages: {e}")

    # Filter and display existing messages
    for message in messages:
        if only_after_timestamp:
            if message.created_at is not None:
                # Handle timezone comparison
                try:
                    message_time = datetime.fromisoformat(message.created_at.replace("Z", "+00:00"))
                    comparison_time = only_after_timestamp
                    if comparison_time.tzinfo is None:
                        comparison_time = comparison_time.replace(tzinfo=timezone.utc)
                    if message_time < comparison_time:
                        continue
                except (ValueError, AttributeError):
                    pass
            messages_to_return.append(message)
            print_task_message(message, print_messages, rich_print)
        else:
            messages_to_return.append(message)
            print_task_message(message, print_messages, rich_print)

    # Subscribe to server-side events using tasks.stream_by_name or stream
    # This is the proper way to get agent responses after sending an event in async agents

    # Ensure task has a name or id for streaming
    if not task.name and not task.id:
        print("Error: Task must have either name or id for streaming")
        raise ValueError("Task name or id is required")

    try:
        # Use stream_by_name (preferred) or fall back to stream (by id)
        # SDK returns Iterator[dict] with TextStreamPart events

        # Track spinner state
        spinner: Yaspin | None = None

        # Track pending tool calls - increment on tool-input-start, decrement on tool-result
        # Stream closes on 'finish' only when pending_tool_calls == 0
        pending_tool_calls: int = 0

        # Choose streaming method based on available task identifier
        stream_iter = (
            client.tasks.stream_by_name(task_name=task.name) if task.name else client.tasks.stream(task_id=task.id)
        )

        try:
            for event in stream_iter:
                # SDK returns typed TextStreamPartWrapper objects directly

                # Handle different event types
                if isinstance(event, TextStreamPartWrapper_Start):
                    # Start spinner when session begins
                    if print_messages and spinner is None:
                        spinner = yaspin(text="🔄 Agent responding...")
                        spinner.start()

                elif isinstance(event, TextStreamPartWrapper_StartStep):
                    # Ensure spinner is running for new step
                    if print_messages and spinner is None:
                        spinner = yaspin(text="🔄 Agent responding...")
                        spinner.start()

                elif isinstance(event, TextStreamPartWrapper_ToolInputStart):
                    # Increment pending tool calls on tool input start
                    pending_tool_calls += 1

                elif isinstance(event, TextStreamPartWrapper_ToolResult):
                    # Decrement pending tool calls on tool result
                    pending_tool_calls = max(0, pending_tool_calls - 1)

                elif isinstance(event, TextStreamPartWrapper_Finish):
                    # Stop spinner
                    if spinner is not None:
                        spinner.stop()
                        spinner = None
                        if print_messages:
                            print()

                    # Exit the streaming loop once finish is received
                    # AND there are no pending tool calls
                    if pending_tool_calls == 0:
                        # Fetch all messages again to get the complete conversation
                        try:
                            response = client.messages_v2.list(task_id=task.id)
                            final_messages = list(response.data) if response.data else []
                            # Add any new messages we haven't seen
                            existing_ids = {m.id for m in messages_to_return}
                            for msg in final_messages:
                                if msg.id not in existing_ids:
                                    messages_to_return.append(msg)
                                    print_task_message(msg, print_messages, rich_print)
                        except Exception:
                            pass
                        break

                elif isinstance(event, TextStreamPartWrapper_Error):
                    if print_messages:
                        print(f"Error: {event.error}")

        finally:
            # Stop spinner if still running
            if spinner is not None:
                spinner.stop()

    except Exception as e:
        # Handle timeout gracefully
        if "timeout" in str(e).lower() or "timed out" in str(e).lower():
            if print_messages:
                print(f"Streaming timed out after {timeout} seconds - returning collected messages")
        else:
            if print_messages:
                print(f"Error subscribing to events: {e}")
                print("Make sure your agent is running and the task exists")

    return messages_to_return


# =============================================================================
# Streaming function: stream_task_events (v2 TextStreamPart format)
# =============================================================================


@dataclass
class StreamToolCallBuffer:
    """Buffer for accumulating tool call data during v2 streaming."""

    tool_call_id: str
    tool_name: str
    input_buffer: str = ""
    output: Any = None
    input_complete: bool = False
    input_printed: bool = False
    output_printed: bool = False


@dataclass
class StreamState:
    """State for tracking v2 streaming session."""

    text_buffers: dict[str, str] = field(default_factory=dict)
    reasoning_buffers: dict[str, str] = field(default_factory=dict)
    tool_calls: dict[str, StreamToolCallBuffer] = field(default_factory=dict)
    is_finished: bool = False
    _output_parts: list[str] = field(default_factory=list)  # O(1) appends for Live display

    def append_output(self, text: str) -> None:
        """Append text to output buffer efficiently (O(1) amortized)."""
        if text:
            self._output_parts.append(text)

    def get_output(self) -> str:
        """Get accumulated output by joining parts (called at render time)."""
        return "".join(self._output_parts)


def stream_task_events(
    client: TerminalUse,
    task: Task | TaskResponse,
    timeout: int = 600,
    debug: bool = False,
    verbose: bool = False,
) -> None:
    """
    Stream task events using the v2 TextStreamPart format.

    This function provides real-time streaming output using the new messaging v2 format:
    - text-start/text-delta/text-end: Text content streaming
    - reasoning-start/reasoning-delta/reasoning-end: Reasoning content streaming
    - tool-input-start/tool-input-delta/tool-input-end: Tool input streaming
    - tool-call: Complete tool call with parsed input
    - tool-result: Tool execution result
    - start/start-step/finish-step/finish: Session lifecycle events
    - error: Error events

    Args:
        client: The TerminalUse client instance
        task: The task to stream events for
        timeout: Connection timeout in seconds (default 600 = 10 minutes)
        debug: If True, print raw SSE events for debugging
        verbose: If True, show full tool arguments and responses without truncation

    Returns:
        None (display-only function)

    Raises:
        ValueError: If task has neither name nor id

    Note:
        Non-TTY Behavior: When stdout is not a terminal (pipes, CI environments),
        Rich Live gracefully degrades to print only the final output. No special
        handling is required.
    """
    console = Console()
    state = StreamState()
    pending_tool_call_ids: set[str] = set()

    if not task.id:
        raise ValueError("Task id is required")

    if debug:
        print(f"DEBUG: Streaming task_id: {task.id}", flush=True)  # noqa: T201

    # Unified renderer architecture: get_renderers() returns TTY-aware renderers
    # RichToolCallRenderer/RichToolOutputRenderer for TTY, PlainToolCallRenderer/PlainToolOutputRenderer for non-TTY
    call_renderer, output_renderer = get_renderers(console, verbose=True)
    tool_extractor = ToolRequestExtractor()
    is_tty = console.is_terminal

    try:
        if debug:
            print("DEBUG: Connected to v2 stream", flush=True)  # noqa: T201

        def _process_events() -> Iterator[tuple[str, str]]:
            """Process streaming events (shared logic for TTY and non-TTY)."""
            nonlocal pending_tool_call_ids

            for event in client.tasks.stream(task_id=task.id):
                if debug:
                    print(f"DEBUG: type={type(event).__name__}", flush=True)  # noqa: T201

                if isinstance(event, TextStreamPartWrapper_Start):
                    pass

                elif isinstance(event, TextStreamPartWrapper_StartStep):
                    pass

                elif isinstance(event, TextStreamPartWrapper_TextStart):
                    state.text_buffers[event.id] = ""

                elif isinstance(event, TextStreamPartWrapper_TextDelta):
                    if event.id in state.text_buffers:
                        state.text_buffers[event.id] += event.text
                    state.append_output(event.text)
                    yield ("text", event.text)

                elif isinstance(event, TextStreamPartWrapper_TextEnd):
                    pass

                elif isinstance(event, TextStreamPartWrapper_ReasoningStart):
                    state.reasoning_buffers[event.id] = ""

                elif isinstance(event, TextStreamPartWrapper_ReasoningDelta):
                    if event.id in state.reasoning_buffers:
                        state.reasoning_buffers[event.id] += event.text
                    if is_tty:
                        state.append_output(f"[dim italic]{event.text}[/dim italic]")
                    else:
                        state.append_output(event.text)
                    yield ("reasoning", event.text)

                elif isinstance(event, TextStreamPartWrapper_ReasoningEnd):
                    pass

                elif isinstance(event, TextStreamPartWrapper_ToolInputStart):
                    tool_id = event.id
                    state.tool_calls[tool_id] = StreamToolCallBuffer(
                        tool_call_id=tool_id,
                        tool_name=event.tool_name,
                    )
                    pending_tool_call_ids.add(tool_id)
                    if debug:
                        print(f"DEBUG: Tool input start: {tool_id} ({event.tool_name})", flush=True)  # noqa: T201

                elif isinstance(event, TextStreamPartWrapper_ToolInputDelta):
                    tool_id = event.id
                    if tool_id in state.tool_calls:
                        state.tool_calls[tool_id].input_buffer += event.delta

                elif isinstance(event, TextStreamPartWrapper_ToolInputEnd):
                    tool_id = event.id
                    if tool_id in state.tool_calls:
                        state.tool_calls[tool_id].input_complete = True

                elif isinstance(event, TextStreamPartWrapper_ToolCall):
                    tool_id = event.tool_call_id
                    if tool_id not in state.tool_calls:
                        state.tool_calls[tool_id] = StreamToolCallBuffer(
                            tool_call_id=tool_id,
                            tool_name=event.tool_name,
                        )
                        pending_tool_call_ids.add(tool_id)
                    buf = state.tool_calls[tool_id]
                    buf.input_buffer = json.dumps(event.input)
                    buf.input_complete = True
                    if not buf.input_printed:
                        # Unified path: extractor creates structured data, renderer handles TTY-aware output
                        tool_data = tool_extractor.extract(buf.tool_name, buf.input_buffer)
                        rendered = call_renderer.render(tool_data)
                        state.append_output(rendered)
                        buf.input_printed = True
                        yield ("tool_call", rendered)

                elif isinstance(event, TextStreamPartWrapper_ToolResult):
                    tool_id = event.tool_call_id
                    if tool_id in state.tool_calls:
                        buf = state.tool_calls[tool_id]
                        buf.output = event.output
                        if not buf.input_printed:
                            # Unified path: extractor creates structured data, renderer handles TTY-aware output
                            tool_data = tool_extractor.extract(buf.tool_name, buf.input_buffer)
                            rendered = call_renderer.render(tool_data)
                            state.append_output(rendered)
                            buf.input_printed = True
                            yield ("tool_call", rendered)
                        if not buf.output_printed:
                            # Unified path: extract → filter → clean → render
                            raw_text = extract_text(buf.output)
                            filtered_text = CLI_TOOL_RESPONSE_FILTERS.apply(raw_text)
                            lines = clean_lines(filtered_text)
                            output_data = ToolOutputData(
                                lines=lines,
                                total_line_count=len(lines),
                                tool_name=buf.tool_name,
                            )
                            rendered = output_renderer.render(output_data)
                            state.append_output(rendered)
                            buf.output_printed = True
                            yield ("tool_result", rendered)
                        pending_tool_call_ids.discard(tool_id)
                        if debug:
                            print(f"DEBUG: Tool result received: {tool_id}", flush=True)  # noqa: T201

                elif isinstance(event, TextStreamPartWrapper_FinishStep):
                    if debug:
                        print(f"DEBUG: Step finished: {event.finish_reason}", flush=True)  # noqa: T201

                elif isinstance(event, TextStreamPartWrapper_Finish):
                    state.is_finished = True
                    if len(pending_tool_call_ids) == 0:
                        if debug:
                            print("DEBUG: Session finished, closing stream", flush=True)  # noqa: T201
                        return
                    else:
                        if debug:
                            print(f"DEBUG: Waiting for {len(pending_tool_call_ids)} pending tool(s)", flush=True)  # noqa: T201

                elif isinstance(event, TextStreamPartWrapper_Error):
                    if is_tty:
                        # Escape brackets in error message to prevent Rich markup injection
                        error_text = f"\n[red]Error: {escape(event.error)}[/red]"
                    else:
                        error_text = f"\nError: {event.error}"
                    state.append_output(error_text)
                    yield ("error", error_text)

        # Print agent header
        if is_tty:
            console.print("\n[blue][Agent][/blue]")
        else:
            print("\n[Agent]", flush=True)  # noqa: T201

        # Stream output directly for scrollable history (both TTY and non-TTY)
        for event_type, content in _process_events():
            if event_type == "text" or event_type == "reasoning":
                # Stream text character-by-character
                print(content, end="", flush=True)  # noqa: T201
            elif is_tty:
                # Tool calls, results, errors in TTY - use console.print for Rich markup
                console.print(content)
            else:
                # Tool calls, results, errors in non-TTY - print as plain text
                print(content, flush=True)  # noqa: T201
        # Final newline
        print(flush=True)  # noqa: T201

        # Print agent footer
        if is_tty:
            console.print("[green]Agent finished.[/green]")
        else:
            print("Agent finished.", flush=True)  # noqa: T201

    except Exception as e:
        if "timeout" in str(e).lower():
            if is_tty:
                console.print("\n[yellow]Connection timed out[/yellow]")
            else:
                print("\nConnection timed out", flush=True)  # noqa: T201
        else:
            if is_tty:
                console.print(f"\n[red]Disconnected: {escape(str(e))}[/red]")
            else:
                print(f"\nDisconnected: {e}", flush=True)  # noqa: T201
