"""Data Agent — processes data, generates chart specs and table data."""

from __future__ import annotations

import json
from typing import Any

from .base import BaseAgent
from app.config import settings


class DataAgent(BaseAgent):
    name = "data"
    description = "处理数据、生成图表规格和表格数据"
    system_prompt = """\
你是一位数据分析专家。你的任务是根据报告草稿中标注的图表和表格需求，
生成具体的数据和图表规格。

输出要求（JSON 格式）：
{
  "charts": [
    {
      "id": "chart_1",
      "title": "图表标题",
      "type": "bar|line|pie|area|scatter",
      "description": "图表说明",
      "data": {
        "labels": ["标签1", "标签2"],
        "datasets": [
          {"label": "数据集名", "data": [100, 200]}
        ]
      }
    }
  ],
  "tables": [
    {
      "id": "table_1",
      "title": "表格标题",
      "headers": ["列1", "列2", "列3"],
      "rows": [["数据1", "数据2", "数据3"]]
    }
  ]
}"""

    def __init__(self):
        super().__init__(model=settings.model_for_domain("fast"))

    async def run(self, context: dict[str, Any]) -> dict[str, Any]:
        draft = context["draft"]
        extra_data = context.get("extra_data", "")

        # Collect chart/table needs from draft
        chart_needs = []
        table_needs = []
        for ch in draft.get("chapters", []):
            chart_needs.extend(ch.get("charts", []))
            table_needs.extend(ch.get("tables", []))

        if not chart_needs and not table_needs:
            return {"data_assets": {"charts": [], "tables": []}}

        prompt = f"""\
## 报告标题
{draft.get("title", "")}

## 需要生成的图表
{json.dumps(chart_needs, ensure_ascii=False)}

## 需要生成的表格
{json.dumps(table_needs, ensure_ascii=False)}

## 补充数据源
{extra_data if extra_data else "（无额外数据，请根据行业常识生成合理的示例数据）"}

请为以上需求生成具体的图表规格和表格数据。输出 JSON。"""

        result = await self.call_llm_json(prompt)
        return {"data_assets": result}