import json as json_mod
from typing import Optional

import typer
from rich.panel import Panel
from rich.table import Table
from rich.tree import Tree

from ..entities.agent import AgentRepository
from ..entities.agent_session import AgentSessionRepository
from ..utils.config import (
    TITLE_STYLE,
    INFO_STYLE,
    ERROR_STYLE,
    ID_STYLE,
    NOT_AVAILABLE,
)
from ..utils.utils import console, print_table, print_json, format_duration_ms, should_output_json

app = typer.Typer(help="View sessions and session data.")


def print_section(title: str, items: list[tuple[str, str]]):
    """Print a section with key-value pairs."""
    table = Table(show_header=False, box=None, padding=(0, 2))
    table.add_column("Key", style=INFO_STYLE, width=35)
    table.add_column("Value")

    for key, value in items:
        table.add_row(key, value)

    console.print(Panel(table, title=f"[{TITLE_STYLE}]{title}[/{TITLE_STYLE}]", border_style="dim"))


@app.command("list")
def session_list(
    agent: str = typer.Argument(help="Agent name or ID"),
    status: Optional[str] = typer.Option(None, "--status", "-s", help="Filter by status"),
    direction: Optional[str] = typer.Option(None, "--direction", "-d", help="Filter by call direction"),
    limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Limit number of results"),
    page: Optional[int] = typer.Option(None, "--page", "-p", help="Page number"),
    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
    table_output: bool = typer.Option(False, "--table", "-t", help="Output as table (default for interactive terminals)"),
):
    """List sessions for an agent."""
    agent_repo = AgentRepository()
    agent_entity = agent_repo.get_by_name_or_id(agent)

    if not agent_entity:
        console.print(f"[{ERROR_STYLE}]Agent not found: {agent}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    repo = AgentSessionRepository(agent_entity.id)
    filters = {}
    if status:
        filters["status"] = status
    if direction:
        filters["direction"] = direction

    sessions = repo.get(filters if filters else None, page=page, limit=limit)

    if should_output_json(json_output, table_output):
        print_json([s.__dict__ for s in sessions])
        return

    if not sessions:
        console.print(f"[{INFO_STYLE}]No sessions found for agent '{agent_entity.name}'.[/{INFO_STYLE}]")
        return

    headers = ["STATUS", "ID", "DIRECTION", "ORIGIN", "DURATION", "CREATED AT"]
    rows = []
    for session in sessions:
        duration = NOT_AVAILABLE
        if session.usage and isinstance(session.usage, dict):
            duration = format_duration_ms(session.usage.get("sessionDuration"))
        rows.append([
            session.status or NOT_AVAILABLE,
            session.id or NOT_AVAILABLE,
            session.direction or NOT_AVAILABLE,
            session.origin or NOT_AVAILABLE,
            duration,
            session.created_at or NOT_AVAILABLE,
        ])
    print_table(headers, rows)


@app.command("info")
def session_info(
    agent: str = typer.Argument(help="Agent name or ID"),
    session_id: str = typer.Argument(help="Session ID"),
    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
    table_output: bool = typer.Option(False, "--table", "-t", help="Output as rich format (default for interactive terminals)"),
):
    """Show detailed information about a session."""
    agent_repo = AgentRepository()
    agent_entity = agent_repo.get_by_name_or_id(agent)

    if not agent_entity:
        console.print(f"[{ERROR_STYLE}]Agent not found: {agent}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    repo = AgentSessionRepository(agent_entity.id)
    session = repo.get_by_id(session_id)

    if not session:
        console.print(f"[{ERROR_STYLE}]Session not found: {session_id}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    if should_output_json(json_output, table_output):
        print_json(session.__dict__)
        return

    console.print()
    console.print(f"[{TITLE_STYLE}]Session: {session.id}[/{TITLE_STYLE}]")
    console.print(f"[{INFO_STYLE}]Agent: {agent_entity.name}[/{INFO_STYLE}]")
    console.print()

    print_section("Session Info", [
        ("ID", f"[{ID_STYLE}]{session.id}[/{ID_STYLE}]"),
        ("Status", session.status or NOT_AVAILABLE),
        ("Direction", session.direction or NOT_AVAILABLE),
        ("Origin", session.origin or NOT_AVAILABLE),
        ("Environment ID", session.agent_environment_id or NOT_AVAILABLE),
        ("Deployment ID", session.deployment_id or NOT_AVAILABLE),
        ("Telephony Call ID", session.telephony_call_id or NOT_AVAILABLE),
    ])

    usage = session.usage if isinstance(session.usage, dict) else {}
    print_section("Usage", [
        ("Session Duration", format_duration_ms(usage.get("sessionDuration"))),
        ("TTS Duration", format_duration_ms(usage.get("ttsDuration"))),
        ("Cached TTS Duration", format_duration_ms(usage.get("cachedTtsDuration"))),
        ("Module Duration", format_duration_ms(usage.get("moduleDuration"))),
    ])

    tags_str = ", ".join(session.tags) if session.tags else "None"
    print_section("Tags", [
        ("Tags", tags_str),
    ])

    print_section("Timestamps", [
        ("Created At", session.created_at or NOT_AVAILABLE),
        ("Updated At", session.updated_at or NOT_AVAILABLE),
    ])


@app.command("trace")
def session_trace(
    agent: str = typer.Argument(help="Agent name or ID"),
    session_id: str = typer.Argument(help="Session ID"),
    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
    table_output: bool = typer.Option(False, "--table", "-t", help="Output as tree format (default for interactive terminals)"),
):
    """Show trace overview for a session as a span tree."""
    agent_repo = AgentRepository()
    agent_entity = agent_repo.get_by_name_or_id(agent)

    if not agent_entity:
        console.print(f"[{ERROR_STYLE}]Agent not found: {agent}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    repo = AgentSessionRepository(agent_entity.id)
    trace_data = repo.get_trace(session_id)

    if not trace_data:
        console.print(f"[{ERROR_STYLE}]Trace not found for session: {session_id}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    if should_output_json(json_output, table_output):
        print_json(trace_data)
        return

    spans = trace_data.get("spans", [])
    trace_info = trace_data.get("trace", {})

    # Summary header
    total_spans = len(spans)
    error_count = sum(1 for s in spans if s.get("statusCode") == "ERROR")
    total_duration = format_duration_ms(trace_info.get("durationMs"))
    console.print(f"[{INFO_STYLE}]Trace for session {session_id}[/{INFO_STYLE}]")
    console.print(f"[{INFO_STYLE}]Spans: {total_spans}  Duration: {total_duration}  Errors: {error_count}[/{INFO_STYLE}]")
    console.print()

    # Build span tree
    spans_by_id = {s.get("spanId"): s for s in spans}
    children_map = {}
    roots = []
    for span in spans:
        parent_id = span.get("parentId")
        if parent_id and parent_id in spans_by_id:
            children_map.setdefault(parent_id, []).append(span)
        else:
            roots.append(span)

    # Sort by startTime
    roots.sort(key=lambda s: s.get("startTime", ""))
    for children in children_map.values():
        children.sort(key=lambda s: s.get("startTime", ""))

    def _span_label(span):
        span_id = span.get("spanId", "")
        name = span.get("name", "unknown")
        duration = format_duration_ms(span.get("durationMs"))
        status = span.get("statusCode", "UNSET")

        if status == "ERROR":
            status_str = f"  [bold red]{status}[/bold red]"
        elif status == "OK":
            status_str = f"  [green]{status}[/green]"
        else:
            status_str = ""

        label = f"{name}  [{ID_STYLE}]{span_id}[/{ID_STYLE}]  {duration}{status_str}"

        # Add preview text if available
        preview = span.get("preview")
        if preview:
            preview = preview[:60] + "..." if len(preview) > 60 else preview
            label += f"  [dim]{preview}[/dim]"

        return label

    def _add_children(tree_node, span_id):
        for child in children_map.get(span_id, []):
            child_node = tree_node.add(_span_label(child))
            _add_children(child_node, child.get("spanId"))

    tree = Tree(f"[bold]Trace[/bold]")
    for root_span in roots:
        root_node = tree.add(_span_label(root_span))
        _add_children(root_node, root_span.get("spanId"))

    console.print(tree)


@app.command("span")
def session_span(
    agent: str = typer.Argument(help="Agent name or ID"),
    session_id: str = typer.Argument(help="Session ID"),
    span_id: str = typer.Argument(help="Span ID"),
    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
    table_output: bool = typer.Option(False, "--table", "-t", help="Output as rich format (default for interactive terminals)"),
):
    """Show detailed information about a span."""
    agent_repo = AgentRepository()
    agent_entity = agent_repo.get_by_name_or_id(agent)

    if not agent_entity:
        console.print(f"[{ERROR_STYLE}]Agent not found: {agent}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    repo = AgentSessionRepository(agent_entity.id)
    span_data = repo.get_span(session_id, span_id)

    if not span_data:
        console.print(f"[{ERROR_STYLE}]Span not found: {span_id}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    if should_output_json(json_output, table_output):
        print_json(span_data)
        return

    console.print()
    console.print(f"[{TITLE_STYLE}]Span: {span_data.get('name', span_id)}[/{TITLE_STYLE}]")
    console.print(f"[{INFO_STYLE}]Agent: {agent_entity.name}[/{INFO_STYLE}]")
    console.print()

    print_section("Span Info", [
        ("Span ID", f"[{ID_STYLE}]{span_data.get('spanId', NOT_AVAILABLE)}[/{ID_STYLE}]"),
        ("Name", span_data.get("name", NOT_AVAILABLE)),
        ("Span Kind", span_data.get("spanKind", NOT_AVAILABLE)),
        ("Event Name", span_data.get("eventName", NOT_AVAILABLE)),
        ("Turn Number", str(span_data.get("turnNumber", NOT_AVAILABLE))),
        ("Parent ID", span_data.get("parentId", NOT_AVAILABLE)),
    ])

    print_section("Timing", [
        ("Start Time", span_data.get("startTime", NOT_AVAILABLE)),
        ("End Time", span_data.get("endTime", NOT_AVAILABLE)),
        ("Duration", format_duration_ms(span_data.get("durationMs"))),
    ])

    status_code = span_data.get("statusCode", NOT_AVAILABLE)
    status_message = span_data.get("statusMessage", NOT_AVAILABLE)
    print_section("Status", [
        ("Status Code", status_code),
        ("Status Message", status_message),
    ])

    attributes = span_data.get("attributes")
    if attributes:
        attrs_str = json_mod.dumps(attributes, indent=2, default=str)
        console.print(Panel(attrs_str, title=f"[{TITLE_STYLE}]Attributes[/{TITLE_STYLE}]", border_style="dim"))


@app.command("transcript")
def session_transcript(
    agent: str = typer.Argument(help="Agent name or ID"),
    session_id: str = typer.Argument(help="Session ID"),
    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
    table_output: bool = typer.Option(False, "--table", "-t", help="Output as conversation format (default for interactive terminals)"),
):
    """Show the transcript for a session."""
    agent_repo = AgentRepository()
    agent_entity = agent_repo.get_by_name_or_id(agent)

    if not agent_entity:
        console.print(f"[{ERROR_STYLE}]Agent not found: {agent}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    repo = AgentSessionRepository(agent_entity.id)
    events = repo.get_transcript(session_id)

    if events is None:
        console.print(f"[{ERROR_STYLE}]Transcript not found for session: {session_id}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    if should_output_json(json_output, table_output):
        print_json(events)
        return

    if not events:
        console.print(f"[{INFO_STYLE}]No transcript events for session '{session_id}'.[/{INFO_STYLE}]")
        return

    for event in events:
        timestamp = event.get("timestamp", "")
        role = event.get("role", "unknown")
        content = event.get("content", "")
        console.print(f"[dim]{timestamp}[/dim]  [{TITLE_STYLE}]{role}[/{TITLE_STYLE}]")
        console.print(f"  {content}")
        console.print()


@app.command("events")
def session_events(
    agent: str = typer.Argument(help="Agent name or ID"),
    session_id: str = typer.Argument(help="Session ID"),
    page: Optional[int] = typer.Option(None, "--page", "-p", help="Page number"),
    limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Limit number of results"),
    json_output: bool = typer.Option(False, "--json", "-j", help="Output as JSON"),
    table_output: bool = typer.Option(False, "--table", "-t", help="Output as table (default for interactive terminals)"),
):
    """Show raw events for a session."""
    agent_repo = AgentRepository()
    agent_entity = agent_repo.get_by_name_or_id(agent)

    if not agent_entity:
        console.print(f"[{ERROR_STYLE}]Agent not found: {agent}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    repo = AgentSessionRepository(agent_entity.id)
    result = repo.get_events(session_id, page=page, limit=limit)

    if result is None:
        console.print(f"[{ERROR_STYLE}]Events not found for session: {session_id}[/{ERROR_STYLE}]")
        raise typer.Exit(1)

    if should_output_json(json_output, table_output):
        print_json(result)
        return

    data = result.get("data", [])
    metadata = result.get("metadata", {})

    if not data:
        console.print(f"[{INFO_STYLE}]No events for session '{session_id}'.[/{INFO_STYLE}]")
        return

    headers = ["TIMESTAMP", "NAME", "DATA"]
    rows = []
    for event in data:
        timestamp = event.get("timestamp", NOT_AVAILABLE)
        name = event.get("name", NOT_AVAILABLE)
        event_data = event.get("data", "")
        if isinstance(event_data, dict):
            event_data = json_mod.dumps(event_data, default=str)
        if isinstance(event_data, str) and len(event_data) > 80:
            event_data = event_data[:77] + "..."
        rows.append([timestamp, name, str(event_data)])
    print_table(headers, rows)

    if metadata:
        total = metadata.get("total", "?")
        current_page = metadata.get("page", "?")
        page_limit = metadata.get("limit", "?")
        console.print(f"\n[{INFO_STYLE}]Page {current_page} | Limit {page_limit} | Total {total}[/{INFO_STYLE}]")


if __name__ == "__main__":
    app()
