"""Suggest chart type from result shape using local LLM."""

from __future__ import annotations

import json

from local_llm.ollama_client import OllamaClient


class ChartAdvisor:
    """Suggests visualization type based on query results — runs locally."""

    def __init__(self, ollama: OllamaClient) -> None:
        self._ollama = ollama

    async def suggest(
        self,
        question: str,
        columns: list[str],
        row_count: int,
        sample_rows: list[dict],
    ) -> dict | None:
        """Suggest a chart configuration for the query results.

        Returns a Plotly-compatible chart config or None if no chart is appropriate.
        """
        if row_count == 0:
            return None

        sample_text = json.dumps(sample_rows[:5], indent=2, default=str)

        prompt = f"""Given this query result, suggest the best chart visualization.

Question: {question}
Columns: {', '.join(columns)}
Row count: {row_count}
Sample data:
{sample_text}

Choose one:
- "bar": For comparing categories
- "line": For time series or trends
- "pie": For proportions (< 8 categories)
- "scatter": For correlations between two numeric columns
- "table_only": If the data is better shown as a table

Respond in JSON:
{{"chart_type": "bar|line|pie|scatter|table_only", "x_column": "column_name", "y_column": "column_name", "title": "Chart Title"}}

If chart_type is "table_only", only include chart_type."""

        response = await self._ollama.generate(
            prompt=prompt,
            system="You are a data visualization expert. Respond only with valid JSON.",
            temperature=0.1,
        )

        try:
            result = json.loads(response.strip())
            if result.get("chart_type") == "table_only":
                return None
            return result
        except (json.JSONDecodeError, ValueError):
            return None
