import json
import requests
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


class StreamResponseHandler:
    """
    A utility class for handling streaming responses from API endpoints.
    Provides rich formatting and real-time updates of the response content.
    """

    def __init__(self, console=None):
        """
        Initialize the stream response handler.

        Args:
            console (Console, optional): A Rich console instance. If not provided, a new one will be created.
        """
        self.console = console or Console()

    def check_server_health(self, health_url="http://localhost:8000/health"):
        """
        Check if the server is running and accessible.

        Args:
            health_url (str, optional): The URL to check server health. Defaults to "http://localhost:8000/health".

        Returns:
            bool: True if the server is running and accessible, False otherwise.
        """
        try:
            self.console.print("Checking server health...", style="bold yellow")
            response = requests.get(health_url)
            if response.status_code == 200:
                self.console.print("[bold green]✓ Server is running and accessible.[/]")
                return True
            else:
                self.console.print(
                    f"[bold red]✗ Server health check failed[/] with status code: {response.status_code}"
                )
                return False
        except requests.exceptions.ConnectionError:
            self.console.print(
                "[bold red]✗ Error:[/] Could not connect to the server. Make sure it's running."
            )
            return False
        except Exception as e:
            self.console.print(f"[bold red]✗ Error checking server health:[/] {e}")
            return False

    def stream_response(
        self, url, payload=None, params=None, method="POST", title="AI Response"
    ):
        """
        Send a request to an endpoint and stream the output to the terminal.

        Args:
            url (str): The URL of the endpoint to send the request to.
            payload (dict, optional): The JSON payload to send in the request body. Defaults to None.
            params (dict, optional): The query parameters to send in the request. Defaults to None.
            method (str, optional): The HTTP method to use. Defaults to "POST".
            title (str, optional): The title to display in the panel. Defaults to "AI Response".

        Returns:
            bool: True if the streaming was successful, False otherwise.
        """
        # Display request information
        self.console.print(f"Sending request to [bold cyan]{url}[/]")
        if payload:
            self.console.print("Payload:", style="bold")
            self.console.print(json.dumps(payload, indent=2))
        if params:
            self.console.print("Parameters:", style="bold")
            self.console.print(json.dumps(params, indent=2))

        try:
            # Prepare the request
            request_kwargs = {"stream": True}
            if payload:
                request_kwargs["json"] = payload
            if params:
                request_kwargs["params"] = params

            # Make the request
            with getattr(requests, method.lower())(url, **request_kwargs) as response:
                # Check if the request was successful
                if response.status_code != 200:
                    self.console.print(
                        f"[bold red]Error:[/] Received status code {response.status_code}"
                    )
                    self.console.print(f"Response: {response.text}")
                    return False

                # Initialize an empty response text
                full_response = ""

                # Use Rich's Live display to update the content in place
                with Live(
                    Panel("Waiting for response...", title=title, border_style="blue"),
                    refresh_per_second=10,
                ) as live:
                    # Process the streaming response
                    for line in response.iter_lines():
                        if line:
                            # Decode the line and parse it as JSON
                            decoded_line = line.decode("utf-8")
                            try:
                                # Parse the JSON
                                data = json.loads(decoded_line)

                                # Extract and display the content
                                if isinstance(data, dict):
                                    if "content" in data:
                                        for content in data["content"]:
                                            if content.get("type") == "text":
                                                text_content = content.get("text", "")
                                                # Append to the full response
                                                full_response += text_content
                                                # Update the live display with the current full response
                                                live.update(
                                                    Panel(
                                                        Markdown(full_response),
                                                        title=title,
                                                        border_style="green",
                                                    )
                                                )
                                            elif content.get("type") == "image_url":
                                                image_url = content.get(
                                                    "image_url", {}
                                                ).get("url", "")
                                                # Add a note about the image URL
                                                image_note = (
                                                    f"\n\n[Image URL: {image_url}]"
                                                )
                                                full_response += image_note
                                                live.update(
                                                    Panel(
                                                        Markdown(full_response),
                                                        title=title,
                                                        border_style="green",
                                                    )
                                                )
                                    elif "edited_image_url" in data:
                                        # Handle edited image URL from edit endpoint
                                        image_url = data.get("edited_image_url", "")
                                        image_note = (
                                            f"\n\n[Edited Image URL: {image_url}]"
                                        )
                                        full_response += image_note
                                        live.update(
                                            Panel(
                                                Markdown(full_response),
                                                title=title,
                                                border_style="green",
                                            )
                                        )
                                    else:
                                        # For other types of data, just show the JSON
                                        live.update(
                                            Panel(
                                                Text(json.dumps(data, indent=2)),
                                                title="Raw JSON Response",
                                                border_style="yellow",
                                            )
                                        )
                                else:
                                    live.update(
                                        Panel(
                                            Text(decoded_line),
                                            title="Raw Response",
                                            border_style="yellow",
                                        )
                                    )
                            except json.JSONDecodeError:
                                # If it's not valid JSON, just show the raw line
                                live.update(
                                    Panel(
                                        Text(f"Raw response: {decoded_line}"),
                                        title="Invalid JSON",
                                        border_style="red",
                                    )
                                )

                self.console.print("[bold green]Stream completed.[/]")
                return True

        except requests.exceptions.ConnectionError:
            self.console.print(
                f"[bold red]Error:[/] Could not connect to the server at {url}",
                style="red",
            )
            self.console.print(
                "Make sure the server is running and accessible.", style="red"
            )
            return False
        except requests.exceptions.RequestException as e:
            self.console.print(f"[bold red]Error:[/] {e}", style="red")
            return False