"""
Annie SDK Response

Structured response types for query results. Follows the same pattern
as the pandas-ai-sdk AgentResponse.
"""

from __future__ import annotations

import json
from enum import Enum
from typing import Any, Iterator


class ResponseType(str, Enum):
    """Types of response content."""

    TEXT = "text"
    CHART = "chart"
    DATA = "data"
    ERROR = "error"


class ChartFormat(str, Enum):
    """Chart output formats."""

    CHARTJS = "chartjs"


class TextContent:
    """Text content in a response."""

    def __init__(self, content: str):
        self.type = ResponseType.TEXT
        self.content = content

    def to_dict(self) -> dict:
        return {"type": self.type.value, "content": self.content}

    def __str__(self) -> str:
        return self.content

    def __repr__(self) -> str:
        return f"TextContent(content={self.content!r})"


class _JSONSafeDict(dict):
    """A dict subclass whose str()/repr() output valid JSON.

    Python's default dict.__repr__ uses True/False/None and single quotes,
    which are invalid JavaScript.  This subclass ensures that print(d) and
    str(d) always produce JSON (lowercase true/false/null, double quotes).
    """

    def __str__(self) -> str:
        return json.dumps(self, default=str)

    def __repr__(self) -> str:
        return json.dumps(self, default=str)


def _make_json_safe(obj: Any) -> Any:
    """Recursively wrap dicts in _JSONSafeDict and coalesce None in lists.

    This ensures that every nested dict prints as valid JSON and that
    None values in data arrays are replaced with sensible defaults.
    """
    if isinstance(obj, dict):
        return _JSONSafeDict({k: _make_json_safe(v) for k, v in obj.items()})
    if isinstance(obj, list):
        return [_make_json_safe(v) for v in obj]
    return obj


class ChartContent:
    """Chart content in a response (ChartJS configuration)."""

    def __init__(
        self,
        content: dict[str, Any],
        format: ChartFormat = ChartFormat.CHARTJS,
        title: str | None = None,
    ):
        self.type = ResponseType.CHART
        self.content = _make_json_safe(content)
        self.format = format
        self.title = title

    def to_dict(self) -> dict:
        result: dict = {
            "type": self.type.value,
            "content": self.content,
            "format": self.format.value,
        }
        if self.title is not None:
            result["title"] = self.title
        return result

    def to_json(self, indent: int | None = None) -> str:
        """Convert to a JSON string safe for JavaScript consumption."""
        return json.dumps(self.to_dict(), indent=indent, default=str)

    def __str__(self) -> str:
        """Return JSON string so print() outputs JS-safe values."""
        return json.dumps(self.content, default=str)

    def __repr__(self) -> str:
        return f"ChartContent(format={self.format.value!r}, title={self.title!r})"


class DataContent:
    """Data content in a response (rows and columns)."""

    def __init__(
        self,
        columns: list[str],
        rows: list[list[Any]],
    ):
        self.type = ResponseType.DATA
        self.columns = columns
        self.rows = rows

    @property
    def row_count(self) -> int:
        return len(self.rows)

    def to_dict(self) -> dict:
        return {
            "type": self.type.value,
            "columns": self.columns,
            "rows": self.rows,
            "row_count": self.row_count,
        }

    def to_json(self, indent: int | None = None) -> str:
        """Convert to a JSON string safe for JavaScript consumption."""
        return json.dumps(self.to_dict(), indent=indent, default=str)

    def to_records(self) -> list[dict[str, Any]]:
        """Convert to list of dictionaries (like pandas to_dict('records'))."""
        return [dict(zip(self.columns, row)) for row in self.rows]

    def __str__(self) -> str:
        """Return JSON string so print() outputs JS-safe values."""
        return json.dumps(self.to_dict(), indent=2, default=str)

    def __repr__(self) -> str:
        return f"DataContent(columns={self.columns!r}, row_count={self.row_count})"


class ErrorContent:
    """Error content in a response."""

    def __init__(self, content: str, code: str | None = None):
        self.type = ResponseType.ERROR
        self.content = content
        self.code = code

    def to_dict(self) -> dict:
        result: dict[str, Any] = {"type": self.type.value, "content": self.content}
        if self.code:
            result["code"] = self.code
        return result

    def __str__(self) -> str:
        return self.content

    def __repr__(self) -> str:
        return f"ErrorContent(content={self.content!r}, code={self.code!r})"


ResponseContent = TextContent | ChartContent | DataContent | ErrorContent


class AgentResponse:
    """Structured response from an Agent query.

    Contains multiple content items (text, chart, data, error) along
    with metadata like SQL, DSL, and recommended visualization type.
    """

    def __init__(
        self,
        items: list[ResponseContent] | None = None,
        sql: str | None = None,
        dsl: dict[str, Any] | None = None,
        visualization: str = "table",
        success: bool = True,
    ):
        self.items: list[ResponseContent] = items or []
        self.sql = sql
        self.dsl = dsl
        self.visualization = visualization
        self.success = success

    # =========================================================================
    # Builder methods
    # =========================================================================

    def add_text(self, content: str) -> AgentResponse:
        """Add a text item to the response."""
        self.items.append(TextContent(content))
        return self

    def add_chart(
        self,
        content: dict[str, Any],
        format: ChartFormat = ChartFormat.CHARTJS,
        title: str | None = None,
    ) -> AgentResponse:
        """Add a chart item to the response."""
        self.items.append(ChartContent(content, format, title))
        return self

    def add_data(
        self,
        columns: list[str],
        rows: list[list[Any]],
    ) -> AgentResponse:
        """Add a data item to the response."""
        self.items.append(DataContent(columns, rows))
        return self

    def add_error(self, content: str, code: str | None = None) -> AgentResponse:
        """Add an error item to the response."""
        self.items.append(ErrorContent(content, code))
        self.success = False
        return self

    # =========================================================================
    # Convenience accessors
    # =========================================================================

    @property
    def text(self) -> str | None:
        """Get the first text content, or None."""
        for item in self.items:
            if isinstance(item, TextContent):
                return item.content
        return None

    @property
    def texts(self) -> list[str]:
        """Get all text contents."""
        return [item.content for item in self.items if isinstance(item, TextContent)]

    @property
    def chart(self) -> ChartContent | None:
        """Get the first chart content, or None."""
        for item in self.items:
            if isinstance(item, ChartContent):
                return item
        return None

    @property
    def charts(self) -> list[ChartContent]:
        """Get all chart contents."""
        return [item for item in self.items if isinstance(item, ChartContent)]

    @property
    def data(self) -> DataContent | None:
        """Get the first data content, or None."""
        for item in self.items:
            if isinstance(item, DataContent):
                return item
        return None

    @property
    def all_data(self) -> list[DataContent]:
        """Get all data contents."""
        return [item for item in self.items if isinstance(item, DataContent)]

    @property
    def error(self) -> ErrorContent | None:
        """Get the first error content, or None."""
        for item in self.items:
            if isinstance(item, ErrorContent):
                return item
        return None

    # =========================================================================
    # Serialization
    # =========================================================================

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "items": [item.to_dict() for item in self.items],
            "sql": self.sql,
            "dsl": self.dsl,
            "visualization": self.visualization,
            "success": self.success,
        }

    def to_json(self, indent: int | None = None) -> str:
        """Convert to JSON string."""
        return json.dumps(self.to_dict(), indent=indent, default=str)

    # =========================================================================
    # Iteration
    # =========================================================================

    def __iter__(self) -> Iterator[ResponseContent]:
        return iter(self.items)

    def __len__(self) -> int:
        return len(self.items)

    def __str__(self) -> str:
        """Return JSON string so print() outputs JS-safe values."""
        return self.to_json(indent=2)

    def __repr__(self) -> str:
        item_types = [item.type.value for item in self.items]
        return f"AgentResponse(success={self.success}, items={item_types})"

    # =========================================================================
    # Factory methods
    # =========================================================================

    @classmethod
    def from_error(
        cls,
        message: str,
        code: str | None = None,
    ) -> AgentResponse:
        """Create an error response."""
        response = cls(success=False)
        response.add_error(message, code)
        return response
