#!/usr/bin/env python3
"""
Step 3 v3: Map query results to template and generate HTML report.

Architecture:
  - Fixed templates (dashboard_overview / dashboard_enterprise):
      step2 already produced placeholder_values; just inject into HTML.

  - Normal templates (lineChart / barChart / lineChartCompare / pieChart /
                      dataTable / infoCard / dashboardKpiBar / dashboardKpiBarCompare):
      v3 NEW APPROACH — "LLM decides, Python extracts":
        1. LLM outputs a lightweight "mapping config" (< 200 tokens):
             - which row fields map to which JS array variables
             - text/scalar values (titles, labels, summaries)
        2. Python builds all JS arrays from raw rows using the mapping config.
             No LLM token generation for large arrays → 10x faster.
        3. Summary LLM call runs concurrently using FULL rows data for
             richer analysis (no extra latency since render is the bottleneck).
"""
import os
import sys
import json
import time
import argparse
import urllib.request
import re
from concurrent.futures import ThreadPoolExecutor
import concurrent.futures

SCRIPT_DIR    = os.path.dirname(os.path.abspath(__file__))
SKILL_DIR     = os.path.dirname(SCRIPT_DIR)
TEMPLATES_DIR = os.path.join(SKILL_DIR, "templates")
OUTPUT_DIR    = os.path.join(SKILL_DIR, "output")

LLM_API_URL   = os.environ.get("LLM_API_URL", "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation")
LLM_API_KEY   = os.environ.get("LLM_API_KEY", "sk-b36adcdd25d64b7c847e7c941b5c1e03")
LLM_MODEL     = os.environ.get("LLM_MODEL",      "qwen-plus")
LLM_MODEL_FAST = os.environ.get("LLM_MODEL_FAST", "qwen-turbo")

# Pie chart color palette
PIE_COLORS = ["#1677FF", "#52c41a", "#faad14", "#f5222d", "#722ed1",
              "#13c2c2", "#eb2f96", "#fa8c16", "#a0d911", "#1890ff"]


# ─────────────────────────────────────────────────────────────────────────────
# Time Format Utilities
# ─────────────────────────────────────────────────────────────────────────────

def smart_format_time(value):
    """
    Intelligently format a time/date string based on its precision.
    - '2024-11-09 00:00:00' or '2024-11-09 00:00:00.0' → '2024-11-09'
    - '2024-11-09 08:30:00' (non-zero time) → '2024-11-09 08:30'
    - '2024-11' → '2024-11'
    - '2024' → '2024'
    - Other strings → unchanged
    """
    if not isinstance(value, str):
        return value
    s = value.strip()
    # Match datetime with time part
    import re as _re
    m = _re.match(r'^(\d{4}-\d{2}-\d{2})[T ](\d{2}):(\d{2}):(\d{2})', s)
    if m:
        date_part = m.group(1)
        hh, mm, ss = m.group(2), m.group(3), m.group(4)
        if hh == '00' and mm == '00' and ss == '00':
            return date_part          # pure date
        elif ss == '00':
            return f"{date_part} {hh}:{mm}"   # hour:minute
        else:
            return f"{date_part} {hh}:{mm}:{ss}"
    return s


def detect_time_fields(rows):
    """
    Detect which fields in rows look like datetime strings
    (contain '-' and ':' or match date pattern).
    Returns a set of field names.
    """
    if not rows:
        return set()
    import re as _re
    time_fields = set()
    sample = rows[0]
    for k, v in sample.items():
        if isinstance(v, str):
            if _re.match(r'^\d{4}-\d{2}-\d{2}', v):
                time_fields.add(k)
    return time_fields


# ─────────────────────────────────────────────────────────────────────────────
# Utilities
# ─────────────────────────────────────────────────────────────────────────────

def call_llm(prompt, system="You are a helpful assistant.", model=None, max_tokens=800):
    use_model = model or LLM_MODEL
    headers = {"Authorization": f"Bearer {LLM_API_KEY}", "Content-Type": "application/json"}
    payload = {
        "model": use_model,
        "input": {"messages": [
            {"role": "system", "content": system},
            {"role": "user",   "content": prompt}
        ]},
        "parameters": {"result_format": "message", "max_tokens": max_tokens}
    }
    req = urllib.request.Request(
        LLM_API_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers=headers, method="POST"
    )
    try:
        resp = urllib.request.urlopen(req, timeout=60)
        return json.loads(resp.read().decode("utf-8"))["output"]["choices"][0]["message"]["content"]
    except Exception as e:
        sys.stderr.write(f"[LLM error] {e}\n")
        return None


def extract_json(text):
    if not text:
        return None
    for pattern in [r'```json\s*(.*?)\s*```', r'```\s*([\[{].*?)\s*```']:
        m = re.search(pattern, text, re.DOTALL)
        if m:
            try: return json.loads(m.group(1))
            except: pass
    try: return json.loads(text.strip())
    except: pass
    m = re.search(r'(\{.*\})', text, re.DOTALL)
    if m:
        try: return json.loads(m.group(1))
        except: pass
    return None


def read_file(path):
    if os.path.exists(path):
        with open(path, 'r', encoding='utf-8') as f:
            return f.read()
    return ""


def inject_placeholders(html, placeholder_values):
    for key, value in placeholder_values.items():
        if value is None:
            value = ""
        html = html.replace("{{" + key + "}}", str(value))
    return html


def inject_js_vars(html, js_vars):
    """Inject JS variables at the // {{INJECT_DATA}} placeholder."""
    if not js_vars:
        return html
    js_lines = [f"var {k} = {json.dumps(v, ensure_ascii=False)};" for k, v in js_vars.items()]
    injected = "\n        ".join(js_lines)
    if "// {{INJECT_DATA}}" in html:
        html = html.replace("// {{INJECT_DATA}}", injected)
    elif "{{INJECT_DATA}}" in html:
        html = html.replace("{{INJECT_DATA}}", injected)
    return html


# ─────────────────────────────────────────────────────────────────────────────
# Python Data Extraction Engine
# ─────────────────────────────────────────────────────────────────────────────

def extract_column(rows, field, format_time=False):
    """Extract a single column from rows as a list."""
    if format_time:
        return [smart_format_time(row.get(field)) for row in rows]
    return [row.get(field) for row in rows]


def build_table_rows_html(rows, columns, pie_mode=False, time_fields=None):
    """Build TABLE_ROWS HTML string from rows and column list."""
    if time_fields is None:
        time_fields = detect_time_fields(rows)
    html_rows = []
    for i, row in enumerate(rows):
        cells = []
        for j, col in enumerate(columns):
            val = row.get(col, "")
            if val is None:
                val = "-"
            # Apply smart time formatting
            if col in time_fields and isinstance(val, str):
                val = smart_format_time(val)
            if pie_mode and j == 0:
                color = PIE_COLORS[i % len(PIE_COLORS)]
                cells.append(f'<td><span class="color-dot" style="background-color: {color};"></span>{val}</td>')
            else:
                cells.append(f"<td>{val}</td>")
        html_rows.append(f"<tr>{''.join(cells)}</tr>")
    return "".join(html_rows)


def aggregate_rows_by_key(rows, key_field, value_fields):
    """
    If rows contain duplicate values for key_field, aggregate (SUM) the value_fields.
    Preserves original row order of first occurrence.
    Returns a new list of aggregated rows.
    """
    from collections import OrderedDict
    seen = OrderedDict()
    for row in rows:
        k = row.get(key_field, "")
        if k not in seen:
            seen[k] = {key_field: k}
            for vf in value_fields:
                seen[k][vf] = 0.0
        for vf in value_fields:
            try:
                seen[k][vf] += float(row.get(vf) or 0)
            except (TypeError, ValueError):
                seen[k][vf] = row.get(vf)  # non-numeric: keep last value
    return list(seen.values())


def needs_aggregation(rows, key_field):
    """Return True if key_field has duplicate values in rows."""
    if not rows or not key_field:
        return False
    vals = [row.get(key_field) for row in rows]
    return len(vals) != len(set(str(v) for v in vals))


def python_extract_js_vars(template_name, mapping, rows):
    """
    Core extraction engine: given the LLM mapping config and raw rows,
    build the complete js_variables dict using Python (no LLM token generation).

    mapping schema (from LLM):
    {
      "type": "line|bar|lineCompare|pie|table|kpiBar|kpiBarCompare|infoCard",

      # For line / bar:
      "x_field": "stat_date",
      "y_field": "daily_charge_energy",
      "y_axis_name": "kWh",
      "series_name": "日充电量",

      # For lineCompare:
      "x_field": "stat_date",
      "y1_field": "daily_charge_energy",
      "y2_field": "daily_discharge_energy",
      "y_axis_name": "kWh",
      "series_name_1": "充电量",
      "series_name_2": "放电量",

      # For pie:
      "name_field": "category",
      "value_field": "amount",
      "series_name": "能源占比",

      # For table:
      "columns": [
        {"label": "日期", "key": "stat_date"},
        {"label": "充电量(kWh)", "key": "daily_charge_energy"}
      ],
      "page_size": 20,

      # For kpiBar:
      "day_x_field": "hour",   "day_y_field": "power",
      "month_x_field": "day",  "month_y_field": "energy",
      "year_x_field": "month", "year_y_field": "energy",
      "power_x_field": "hour",
      "power_today_field": "power_today",
      "power_yesterday_field": "power_yesterday",
      "power_legend_1": "今日",
      "power_legend_2": "昨日",

      # For kpiBarCompare:
      "x_field": "date",
      "period_a_field": "value_a",
      "period_b_field": "value_b",
      "period_a_name": "本期",
      "period_b_name": "上期",
      "y_axis_name": "kWh",
    }
    """
    t = template_name
    js = {}

    # Detect time fields for smart formatting
    time_fields = detect_time_fields(rows)

    if t in ("lineChart", "barChart"):
        x_field  = mapping.get("x_field", "")
        y_field  = mapping.get("y_field", "")
        # Auto-aggregate if x_field has duplicates (e.g. SQL grouped by extra dimension)
        if needs_aggregation(rows, x_field):
            rows = aggregate_rows_by_key(rows, x_field, [y_field])
            time_fields = detect_time_fields(rows)  # re-detect after aggregation
        js["CHART_X_AXIS_DATA"] = [smart_format_time(str(r.get(x_field, ""))) if x_field in time_fields else str(r.get(x_field, "")) for r in rows]
        js["CHART_SERIES_DATA"] = [r.get(y_field) for r in rows]
        js["CHART_Y_AXIS_NAME"] = mapping.get("y_axis_name", "")
        js["CHART_SERIES_NAME"] = mapping.get("series_name", y_field)

    elif t == "lineChartCompare":
        x_field  = mapping.get("x_field", "")
        y1_field = mapping.get("y1_field", "")
        y2_field = mapping.get("y2_field", "")
        if needs_aggregation(rows, x_field):
            rows = aggregate_rows_by_key(rows, x_field, [y1_field, y2_field])
            time_fields = detect_time_fields(rows)
        js["CHART_X_AXIS_DATA"]   = [smart_format_time(str(r.get(x_field, ""))) if x_field in time_fields else str(r.get(x_field, "")) for r in rows]
        js["CHART_SERIES_DATA_1"] = [r.get(y1_field) for r in rows]
        js["CHART_SERIES_DATA_2"] = [r.get(y2_field) for r in rows]
        js["CHART_Y_AXIS_NAME"]   = mapping.get("y_axis_name", "")
        js["CHART_SERIES_NAME_1"] = mapping.get("series_name_1", y1_field)
        js["CHART_SERIES_NAME_2"] = mapping.get("series_name_2", y2_field)

    elif t == "pieChart":
        name_field  = mapping.get("name_field", "")
        value_field = mapping.get("value_field", "")
        # Auto-aggregate duplicate name_field values (e.g. park_name with multiple enterprises)
        if needs_aggregation(rows, name_field):
            rows = aggregate_rows_by_key(rows, name_field, [value_field])
        js["CHART_SERIES_NAME"] = mapping.get("series_name", "占比")
        js["CHART_DATA"] = [
            {
                "value": r.get(value_field),
                "name":  str(r.get(name_field, "")),
                "itemStyle": {"color": PIE_COLORS[i % len(PIE_COLORS)]}
            }
            for i, r in enumerate(rows)
        ]

    elif t == "dataTable":
        columns = mapping.get("columns", [])
        if not columns and rows:
            # fallback: use all fields from first row
            columns = [{"label": k, "key": k} for k in rows[0].keys()]
        js["TABLE_TITLE"]    = mapping.get("table_title", "数据列表")
        js["TABLE_COLUMNS"]  = columns
        js["TABLE_ROWS_DATA"] = rows
        js["DEFAULT_PAGE_SIZE"] = mapping.get("page_size", 20)

    elif t == "dashboardKpiBar":
        js["TREND_DATA_DAY_X"]       = extract_column(rows, mapping.get("day_x_field", ""))
        js["TREND_DATA_DAY_Y"]       = extract_column(rows, mapping.get("day_y_field", ""))
        js["TREND_DATA_MONTH_X"]     = extract_column(rows, mapping.get("month_x_field", ""))
        js["TREND_DATA_MONTH_Y"]     = extract_column(rows, mapping.get("month_y_field", ""))
        js["TREND_DATA_YEAR_X"]      = extract_column(rows, mapping.get("year_x_field", ""))
        js["TREND_DATA_YEAR_Y"]      = extract_column(rows, mapping.get("year_y_field", ""))
        js["POWER_DATA_X"]           = extract_column(rows, mapping.get("power_x_field", ""))
        js["POWER_DATA_TODAY"]       = extract_column(rows, mapping.get("power_today_field", ""))
        js["POWER_DATA_YESTERDAY"]   = extract_column(rows, mapping.get("power_yesterday_field", ""))
        js["POWER_LEGEND_1"]         = mapping.get("power_legend_1", "今日")
        js["POWER_LEGEND_2"]         = mapping.get("power_legend_2", "昨日")

    elif t == "dashboardKpiBarCompare":
        x_field = mapping.get("x_field", "")
        js["CHART_X_AXIS_DATA"]    = [smart_format_time(str(r.get(x_field, ""))) if x_field in time_fields else str(r.get(x_field, "")) for r in rows]
        js["CHART_PERIOD_A_DATA"]  = extract_column(rows, mapping.get("period_a_field", ""))
        js["CHART_PERIOD_B_DATA"]  = extract_column(rows, mapping.get("period_b_field", ""))
        js["CHART_PERIOD_A_NAME"]  = mapping.get("period_a_name", "本期")
        js["CHART_PERIOD_B_NAME"]  = mapping.get("period_b_name", "上期")
        js["CHART_Y_AXIS_NAME"]    = mapping.get("y_axis_name", "")

    # infoCard: no JS arrays needed
    return js


def build_table_rows_for_text(template_name, mapping, rows):
    """
    Build TABLE_ROWS HTML string for templates that have a text_values TABLE_ROWS field.
    Time fields are automatically detected and formatted.
    Duplicate x/name field values are automatically aggregated (SUM).
    """
    t = template_name
    time_fields = detect_time_fields(rows)

    def get_cols_from_mapping(table_columns, fallback_cols, rows):
        """Extract column key list from LLM table_columns, fallback to fallback_cols."""
        if table_columns and isinstance(table_columns, list):
            # LLM returned table_columns: [{"key": ..., "label": ...}, ...]
            valid_keys = list(rows[0].keys()) if rows else []
            cols = [c["key"] for c in table_columns
                    if isinstance(c, dict) and c.get("key") in valid_keys]
            if cols:
                return cols
        return fallback_cols

    if t in ("lineChart", "barChart"):
        x_field = mapping.get("x_field", "")
        y_field = mapping.get("y_field", "")
        all_fields = list(rows[0].keys()) if rows else []
        value_fields = [f for f in all_fields if f != x_field]
        if needs_aggregation(rows, x_field):
            rows = aggregate_rows_by_key(rows, x_field, value_fields)
            time_fields = detect_time_fields(rows)
        # Use LLM-decided table_columns; fallback to all fields
        table_columns = mapping.get("table_columns", [])
        cols = get_cols_from_mapping(table_columns, [x_field] + value_fields, rows)
        return build_table_rows_html(rows, cols, time_fields=time_fields)

    elif t == "lineChartCompare":
        x_field  = mapping.get("x_field", "")
        y1_field = mapping.get("y1_field", "")
        y2_field = mapping.get("y2_field", "")
        all_fields = list(rows[0].keys()) if rows else []
        value_fields = [f for f in all_fields if f != x_field]
        if needs_aggregation(rows, x_field):
            rows = aggregate_rows_by_key(rows, x_field, value_fields)
            time_fields = detect_time_fields(rows)
        table_columns = mapping.get("table_columns", [])
        cols = get_cols_from_mapping(table_columns, [x_field] + value_fields, rows)
        return build_table_rows_html(rows, cols, time_fields=time_fields)

    elif t == "pieChart":
        name_field  = mapping.get("name_field", "")
        value_field = mapping.get("value_field", "")
        all_fields = list(rows[0].keys()) if rows else []
        value_fields_all = [f for f in all_fields if f != name_field]
        if needs_aggregation(rows, name_field):
            rows = aggregate_rows_by_key(rows, name_field, value_fields_all)
        table_columns = mapping.get("table_columns", [])
        cols = get_cols_from_mapping(table_columns, [name_field, value_field], rows)
        return build_table_rows_html(rows, cols, pie_mode=True, time_fields=time_fields)

    elif t == "dashboardKpiBarCompare":
        # Build comparison table with diff and change rate columns
        x_field = mapping.get("x_field", "")
        a_field = mapping.get("period_a_field", "")
        b_field = mapping.get("period_b_field", "")
        html_rows = []
        for row in rows:
            x_val = row.get(x_field, "")
            if x_field in time_fields and isinstance(x_val, str):
                x_val = smart_format_time(x_val)
            a_val = row.get(a_field)
            b_val = row.get(b_field)
            try:
                a_num = float(a_val) if a_val is not None else 0
                b_num = float(b_val) if b_val is not None else 0
                diff  = round(a_num - b_num, 3)
                rate  = round((diff / b_num * 100), 1) if b_num != 0 else 0
                diff_str = f"+{diff}" if diff >= 0 else str(diff)
                if rate > 0:
                    rate_html = f'<span class="tag-up">+{rate}% ↑</span>'
                elif rate < 0:
                    rate_html = f'<span class="tag-down">{rate}% ↓</span>'
                else:
                    rate_html = f'{rate}%'
            except:
                diff_str  = "-"
                rate_html = "-"
            html_rows.append(
                f"<tr><td>{x_val}</td><td>{a_val}</td><td>{b_val}</td>"
                f"<td>{diff_str}</td><td>{rate_html}</td></tr>"
            )
        return "".join(html_rows)

    return ""


# ─────────────────────────────────────────────────────────────────────────────
# LLM Mapping Prompt Builder
# ─────────────────────────────────────────────────────────────────────────────

MAPPING_PROMPTS = {
    "lineChart": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

Return ONLY valid JSON with this exact structure:
{{
  "x_field": "<field name for X axis (usually a date/time/category field)>",
  "y_field": "<field name for Y axis (the main numeric metric to plot on chart)>",
  "y_axis_name": "<unit label, e.g. kWh, t, 次>",
  "series_name": "<Chinese series name for the chart legend>",
  "table_columns": [
    {{"key": "<field name>", "label": "<Chinese column header>"}},
    ... ONLY include fields that actually exist in the available fields list above. DO NOT invent fields not present in the sample row ...
  ]
}}""",

    "barChart": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

Return ONLY valid JSON with this exact structure:
{{
  "x_field": "<field name for X axis (usually a date/time/category field)>",
  "y_field": "<field name for Y axis (the main numeric metric to plot on chart)>",
  "y_axis_name": "<unit label, e.g. kWh, t, 次>",
  "series_name": "<Chinese series name for the chart legend>",
  "table_columns": [
    {{"key": "<field name>", "label": "<Chinese column header>"}},
    ... ONLY include fields that actually exist in the available fields list above. DO NOT invent fields not present in the sample row ...
  ]
}}""",

    "lineChartCompare": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config for a dual-series line chart.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

Return ONLY valid JSON with this exact structure:
{{
  "x_field": "<field name for X axis>",
  "y1_field": "<field name for first Y series>",
  "y2_field": "<field name for second Y series>",
  "y_axis_name": "<unit label>",
  "series_name_1": "<Chinese name for series 1>",
  "series_name_2": "<Chinese name for series 2>",
  "table_columns": [
    {{"key": "<field name>", "label": "<Chinese column header>"}},
    ... ONLY include fields that actually exist in the available fields list above. DO NOT invent fields not present in the sample row ...
  ]
}}""",

    "pieChart": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config for a pie chart.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

Return ONLY valid JSON with this exact structure:
{{
  "name_field": "<field name for slice labels (category names)>",
  "value_field": "<field name for slice values (numeric)>",
  "series_name": "<Chinese series name>",
  "table_columns": [
    {{"key": "<field name>", "label": "<Chinese column header>"}},
    ... ONLY include fields that actually exist in the available fields list above. DO NOT invent fields not present in the sample row ...
  ]
}}""",

    "dataTable": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config for a data table.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

Return ONLY valid JSON with this exact structure:
{{
  "table_title": "<Chinese table title>",
  "page_size": 20,
  "columns": [
    {{"label": "<Chinese column header>", "key": "<field name>"}},
    ...include ALL relevant fields...
  ]
}}""",

    "infoCard": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config for an info card.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

Return ONLY valid JSON with this exact structure:
{{
  "card_title": "<Chinese card title>",
  "items": [
    {{"label": "<Chinese label>", "field": "<field name>", "unit": "<unit or empty string>"}},
    ...include all relevant fields...
  ]
}}""",

    "dashboardKpiBar": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config for a KPI dashboard.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

The dashboard has: day trend chart, month trend chart, year trend chart, and power comparison chart.
Map the available fields to the best matching chart slots. If a slot has no matching field, use the same field as the closest match.

Return ONLY valid JSON with this exact structure:
{{
  "day_x_field": "<time/hour field for day trend X axis>",
  "day_y_field": "<numeric field for day trend Y axis>",
  "month_x_field": "<date field for month trend X axis>",
  "month_y_field": "<numeric field for month trend Y axis>",
  "year_x_field": "<month field for year trend X axis>",
  "year_y_field": "<numeric field for year trend Y axis>",
  "power_x_field": "<time field for power chart X axis>",
  "power_today_field": "<numeric field for today power>",
  "power_yesterday_field": "<numeric field for yesterday power, or same as today if unavailable>",
  "power_legend_1": "今日",
  "power_legend_2": "昨日"
}}""",

    "dashboardKpiBarCompare": """You are a data mapping assistant.
Given the user query and the available row fields, return a JSON mapping config for a comparison bar chart.

User query: {query}
Available fields (from SQL result): {fields}
Sample row: {sample}

Return ONLY valid JSON with this exact structure:
{{
  "x_field": "<field name for X axis>",
  "period_a_field": "<field name for period A data>",
  "period_b_field": "<field name for period B data>",
  "period_a_name": "<Chinese name for period A>",
  "period_b_name": "<Chinese name for period B>",
  "y_axis_name": "<unit label>"
}}""",
}


def get_mapping_from_llm(template_name, query, rows):
    """Ask LLM for field mapping config only (< 200 tokens output)."""
    if not rows:
        return {}
    fields = list(rows[0].keys())
    sample = json.dumps(rows[0], ensure_ascii=False)
    prompt_tpl = MAPPING_PROMPTS.get(template_name)
    if not prompt_tpl:
        return {}
    prompt = prompt_tpl.format(
        query=query,
        fields=json.dumps(fields, ensure_ascii=False),
        sample=sample
    )
    resp = call_llm(
        prompt,
        system="You are a data mapping assistant. Return ONLY valid JSON, no explanation. For table_columns, ONLY include fields that actually exist in the sample row. DO NOT invent or add fields that are not present in the available fields list.",
        model=LLM_MODEL_FAST,
        max_tokens=1000
    )
    mapping = extract_json(resp)
    return mapping or {}


# ─────────────────────────────────────────────────────────────────────────────
# Text Values Builder (LLM for titles/labels/summaries)
# ─────────────────────────────────────────────────────────────────────────────

TEXT_FIELDS_BY_TEMPLATE = {
    "lineChart": [
        "PAGE_TITLE", "CHART_TITLE", "STAT_TITLE",
        "STAT_LABEL_1_1", "STAT_VALUE_1_1", "STAT_LABEL_1_2", "STAT_VALUE_1_2",
        "STAT_LABEL_1_3", "STAT_VALUE_1_3", "STAT_VALUE_1_3_CLASS",
        "STAT_LABEL_2_1", "STAT_VALUE_2_1", "STAT_LABEL_2_2", "STAT_VALUE_2_2",
        "STAT_LABEL_2_3", "STAT_VALUE_2_3", "STAT_VALUE_2_3_CLASS",
        "TABLE_TITLE", "TABLE_HEADER_1", "TABLE_HEADER_2", "TABLE_HEADER_3", "TABLE_HEADER_4",
        "TABLE_ROWS",
        "SUGGESTION_TITLE", "ANALYSIS_SUMMARY", "OPTIMIZATION_SUGGESTIONS"
    ],
    "barChart": [
        "PAGE_TITLE", "CHART_TITLE", "STAT_TITLE",
        "STAT_LABEL_1_1", "STAT_VALUE_1_1", "STAT_LABEL_1_2", "STAT_VALUE_1_2",
        "STAT_LABEL_1_3", "STAT_VALUE_1_3", "STAT_VALUE_1_3_CLASS",
        "STAT_LABEL_2_1", "STAT_VALUE_2_1", "STAT_LABEL_2_2", "STAT_VALUE_2_2",
        "STAT_LABEL_2_3", "STAT_VALUE_2_3", "STAT_VALUE_2_3_CLASS",
        "TABLE_TITLE", "TABLE_HEADER_1", "TABLE_HEADER_2", "TABLE_HEADER_3", "TABLE_HEADER_4",
        "TABLE_ROWS",
        "SUGGESTION_TITLE", "ANALYSIS_SUMMARY", "OPTIMIZATION_SUGGESTIONS"
    ],
    "lineChartCompare": [
        "PAGE_TITLE", "CHART_TITLE", "STAT_TITLE",
        "STAT_LABEL_1_1", "STAT_VALUE_1_1", "STAT_LABEL_1_2", "STAT_VALUE_1_2",
        "STAT_LABEL_1_3", "STAT_VALUE_1_3", "STAT_VALUE_1_3_CLASS",
        "STAT_LABEL_2_1", "STAT_VALUE_2_1", "STAT_LABEL_2_2", "STAT_VALUE_2_2",
        "STAT_LABEL_2_3", "STAT_VALUE_2_3", "STAT_VALUE_2_3_CLASS",
        "TABLE_TITLE", "TABLE_HEADER_1", "TABLE_HEADER_2", "TABLE_HEADER_3",
        "TABLE_HEADER_4", "TABLE_HEADER_5",
        "TABLE_ROWS",
        "SUGGESTION_TITLE", "ANALYSIS_SUMMARY", "OPTIMIZATION_SUGGESTIONS"
    ],
    "pieChart": [
        "PAGE_TITLE", "CHART_TITLE", "STAT_TITLE",
        "STAT_LABEL_1_1", "STAT_VALUE_1_1", "STAT_LABEL_1_2", "STAT_VALUE_1_2",
        "STAT_LABEL_1_3", "STAT_VALUE_1_3",
        "STAT_LABEL_2_1", "STAT_VALUE_2_1", "STAT_LABEL_2_2", "STAT_VALUE_2_2",
        "STAT_LABEL_2_3", "STAT_VALUE_2_3",
        "TABLE_TITLE", "TABLE_HEADER_1", "TABLE_HEADER_2", "TABLE_HEADER_3",
        "TABLE_HEADER_4", "TABLE_HEADER_5",
        "TABLE_ROWS",
        "SUGGESTION_TITLE", "ANALYSIS_SUMMARY", "OPTIMIZATION_SUGGESTIONS"
    ],
    "dashboardKpiBar": [
        "PAGE_TITLE", "KPI_TITLE",
        "KPI_VALUE_1", "KPI_LABEL_1", "KPI_VALUE_2", "KPI_LABEL_2",
        "KPI_VALUE_3", "KPI_LABEL_3", "KPI_VALUE_4", "KPI_LABEL_4",
        "KPI_VALUE_5", "KPI_LABEL_5", "KPI_VALUE_6", "KPI_LABEL_6",
        "KPI_VALUE_7", "KPI_LABEL_7", "KPI_VALUE_8", "KPI_LABEL_8",
        "KPI_VALUE_9", "KPI_LABEL_9",
        "TREND_TITLE", "TAB_LABEL_1", "TAB_LABEL_2", "TAB_LABEL_3", "POWER_TITLE"
    ],
    "dashboardKpiBarCompare": [
        "PAGE_TITLE", "QUERY_INFO_TEXT",
        "PEAK_TITLE", "PEAK_LABEL_A", "PEAK_VALUE_A", "PEAK_LABEL_B", "PEAK_VALUE_B",
        "TREND_LABEL", "TREND_VALUE",
        "TABLE_HEADER_1", "TABLE_HEADER_2", "TABLE_HEADER_3", "TABLE_HEADER_4", "TABLE_HEADER_5",
        "TABLE_ROWS",
    ],
    "infoCard": ["PAGE_TITLE", "CARD_TITLE", "INFO_ITEMS"],
    "dataTable": [],  # dataTable has no text_values, everything is in js_variables
}


def get_text_values_from_llm(template_name, query, rows, mapping):
    """
    Ask LLM to fill in text/label/summary fields only.
    Pass full rows for richer analysis (since this runs concurrently with fast mapping).
    """
    fields_needed = TEXT_FIELDS_BY_TEMPLATE.get(template_name, [])
    if not fields_needed:
        return {}

    # Build compact stats from full rows for analysis
    total = len(rows)
    numeric_keys = [k for k, v in rows[0].items() if isinstance(v, (int, float))] if rows else []
    stats = {}
    for k in numeric_keys[:8]:
        vals = [r[k] for r in rows if isinstance(r.get(k), (int, float))]
        if vals:
            stats[k] = {
                "sum": round(sum(vals), 3),
                "max": round(max(vals), 3),
                "min": round(min(vals), 3),
                "avg": round(sum(vals) / len(vals), 3)
            }

    # For TABLE_ROWS in chart templates, build it with Python (no LLM needed)
    python_built = {}
    if "TABLE_ROWS" in fields_needed:
        table_rows_html = build_table_rows_for_text(template_name, mapping, rows)
        python_built["TABLE_ROWS"] = table_rows_html
        fields_needed = [f for f in fields_needed if f != "TABLE_ROWS"]

    # Build TABLE_HEADERS_HTML directly from mapping.table_columns (no LLM needed)
    if "TABLE_HEADERS_HTML" in fields_needed or True:  # always build it for chart templates
        table_columns = mapping.get("table_columns", [])
        valid_keys = list(rows[0].keys()) if rows else []
        if table_columns and isinstance(table_columns, list):
            resolved_cols = [
                c for c in table_columns
                if isinstance(c, dict) and c.get("key") in valid_keys
            ]
        else:
            x_field = mapping.get("x_field") or mapping.get("name_field", "")
            all_fields = list(rows[0].keys()) if rows else []
            fallback_keys = ([x_field] + [f for f in all_fields if f != x_field]) if x_field else all_fields
            resolved_cols = [{"key": k, "label": k} for k in fallback_keys]
        headers_html = "".join(f"<th>{c.get('label', c.get('key', ''))}</th>" for c in resolved_cols)
        python_built["TABLE_HEADERS_HTML"] = headers_html
        fields_needed = [f for f in fields_needed if f != "TABLE_HEADERS_HTML"]

    # ── FIX: Build TABLE_HEADER_* directly from mapping.table_columns ──────────
    # This ensures headers always match the exact columns rendered in TABLE_ROWS.
    # Previously, TABLE_HEADER_* was filled by a separate LLM call which could
    # return a different number/order of headers, causing header/body mismatch.
    header_fields = [f for f in fields_needed if f.startswith("TABLE_HEADER_")]
    if header_fields:
        # Resolve the same column list that build_table_rows_for_text uses
        table_columns = mapping.get("table_columns", [])
        if table_columns and isinstance(table_columns, list):
            valid_keys = list(rows[0].keys()) if rows else []
            resolved_cols = [
                {"key": c["key"], "label": c.get("label", c["key"])}
                for c in table_columns
                if isinstance(c, dict) and c.get("key") in valid_keys
            ]
        else:
            # fallback: derive from mapping x_field / value fields
            x_field = mapping.get("x_field") or mapping.get("name_field", "")
            all_fields = list(rows[0].keys()) if rows else []
            fallback_keys = ([x_field] + [f for f in all_fields if f != x_field]) if x_field else all_fields
            resolved_cols = [{"key": k, "label": k} for k in fallback_keys]
        # Assign TABLE_HEADER_1 .. TABLE_HEADER_N from resolved_cols
        for i, col in enumerate(resolved_cols[:len(header_fields)], start=1):
            key = f"TABLE_HEADER_{i}"
            if key in header_fields:
                python_built[key] = col.get("label", col.get("key", ""))
        # Remove header fields from LLM fields_needed (Python already handled them)
        fields_needed = [f for f in fields_needed if f not in header_fields]
    # ── END FIX ────────────────────────────────────────────────────────────────

    # For infoCard INFO_ITEMS, build from mapping
    if template_name == "infoCard" and "INFO_ITEMS" in fields_needed:
        items = mapping.get("items", [])
        if items:
            info_html = "".join(
                f'<div class="info-item"><span class="label">{it.get("label","")}</span>'
                f'<span class="value">{rows[0].get(it.get("field",""), "-") if rows else "-"}'
                f'{it.get("unit","")}</span></div>'
                for it in items
            )
            python_built["INFO_ITEMS"] = info_html
            fields_needed = [f for f in fields_needed if f != "INFO_ITEMS"]

    if not fields_needed:
        return python_built

    prompt = f"""User query: {query}

Data summary (full {total} rows):
{json.dumps(stats, ensure_ascii=False)}

Sample rows (first 3):
{json.dumps(rows[:3], ensure_ascii=False)}

Field mapping used: {json.dumps(mapping, ensure_ascii=False)}

Fill in the following text fields for the report template.
For ANALYSIS_SUMMARY: write 2-4 sentences of Chinese data analysis based on the FULL dataset statistics above.
For OPTIMIZATION_SUGGESTIONS: write 3-4 Chinese suggestions as HTML: <li>· suggestion</li>
For TABLE_HEADER_*: use appropriate Chinese column headers matching the data fields.
For STAT_VALUE_*: compute real values from the data summary above.
For STAT_VALUE_*_CLASS: use "up" for increases/positive, "down" for decreases/negative, "" otherwise.

Fields to fill: {json.dumps(fields_needed, ensure_ascii=False)}

Return ONLY valid JSON with exactly these keys and Chinese string values."""

    resp = call_llm(
        prompt,
        system="You are a Chinese data analyst. Fill in template text fields based on real data. Return ONLY valid JSON.",
        model=LLM_MODEL_FAST,
        max_tokens=1200
    )
    text_vals = extract_json(resp) or {}

    # Merge python-built fields
    text_vals.update(python_built)
    return text_vals


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--query",       required=True)
    parser.add_argument("--template",    required=True)
    parser.add_argument("--result_file", required=True)
    args = parser.parse_args()

    template_name = args.template

    # Read result file
    try:
        with open(args.result_file, 'r', encoding='utf-8') as f:
            result_data = json.load(f)
    except Exception as e:
        print(json.dumps({"status": "error", "message": f"Failed to read result file: {e}"}))
        sys.exit(1)

    # Load HTML template
    html_path = os.path.join(TEMPLATES_DIR, f"{template_name}.html")
    if not os.path.exists(html_path):
        print(json.dumps({"status": "error", "message": f"Template HTML not found: {html_path}"}))
        sys.exit(1)
    html = read_file(html_path)

    # ── Fixed template branch ────────────────────────────────────────────────
    if isinstance(result_data, dict) and result_data.get("is_fixed_template"):
        placeholder_values = result_data.get("placeholder_values", {})
        summary            = result_data.get("summary", "数据分析报告已生成。")
        if not placeholder_values:
            print(json.dumps({"status": "error", "message": "Fixed template: placeholder_values is empty."}))
            sys.exit(1)
        html = inject_placeholders(html, placeholder_values)

    # ── Normal template branch (v3: LLM decides, Python extracts) ───────────
    else:
        rows = result_data.get("rows", [])
        if not rows:
            print(json.dumps({"status": "error", "message": "Normal template: rows is empty."}))
            sys.exit(1)

        t0 = time.time()

        # Step A: Get field mapping from LLM (fast, < 400 tokens output)
        mapping = get_mapping_from_llm(template_name, args.query, rows)
        t_mapping = round(time.time() - t0, 2)
        sys.stderr.write(f"[timing] mapping LLM: {t_mapping}s\n")

        # Step B: Concurrently:
        #   B1 - Python extracts all JS arrays (instant, no LLM)
        #   B2 - LLM fills text/label/summary fields (uses FULL rows for analysis)
        def build_js():
            return python_extract_js_vars(template_name, mapping, rows)

        def build_text():
            return get_text_values_from_llm(template_name, args.query, rows, mapping)

        js_vars    = None
        text_vals  = None
        with ThreadPoolExecutor(max_workers=2) as executor:
            future_js   = executor.submit(build_js)
            future_text = executor.submit(build_text)
            js_vars   = future_js.result()
            text_vals = future_text.result()

        t_total = round(time.time() - t0, 2)
        sys.stderr.write(f"[timing] total step3: {t_total}s\n")

        # Inject text values
        if text_vals:
            for key, val in text_vals.items():
                if val is not None:
                    html = html.replace("{{" + key + "}}", str(val))

        # Inject JS variables (Python-built arrays)
        html = inject_js_vars(html, js_vars)

        # Extract summary for output
        summary = text_vals.get("ANALYSIS_SUMMARY", "数据分析报告已生成。") if text_vals else "数据分析报告已生成。"

    # ── Write output HTML ────────────────────────────────────────────────────
    timestamp = int(time.time())
    out_file  = f"report_{timestamp}.html"
    out_path  = os.path.join(OUTPUT_DIR, out_file)
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write(html)

    try:
        os.remove(args.result_file)
    except:
        pass

    print(json.dumps({
        "status":   "success",
        "message":  "Report generated successfully",
        "template": template_name,
        "address":  f"/skill/sql-report-generator/output/{out_file}",
        "summary":  summary.strip() if isinstance(summary, str) else str(summary)
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()