from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from backend.services.llm_service import LLMService


@dataclass
class VisualizationOutput:
    chart_type: str
    x_axis: str
    y_axis: str
    title: str
    reason: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "chart_type": self.chart_type,
            "x_axis": self.x_axis,
            "y_axis": self.y_axis,
            "title": self.title,
            "reason": self.reason,
        }


class VisualizationAgent:
    def __init__(self, llm: LLMService) -> None:
        self.llm = llm

    def run(self, question: str, sql_result: dict[str, Any]) -> VisualizationOutput:
        columns = sql_result.get("columns", [])
        rows = sql_result.get("rows", [])

        fallback = self._fallback_visualization(question=question, columns=columns, rows=rows)
        if not columns:
            return fallback

        sample_rows = rows[:8]
        payload = self.llm.invoke_json(
            system_prompt=(
                "You are a BI visualization agent. Return strict JSON with keys "
                "chart_type, x_axis, y_axis, title, reason. "
                "Choose only from chart_type values: line, bar, scatter, table. "
                "Use only axes present in the provided columns. "
                "Prefer table when the result is not clearly chartable."
            ),
            user_prompt=(
                f"Question: {question}\n"
                f"Columns: {columns}\n"
                f"Sample rows: {sample_rows}\n"
                "Pick the best visualization for a Streamlit frontend."
            ),
            fallback=fallback.to_dict(),
        )

        chart_type = payload.get("chart_type", fallback.chart_type)
        x_axis = payload.get("x_axis", fallback.x_axis)
        y_axis = payload.get("y_axis", fallback.y_axis)
        title = payload.get("title", question or fallback.title)
        reason = payload.get("reason", fallback.reason)

        if chart_type not in {"line", "bar", "scatter", "table"}:
            return fallback
        if chart_type != "table" and (x_axis not in columns or y_axis not in columns):
            return fallback

        return VisualizationOutput(
            chart_type=chart_type,
            x_axis=x_axis,
            y_axis=y_axis,
            title=title,
            reason=reason,
        )

    def _fallback_visualization(
        self,
        question: str,
        columns: list[str],
        rows: list[dict[str, Any]],
    ) -> VisualizationOutput:
        if not columns:
            return VisualizationOutput(
                chart_type="table",
                x_axis="",
                y_axis="",
                title="Query Result",
                reason="No chartable columns were returned, so the table is the primary output.",
            )

        sample = rows[0] if rows else {}
        numeric_columns = [column for column in columns if isinstance(sample.get(column), (int, float))]
        time_columns = [
            column
            for column in columns
            if any(keyword in column.lower() for keyword in {"date", "time", "day", "month", "year", "period"})
        ]
        category_columns = [column for column in columns if column not in numeric_columns]

        if time_columns and numeric_columns:
            return VisualizationOutput(
                chart_type="line",
                x_axis=time_columns[0],
                y_axis=numeric_columns[0],
                title=question,
                reason="A time column and numeric measure were returned, so a line chart best shows the trend.",
            )
        if category_columns and numeric_columns:
            return VisualizationOutput(
                chart_type="bar",
                x_axis=category_columns[0],
                y_axis=numeric_columns[0],
                title=question,
                reason="A categorical dimension and numeric measure were returned, so a bar chart compares categories well.",
            )
        if len(numeric_columns) >= 2:
            return VisualizationOutput(
                chart_type="scatter",
                x_axis=numeric_columns[0],
                y_axis=numeric_columns[1],
                title=question,
                reason="Two numeric columns were returned, so a scatter plot can show their relationship.",
            )
        return VisualizationOutput(
            chart_type="table",
            x_axis="",
            y_axis="",
            title="Query Result",
            reason="The returned data shape is best explored directly as a table.",
        )