"""Chart tool: generate charts and visualizations as images."""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

from loguru import logger

from madpod_agent.agent.tools.base import Tool


class ChartTool(Tool):
    """Generate charts and data visualizations saved as image files."""

    def __init__(self, workspace: Path):
        self._workspace = workspace

    @property
    def name(self) -> str:
        return "chart"

    @property
    def description(self) -> str:
        return (
            "Generate a chart or data visualization and save it as a PNG image. "
            "Supports: bar, line, pie, scatter, histogram, heatmap, and area charts. "
            "Provide data as arrays. Returns the file path of the generated image."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "chart_type": {
                    "type": "string",
                    "enum": ["bar", "line", "pie", "scatter", "histogram", "heatmap", "area"],
                    "description": "Type of chart to generate",
                },
                "title": {
                    "type": "string",
                    "description": "Chart title",
                },
                "x_label": {
                    "type": "string",
                    "description": "X-axis label",
                },
                "y_label": {
                    "type": "string",
                    "description": "Y-axis label",
                },
                "labels": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Category labels (x-axis for bar/line, slice names for pie)",
                },
                "datasets": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "name": {"type": "string"},
                            "values": {"type": "array", "items": {"type": "number"}},
                        },
                    },
                    "description": "Data series. Each has a name and values array.",
                },
                "x_values": {
                    "type": "array",
                    "items": {"type": "number"},
                    "description": "X values for scatter plots",
                },
                "y_values": {
                    "type": "array",
                    "items": {"type": "number"},
                    "description": "Y values for scatter plots",
                },
                "filename": {
                    "type": "string",
                    "description": "Output filename (default: chart.png)",
                },
                "width": {
                    "type": "integer",
                    "description": "Chart width in inches (default: 10)",
                },
                "height": {
                    "type": "integer",
                    "description": "Chart height in inches (default: 6)",
                },
            },
            "required": ["chart_type"],
        }

    async def execute(
        self,
        chart_type: str,
        title: str = "",
        x_label: str = "",
        y_label: str = "",
        labels: list[str] | None = None,
        datasets: list[dict[str, Any]] | None = None,
        x_values: list[float] | None = None,
        y_values: list[float] | None = None,
        filename: str = "chart.png",
        width: int = 10,
        height: int = 6,
        **kwargs: Any,
    ) -> str:
        try:
            import matplotlib
            matplotlib.use("Agg")  # Non-interactive backend
            import matplotlib.pyplot as plt
        except ImportError:
            return json.dumps({
                "error": "matplotlib not installed. Install: pip install matplotlib"
            })

        output_dir = self._workspace / "charts"
        output_dir.mkdir(parents=True, exist_ok=True)
        if not filename.endswith(".png"):
            filename += ".png"
        output_path = output_dir / filename

        try:
            fig, ax = plt.subplots(figsize=(width, height))

            if chart_type == "bar":
                self._render_bar(ax, labels, datasets)
            elif chart_type == "line":
                self._render_line(ax, labels, datasets)
            elif chart_type == "pie":
                self._render_pie(ax, labels, datasets)
            elif chart_type == "scatter":
                self._render_scatter(ax, x_values, y_values, datasets)
            elif chart_type == "histogram":
                self._render_histogram(ax, datasets)
            elif chart_type == "area":
                self._render_area(ax, labels, datasets)
            elif chart_type == "heatmap":
                self._render_heatmap(ax, labels, datasets)
            else:
                plt.close(fig)
                return json.dumps({"error": f"Unsupported chart type: {chart_type}"})

            if title:
                ax.set_title(title, fontsize=14, fontweight="bold")
            if x_label:
                ax.set_xlabel(x_label)
            if y_label:
                ax.set_ylabel(y_label)

            if chart_type != "pie":
                ax.legend()

            fig.tight_layout()
            fig.savefig(output_path, dpi=150, bbox_inches="tight")
            plt.close(fig)

            logger.info("Generated {} chart: {}", chart_type, output_path)
            return json.dumps({
                "ok": True,
                "path": str(output_path),
                "chart_type": chart_type,
                "width": width,
                "height": height,
            })

        except Exception as exc:
            plt.close("all")
            logger.error("Chart generation failed: {}", exc)
            return json.dumps({"error": f"Chart generation failed: {exc}"})

    def _render_bar(
        self, ax: Any, labels: list[str] | None, datasets: list[dict[str, Any]] | None,
    ) -> None:
        import numpy as np
        labels = labels or []
        datasets = datasets or []
        x = np.arange(len(labels))
        n = len(datasets)
        bar_width = 0.8 / max(n, 1)
        for i, ds in enumerate(datasets):
            offset = (i - n / 2 + 0.5) * bar_width
            ax.bar(x + offset, ds.get("values", []), bar_width, label=ds.get("name", f"Series {i+1}"))
        ax.set_xticks(x)
        ax.set_xticklabels(labels, rotation=45, ha="right")

    def _render_line(
        self, ax: Any, labels: list[str] | None, datasets: list[dict[str, Any]] | None,
    ) -> None:
        datasets = datasets or []
        for ds in datasets:
            values = ds.get("values", [])
            x = list(range(len(values))) if not labels else range(len(values))
            ax.plot(x, values, marker="o", markersize=4, label=ds.get("name", ""))
        if labels:
            ax.set_xticks(range(len(labels)))
            ax.set_xticklabels(labels, rotation=45, ha="right")

    def _render_pie(
        self, ax: Any, labels: list[str] | None, datasets: list[dict[str, Any]] | None,
    ) -> None:
        datasets = datasets or []
        values = datasets[0].get("values", []) if datasets else []
        labels = labels or [f"Slice {i+1}" for i in range(len(values))]
        ax.pie(values, labels=labels, autopct="%1.1f%%", startangle=90)
        ax.axis("equal")

    def _render_scatter(
        self, ax: Any,
        x_values: list[float] | None,
        y_values: list[float] | None,
        datasets: list[dict[str, Any]] | None,
    ) -> None:
        if x_values and y_values:
            ax.scatter(x_values, y_values, alpha=0.7, label="data")
        elif datasets:
            for ds in datasets:
                vals = ds.get("values", [])
                ax.scatter(range(len(vals)), vals, alpha=0.7, label=ds.get("name", ""))

    def _render_histogram(
        self, ax: Any, datasets: list[dict[str, Any]] | None,
    ) -> None:
        datasets = datasets or []
        for ds in datasets:
            ax.hist(ds.get("values", []), bins="auto", alpha=0.7, label=ds.get("name", ""))

    def _render_area(
        self, ax: Any, labels: list[str] | None, datasets: list[dict[str, Any]] | None,
    ) -> None:
        datasets = datasets or []
        for ds in datasets:
            values = ds.get("values", [])
            x = range(len(values))
            ax.fill_between(x, values, alpha=0.4, label=ds.get("name", ""))
            ax.plot(x, values, linewidth=1)
        if labels:
            ax.set_xticks(range(len(labels)))
            ax.set_xticklabels(labels, rotation=45, ha="right")

    def _render_heatmap(
        self, ax: Any, labels: list[str] | None, datasets: list[dict[str, Any]] | None,
    ) -> None:
        import numpy as np
        datasets = datasets or []
        data = [ds.get("values", []) for ds in datasets]
        if not data:
            return
        arr = np.array(data)
        im = ax.imshow(arr, aspect="auto", cmap="YlOrRd")
        ax.figure.colorbar(im, ax=ax)
        if labels:
            ax.set_xticks(range(len(labels)))
            ax.set_xticklabels(labels, rotation=45, ha="right")
        row_labels = [ds.get("name", f"Row {i}") for i, ds in enumerate(datasets)]
        ax.set_yticks(range(len(row_labels)))
        ax.set_yticklabels(row_labels)
