#!/usr/bin/env python3
"""
chart.py — Universal Chart & Graph Generator for OpenClaw / Telegram

Generates clean, publication-grade chart images from JSON/CSV data using:
1. Primary: dbt-charts (`dct render`) — local, offline, Vega-Lite/resvg based.
2. Fallback: QuickChart.io — fast HTTP API (Chart.js), zero local dependencies.

Outputs `MEDIA:<path>` on stdout for native Telegram image delivery.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import subprocess
import sys
import tempfile
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union


KNOWN_DCT_PATHS = [
    Path("/home/node/.openclaw/.local/bin/dct"),
    Path("/home/openclaw/.local/bin/dct"),
    Path("/home/node/.openclaw/.local/share/uv/tools/dbt-charts/bin/dct"),
    Path("/home/openclaw/.local/share/uv/tools/dbt-charts/bin/dct"),
    Path.home() / ".local/bin/dct",
]


def find_dct_executable() -> Optional[str]:
    """Locate the dct executable in PATH or well-known install locations."""
    # 1. System PATH
    found = shutil.which("dct")
    if found:
        return found

    # 2. Well-known absolute container/local paths
    for candidate in KNOWN_DCT_PATHS:
        if candidate.is_file() and os.access(candidate, os.X_OK):
            return str(candidate)

    # 3. Check if uv tool run dct works
    uv_path = shutil.which("uv") or "/home/node/.openclaw/bin/uv"
    if os.path.exists(uv_path) and os.access(uv_path, os.X_OK):
        try:
            res = subprocess.run(
                [uv_path, "tool", "run", "dct", "--version"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            if res.returncode == 0:
                return "uv_tool_run"
        except Exception:
            pass

    return None


def parse_input_data(raw_data: str) -> Tuple[List[str], List[List[Any]], Optional[str], Optional[str]]:
    """
    Parse input string (JSON or CSV) into (columns, rows, suggested_x, suggested_y).
    Supports:
    - List of dicts: [{"label": "A", "val": 10}, {"label": "B", "val": 20}]
    - Dict of values: {"A": 10, "B": 20}
    - Explicit table: {"columns": ["k", "v"], "values": [["A", 10], ["B", 20]]}
    - Single KPI: {"value": 1284, "label": "Revenue", "delta": "+12%"}
    """
    raw_data = raw_data.strip()
    if not raw_data:
        raise ValueError("Input data is empty.")

    # Try JSON
    parsed = json.loads(raw_data)

    # Case 1: Dict with explicit columns & values
    if isinstance(parsed, dict) and "columns" in parsed and ("values" in parsed or "rows" in parsed):
        cols = [str(c) for c in parsed["columns"]]
        rows = parsed.get("values") or parsed.get("rows") or []
        x_col = cols[0] if len(cols) > 0 else None
        y_col = cols[1] if len(cols) > 1 else None
        return cols, rows, x_col, y_col

    # Case 2: Dict of key-value pairs (e.g. {"Mo": 10, "Di": 15})
    if isinstance(parsed, dict) and not any(k in parsed for k in ["columns", "values", "rows", "type", "chart"]):
        cols = ["category", "value"]
        rows = [[k, v] for k, v in parsed.items()]
        return cols, rows, "category", "value"

    # Case 3: List of dicts (most common agent format)
    if isinstance(parsed, list):
        if not parsed:
            raise ValueError("Data list is empty.")
        if isinstance(parsed[0], dict):
            # Collect unique keys maintaining insertion order
            cols: List[str] = []
            for item in parsed:
                for k in item.keys():
                    if k not in cols:
                        cols.append(k)

            rows: List[List[Any]] = []
            for item in parsed:
                rows.append([item.get(c) for c in cols])

            # Infer x (category/date) and y (measure)
            x_col = None
            y_col = None
            for c in cols:
                low = c.lower()
                if low in ["x", "label", "name", "date", "datum", "month", "monat", "day", "tag", "category", "kategorie", "typ", "woche", "week"]:
                    x_col = c
                    break
            if not x_col and len(cols) > 0:
                x_col = cols[0]

            for c in cols:
                if c == x_col:
                    continue
                low = c.lower()
                if low in ["y", "value", "wert", "count", "anzahl", "amount", "menge", "total", "summe", "preis", "price", "revenue", "umsatz", "km", "kmh", "watt", "hr", "puls"]:
                    y_col = c
                    break
            if not y_col and len(cols) > 1:
                y_col = cols[1]

            return cols, rows, x_col, y_col

        elif isinstance(parsed[0], list):
            # List of lists without explicit headers
            cols = [f"col_{i+1}" for i in range(len(parsed[0]))]
            return cols, parsed, cols[0], cols[1] if len(cols) > 1 else None

    raise ValueError(f"Unrecognized data structure: {type(parsed)}")


def render_dbt_chart(
    dct_cmd: str,
    title: str,
    chart_type: str,
    columns: List[str],
    rows: List[List[Any]],
    x_col: str,
    y_col: str,
    output_path: Path,
    notes: Optional[str] = None,
    theme: Optional[str] = None,
) -> bool:
    """Render a chart using dbt-charts (dct render)."""
    scratch_dir = Path(tempfile.mkdtemp(prefix="dct_render_"))
    try:
        # 1. dbt_charts.yml project file (mandatory for dct)
        (scratch_dir / "dbt_charts.yml").write_text("# dbt charts project\nversion: 1\n", encoding="utf-8")

        # Normalize chart type
        ct = chart_type.lower()
        if ct in ["hbar", "horizontal_bar"]:
            ct = "bar"
            is_horizontal = True
        else:
            is_horizontal = False

        if ct in ["donut", "pie"]:
            dbt_type = "donut"
        elif ct in ["line", "trend"]:
            dbt_type = "line"
        elif ct in ["area"]:
            dbt_type = "area"
        elif ct in ["kpi"]:
            dbt_type = "kpi"
        elif ct in ["table"]:
            dbt_type = "table"
        else:
            dbt_type = "bar"

        # Build board YAML
        board: Dict[str, Any] = {
            "title": title or "Dashboard",
            "queries": {
                "dataset": {
                    "columns": columns,
                    "values": rows,
                }
            },
            "charts": {},
            "rows": ["main_chart"],
        }

        if notes:
            board["notes"] = notes

        chart_spec: Dict[str, Any] = {
            "type": dbt_type,
            "query": "dataset",
        }
        if title:
            chart_spec["title"] = title

        if dbt_type in ["donut", "pie"]:
            chart_spec["theta"] = y_col
            chart_spec["color"] = x_col
        elif dbt_type == "kpi":
            chart_spec["value"] = y_col or columns[0]
            chart_spec["label"] = x_col or "Value"
        elif dbt_type == "table":
            # Table automatically uses query columns
            pass
        else:
            # Bar / Line / Area
            if is_horizontal:
                chart_spec["x"] = y_col
                chart_spec["y"] = x_col
            else:
                chart_spec["x"] = x_col
                chart_spec["y"] = y_col

        board["charts"]["main_chart"] = chart_spec

        import yaml  # PyYAML is standard in python / uv environments
        board_file = scratch_dir / "board.yml"
        board_file.write_text(yaml.safe_dump(board, sort_keys=False), encoding="utf-8")

        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Build command
        if dct_cmd == "uv_tool_run":
            cmd = ["uv", "tool", "run", "dct", "render", str(board_file), "--format", "png", "--output", str(output_path)]
        else:
            cmd = [dct_cmd, "render", str(board_file), "--format", "png", "--output", str(output_path)]

        proc = subprocess.run(
            cmd,
            cwd=str(scratch_dir),
            capture_output=True,
            text=True,
            timeout=30,
        )

        if proc.returncode == 0 and output_path.exists() and output_path.stat().st_size > 0:
            return True
        else:
            sys.stderr.write(f"[dct render warning] exit={proc.returncode} stderr={proc.stderr.strip()}\n")
            return False

    except Exception as e:
        sys.stderr.write(f"[dct render error] {e}\n")
        return False
    finally:
        shutil.rmtree(scratch_dir, ignore_errors=True)


def render_quickchart(
    title: str,
    chart_type: str,
    columns: List[str],
    rows: List[List[Any]],
    x_col: str,
    y_col: str,
    output_path: Path,
    width: int = 800,
    height: int = 450,
) -> bool:
    """Render chart via QuickChart.io API as fallback."""
    try:
        x_idx = columns.index(x_col)
        y_idx = columns.index(y_col)

        labels = [str(r[x_idx]) for r in rows]
        values = []
        for r in rows:
            v = r[y_idx]
            try:
                values.append(float(v) if v is not None else 0)
            except (ValueError, TypeError):
                values.append(0)

        ct = chart_type.lower()
        if ct in ["donut", "pie"]:
            qc_type = "doughnut" if ct == "donut" else "pie"
            bg_colors = [
                "#3366cc", "#dc3912", "#ff9900", "#109618", "#990099",
                "#3b3eac", "#0099c6", "#dd4477", "#66aa00", "#b82e2e",
            ]
            dataset = {
                "data": values,
                "backgroundColor": bg_colors[:len(values)],
            }
        elif ct in ["line", "trend"]:
            qc_type = "line"
            dataset = {
                "label": y_col,
                "data": values,
                "borderColor": "#3e64ff",
                "backgroundColor": "rgba(62, 100, 255, 0.1)",
                "fill": False,
                "tension": 0.2,
            }
        elif ct in ["area"]:
            qc_type = "line"
            dataset = {
                "label": y_col,
                "data": values,
                "borderColor": "#3e64ff",
                "backgroundColor": "rgba(62, 100, 255, 0.25)",
                "fill": True,
                "tension": 0.2,
            }
        else:
            qc_type = "bar"
            dataset = {
                "label": y_col,
                "data": values,
                "backgroundColor": "#3e64ff",
            }

        chart_config = {
            "type": qc_type,
            "data": {
                "labels": labels,
                "datasets": [dataset],
            },
            "options": {
                "title": {
                    "display": bool(title),
                    "text": title,
                    "fontSize": 18,
                },
                "legend": {
                    "display": qc_type in ["pie", "doughnut"],
                    "position": "top",
                },
                "plugins": {
                    "datalabels": {
                        "display": False,
                    }
                }
            }
        }

        query_params = {
            "w": str(width),
            "h": str(height),
            "devicePixelRatio": "2.0",
            "c": json.dumps(chart_config),
        }
        url = "https://quickchart.io/chart?" + urllib.parse.urlencode(query_params)

        req = urllib.request.Request(url, headers={"User-Agent": "OpenClaw-Chart/1.0"})
        with urllib.request.urlopen(req, timeout=15) as resp:
            if resp.status == 200:
                output_path.parent.mkdir(parents=True, exist_ok=True)
                output_path.write_bytes(resp.read())
                return True

    except Exception as e:
        sys.stderr.write(f"[quickchart error] {e}\n")

    return False


def main() -> None:
    parser = argparse.ArgumentParser(description="OpenClaw Universal Chart Generator")
    parser.add_argument("--data", help="JSON data string, or '-' for stdin")
    parser.add_argument("--file", help="Path to JSON or CSV file")
    parser.add_argument("--type", default="bar", choices=["bar", "hbar", "line", "area", "donut", "pie", "kpi", "table"], help="Chart type")
    parser.add_argument("--title", default="Übersicht", help="Chart or Dashboard title")
    parser.add_argument("--notes", help="Subtitle or explanatory notes")
    parser.add_argument("--x", help="X-axis / category column name")
    parser.add_argument("--y", help="Y-axis / measure column name")
    parser.add_argument("--output", help="Target PNG path (default: /tmp/charts/chart_<hash>.png)")
    parser.add_argument("--engine", default="auto", choices=["auto", "dct", "quickchart"], help="Rendering engine")
    parser.add_argument("--width", type=int, default=800, help="Width in pixels (QuickChart)")
    parser.add_argument("--height", type=int, default=450, help="Height in pixels (QuickChart)")

    args = parser.parse_args()

    # Read data
    raw_data = None
    if args.data and args.data != "-":
        raw_data = args.data
    elif args.file:
        raw_data = Path(args.file).read_text(encoding="utf-8")
    elif args.data == "-" or not sys.stdin.isatty():
        raw_data = sys.stdin.read()

    if not raw_data:
        parser.error("No input data provided. Use --data, --file, or pipe data via STDIN.")

    try:
        columns, rows, inferred_x, inferred_y = parse_input_data(raw_data)
    except Exception as e:
        sys.stderr.write(f"Error parsing input data: {e}\n")
        sys.exit(1)

    x_col = args.x or inferred_x or (columns[0] if columns else "x")
    y_col = args.y or inferred_y or (columns[1] if len(columns) > 1 else columns[0])

    # Destination path
    if args.output:
        output_path = Path(args.output)
    else:
        content_hash = hashlib.md5(f"{args.title}_{args.type}_{raw_data}".encode("utf-8")).hexdigest()[:8]
        output_path = Path(f"/tmp/charts/chart_{content_hash}.png")

    rendered = False

    # Attempt 1: dbt Charts (unless engine is explicitly quickchart)
    if args.engine in ["auto", "dct"]:
        dct_bin = find_dct_executable()
        if dct_bin:
            rendered = render_dbt_chart(
                dct_cmd=dct_bin,
                title=args.title,
                chart_type=args.type,
                columns=columns,
                rows=rows,
                x_col=x_col,
                y_col=y_col,
                output_path=output_path,
                notes=args.notes,
            )
        elif args.engine == "dct":
            sys.stderr.write("Error: dct executable not found and --engine dct was requested.\n")
            sys.exit(1)

    # Attempt 2: QuickChart Fallback
    if not rendered and args.engine in ["auto", "quickchart"]:
        rendered = render_quickchart(
            title=args.title,
            chart_type=args.type,
            columns=columns,
            rows=rows,
            x_col=x_col,
            y_col=y_col,
            output_path=output_path,
            width=args.width,
            height=args.height,
        )

    if not rendered:
        sys.stderr.write("Error: Chart rendering failed on all available engines.\n")
        sys.exit(1)

    # Deliver MEDIA line for Telegram/OpenClaw
    print(f"MEDIA:{output_path.resolve()}")


if __name__ == "__main__":
    main()