"""Data classes for LLM Sandbox.

This module defines the data structures used to represent various types of output
and results from sandbox sessions, such as console output, execution results,
and extracted plots or other file artifacts.
"""

import json
import warnings
from dataclasses import dataclass, field
from enum import Enum


class FileType(Enum):
    r"""Enumeration of file types supported by artifact extractors.

    This enum lists common file formats that can be generated by code running
    in the sandbox and subsequently extracted, such as images, data files,
    and documents.
    """

    PNG = "png"
    JPEG = "jpeg"
    PDF = "pdf"
    SVG = "svg"
    CSV = "csv"
    JSON = "json"
    TXT = "txt"
    HTML = "html"


@dataclass(frozen=True)
class PlotOutput:
    r"""Represents a plot, chart, or other visual artifact output from code execution.

    Attributes:
        format (FileType): The format of the plot (e.g., PNG, SVG, PDF).
        content_base64 (str): The raw content of the plot, base64 encoded.
        width (int | None, optional): The width of the plot in pixels. Defaults to None.
        height (int | None, optional): The height of the plot in pixels. Defaults to None.
        dpi (int | None, optional): The dots per inch (resolution) of the plot. Defaults to None.

    """

    format: FileType
    content_base64: str
    width: int | None = None
    height: int | None = None
    dpi: int | None = None


@dataclass(frozen=True)
class ConsoleOutput:
    r"""Represents the standard output and standard error from code execution or a command.

    Attributes:
        exit_code (int): The exit code of the executed code or command. 0 typically indicates success.
        stderr (str): The content written to the standard error stream.
        stdout (str): The content written to the standard output stream.

    """

    exit_code: int = 0
    stderr: str = ""
    stdout: str = ""

    def text(self) -> str:
        r"""Get the text representation of the console output (stdout).

        .. deprecated:: 0.1.0
            The `text` property is deprecated and will be removed in a future version.
            Use the `stdout` attribute directly instead.

        Returns:
            str: The content of the standard output stream.

        """
        warnings.warn(
            "The 'text' property is deprecated and will be removed in a future version. "
            "Use 'stdout' attribute directly instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.stdout

    def success(self) -> bool:
        r"""Check if the execution was successful (exit code is 0).

        Returns:
            bool: True if `exit_code` is 0, False otherwise.

        """
        return not self.exit_code

    def to_json(self, include_plots: bool = False) -> str:
        r"""Get the JSON representation of the execution result.

        Args:
            include_plots (bool): Whether to include the plots in the JSON representation.

        Returns:
            str: The JSON representation of the execution result.

        """
        result = self.__dict__.copy()
        if not include_plots and "plots" in result:
            result.pop("plots", None)

        return json.dumps(result, indent=2)


@dataclass(frozen=True)
class ExecutionResult(ConsoleOutput):
    r"""Represents the comprehensive result of code execution within a sandbox session.

    This class extends `ConsoleOutput` to include any plots or other file artifacts
    that were generated and captured during the execution.

    Attributes:
        plots (list[PlotOutput]): A list of `PlotOutput` objects, each representing a
                                    captured plot or visual artifact. Defaults to an empty list.

    """

    plots: list[PlotOutput] = field(default_factory=list)
