from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.panel import Panel
from rich.text import Text

from pyra.security.injection import redact_api_keys

console = Console()


def render_streaming_response(stream) -> str:
    """Consume a litellm streaming response, render markdown progressively, return full text."""
    full_text = ""
    with Live(console=console, refresh_per_second=8) as live:
        for chunk in stream:
            delta = chunk.choices[0].delta.content or ""
            full_text += delta
            safe_text = redact_api_keys(full_text)
            live.update(Markdown(safe_text))

    return redact_api_keys(full_text)


def render_text_response(text: str) -> str:
    """Render a complete (non-streaming) AI response as markdown. Returns redacted text."""
    safe_text = redact_api_keys(text)
    if safe_text.strip():
        console.print(Markdown(safe_text))
    return safe_text


def render_injection_warning(warnings) -> None:
    labels = ", ".join(w.pattern_label for w in warnings)
    console.print(Panel(
        f"[yellow]Possible prompt injection detected[/yellow]\n"
        f"Pattern(s): [bold]{labels}[/bold]\n\n"
        "[dim]The response is shown, but treat it with caution.\n"
        "Details logged to ~/.pyra/security.log[/dim]",
        border_style="yellow",
        title="Security Warning",
    ))


def render_error(message: str) -> None:
    console.print(Panel(f"[red]{message}[/red]", border_style="red"))


def render_info(message: str) -> None:
    console.print(f"[dim]{message}[/dim]")


def render_system(message: str) -> None:
    console.print(Panel(message, border_style="cyan"))