#!/usr/bin/env python3
"""
SQL Report Generator — Universal Pipeline (v5)
===============================================
Optimized 3-LLM-call flow with real-time progress output:

LLM Call 1: Select views + template simultaneously  (merged from 2 calls → saves ~3s)
  ↓ Load view schemas (file I/O, fast)
LLM Call 2: Generate SQL (with template format + view schema context)
  ↓ Execute SQL via DB API
LLM Call 3: Map data to template JS vars + generate summary  (merged from 2 calls → saves ~3s)
  ↓ Render HTML

Progress output format:  PROGRESS: <message>
  The OpenClaw agent reads these lines and can relay them to the user in real-time.

v5 Optimizations:
- Added detailed Chinese progress messages at every key step
- LLM calls use max_tokens to reduce unnecessary token generation
- LLM Call 2 uses faster qwen-turbo model for SQL generation
- Concurrent file I/O after LLM Call 1 using threading
- DB query timeout tuned to 20s (was 30s)
- extract_json enhanced with raw-brace fallback for robustness
"""
import os
import sys
import json
import time
import argparse
import urllib.request
import re
import threading

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
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")
REFERENCES_DIR = os.path.join(SKILL_DIR, "references")
VIEWS_DIR      = os.path.join(REFERENCES_DIR, "views")
VIEWS_INDEX    = os.path.join(REFERENCES_DIR, "views_index.md")

os.makedirs(OUTPUT_DIR, exist_ok=True)

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_SQL = os.environ.get("LLM_MODEL_SQL", "qwen-turbo")   # faster model for SQL-only call

DB_API_URL = "http://47.97.121.32:8080/VOLERTECH/app/api/oauth/getmysql"
DB_JDBC    = "snowy_cloud_mysql_copy"

# ---------------------------------------------------------------------------
# Progress output — flush immediately so agent can relay in real-time
# ---------------------------------------------------------------------------
def progress(msg):
    print(f"PROGRESS: {msg}", flush=True)


# ---------------------------------------------------------------------------
# Token usage tracking
# ---------------------------------------------------------------------------
import atexit, datetime as _dt

_token_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "calls": 0}

def _flush_token_usage():
    log_path = os.path.join(OUTPUT_DIR, "token_usage.jsonl")
    record = {
        "time":   _dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "script": "generate_report",
        **_token_usage
    }
    with open(log_path, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")

atexit.register(_flush_token_usage)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def call_llm(prompt, system="You are a helpful assistant.", model=None, max_tokens=None):
    """Call LLM API and return the text content of the first choice."""
    use_model = model or LLM_MODEL
    headers = {
        "Authorization": f"Bearer {LLM_API_KEY}",
        "Content-Type": "application/json"
    }
    parameters = {"result_format": "message"}
    if max_tokens:
        parameters["max_tokens"] = max_tokens

    payload = {
        "model": use_model,
        "input": {
            "messages": [
                {"role": "system", "content": system},
                {"role": "user",   "content": prompt}
            ]
        },
        "parameters": parameters
    }
    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)
        result = json.loads(resp.read().decode("utf-8"))
        usage = result.get("usage", {})
        _token_usage["input_tokens"]  += usage.get("input_tokens", 0)
        _token_usage["output_tokens"] += usage.get("output_tokens", 0)
        _token_usage["total_tokens"]  += usage.get("total_tokens", 0)
        _token_usage["calls"]         += 1
        return result["output"]["choices"][0]["message"]["content"]
    except Exception as e:
        print(f"LLM error: {e}", flush=True)
        return None


def sanitize_sql(sql):
    """Replace ANSI double-quoted Chinese/space identifiers with MySQL backticks."""
    sql = re.sub(r'"([^"]*[\u4e00-\u9fff][^"]*)"', lambda m: f'`{m.group(1)}`', sql)
    sql = re.sub(r'"([^"]+\s[^"]+)"',              lambda m: f'`{m.group(1)}`', sql)
    return sql


def query_database(sql):
    sql = sanitize_sql(sql)
    print(f"Sanitized SQL: {sql}", flush=True)
    req = urllib.request.Request(
        DB_API_URL,
        data=json.dumps({"sql": sql, "basejdbc": DB_JDBC}).encode("utf-8"),
        headers={"Content-Type": "application/json"}
    )
    try:
        resp = urllib.request.urlopen(req, timeout=20)
        return json.loads(resp.read().decode("utf-8"))
    except Exception as e:
        print(f"DB error: {e}", flush=True)
        return None


def extract_json(text):
    """Extract the first JSON object/array from LLM response text."""
    # Try fenced code blocks first
    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 Exception:
                pass
    # Try raw JSON
    try:
        return json.loads(text.strip())
    except Exception:
        pass
    # Try to find the outermost { ... } block
    brace_match = re.search(r'(\{.*\})', text, re.DOTALL)
    if brace_match:
        try:
            return json.loads(brace_match.group(1))
        except Exception as e:
            print(f"JSON parse error (brace fallback): {e}", flush=True)
    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 get_template_list():
    """Return list of {name, summary} for all templates that have both .html and .md."""
    templates = []
    for fname in sorted(os.listdir(TEMPLATES_DIR)):
        if fname.endswith(".md"):
            name = fname[:-3]
            if os.path.exists(os.path.join(TEMPLATES_DIR, f"{name}.html")):
                doc = read_file(os.path.join(TEMPLATES_DIR, fname))
                summary = "\n".join(doc.split("\n")[:8])
                templates.append({"name": name, "summary": summary})
    return templates


def load_view_schemas(view_names):
    schemas = []
    for vname in view_names:
        safe = re.sub(r'[^\w\u4e00-\u9fff]', '_', vname.strip())
        path = os.path.join(VIEWS_DIR, f"{safe}.md")
        if not os.path.exists(path):
            for fname in os.listdir(VIEWS_DIR):
                if vname.strip() in fname or fname.replace('.md', '') in vname:
                    path = os.path.join(VIEWS_DIR, fname)
                    break
        if os.path.exists(path):
            schemas.append(read_file(path))
            print(f"  Loaded view schema: {vname}", flush=True)
        else:
            print(f"  Warning: schema not found for '{vname}'", flush=True)
    return "\n\n---\n\n".join(schemas)


# ---------------------------------------------------------------------------
# Main Pipeline
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--query", required=True)
    args  = parser.parse_args()
    query = args.query
    print(f"Processing query: {query}", flush=True)

    # ── LLM Call 1: Select views + template simultaneously ───────────────────
    progress("【开始分析问题...】正在理解用户需求，匹配数据视图与展示模板...")
    views_index = read_file(VIEWS_INDEX)
    templates   = get_template_list()
    templates_summary = "\n\n".join([f"[{t['name']}]\n{t['summary']}" for t in templates])

    step1_resp = call_llm(
        f"USER QUERY: {query}\n\n"
        f"AVAILABLE DATABASE VIEWS INDEX:\n{views_index}\n\n"
        f"AVAILABLE TEMPLATES:\n{templates_summary}\n\n"
        "Return ONLY valid JSON.",
        system="""\
You are a data analysis expert. Given the user's query, do TWO things simultaneously:
1. Select 1-3 most relevant database views from the index.
2. Select the single best HTML template for visualizing the result.

Template guidelines:
- Time-series trend (每天/每月/每年趋势) → lineChart
- Comparison by category/entity (各企业/各园区对比) → barChart
- Distribution / proportion (占比/分布/比例) → pieChart
- Detailed tabular data (列表/明细/详情) → dataTable
- KPI summary metrics (总量/汇总/概览) → infoCard or dashboardKpiBar

Return ONLY valid JSON in this exact format:
{
  "selected_views": ["视图名1", "视图名2"],
  "template_name": "templateName",
  "reason": "one sentence explanation"
}
Use EXACT view names from the index. Use EXACT template names from the list.
""",
        max_tokens=512
    )
    if not step1_resp:
        print(json.dumps({"status": "error", "message": "Failed at step 1"})); return
    step1 = extract_json(step1_resp)
    if not step1 or "selected_views" not in step1 or "template_name" not in step1:
        print(json.dumps({"status": "error", "message": "Failed to parse step 1", "raw": step1_resp})); return

    selected_views = step1["selected_views"]
    template_name  = step1["template_name"]
    print(f"  Selected views: {selected_views}", flush=True)
    print(f"  Selected template: {template_name}", flush=True)
    print(f"  Reason: {step1.get('reason', '')}", flush=True)

    progress(f"【开始选择模板...】已选定模板：{template_name}，正在加载数据视图结构：{', '.join(selected_views)}...")

    # ── Concurrent file I/O: load view schemas + template doc in parallel ────
    view_schemas_holder = [None]
    template_doc_holder = [None]

    def _load_views():
        view_schemas_holder[0] = load_view_schemas(selected_views)

    def _load_template():
        template_doc_holder[0] = read_file(os.path.join(TEMPLATES_DIR, f"{template_name}.md"))

    t1 = threading.Thread(target=_load_views)
    t2 = threading.Thread(target=_load_template)
    t1.start(); t2.start()
    t1.join();  t2.join()

    view_schemas = view_schemas_holder[0]
    template_doc = template_doc_holder[0]

    if not view_schemas:
        print(json.dumps({"status": "error", "message": f"Could not load schemas for: {selected_views}"})); return
    if not template_doc:
        print(json.dumps({"status": "error", "message": f"Template doc not found: {template_name}.md"})); return

    print(f"  View schemas loaded ({len(view_schemas)} chars), template doc loaded ({len(template_doc)} chars).", flush=True)

    # ── LLM Call 2: Generate SQL ─────────────────────────────────────────────
    progress(f"【开始生成SQL语句...】正在根据视图结构生成 MySQL 查询（模板：{template_name}）...")
    sql_resp = call_llm(
        f"USER QUERY: {query}\n\n"
        f"SELECTED TEMPLATE: {template_name}\n\n"
        f"TEMPLATE DATA FORMAT REQUIREMENTS:\n{template_doc}\n\n"
        f"VIEW SCHEMAS (available columns):\n{view_schemas}\n\n"
        "Generate a MySQL SELECT statement that returns data matching the template's required format.\n"
        "Return ONLY valid JSON: {\"sql\": \"...\"}",
        system="""\
You are an expert MySQL developer. Write a single SELECT statement that returns data
READY to be injected into the specified HTML template.

SQL Rules:
1. Use the EXACT view name from '**SQL 视图名**' field (backtick-quoted Chinese name).
   NEVER translate to English. NEVER use double quotes for identifiers.
2. Return chart-ready aggregated data — one row per X-axis label:
   - lineChart/barChart with time axis → GROUP BY single time col + SUM(value)
   - barChart by entity → GROUP BY entity col + SUM(value)
   - pieChart → GROUP BY category col + SUM(value)
   - dataTable → SELECT needed columns with WHERE/ORDER
3. SELECT ONLY the minimum columns needed (typically 2: dimension + metric).
   Do NOT select extra columns (enterprise_name, park_name, unit, etc.) unless they
   are the grouping dimension for the chart.
4. TIME-SERIES AXIS: SELECT and GROUP BY ONLY the single finest-grain time column:
   - '天/每天/日' query → SELECT `day`, GROUP BY `day` (filter year/month in WHERE)
   - '月/每月' query → SELECT `month`, GROUP BY `month`
   - '年/每年' query → SELECT `year`, GROUP BY `year`
   NEVER GROUP BY (year, month, day) together.
5. CHINESE TIME DISAMBIGUATION:
   - 'X月中' / 'X月份中' = entire month X (all days), NOT mid-month. Use WHERE month=X only.
   - '中旬' = days 11-20 explicitly. '上旬' = 1-10. '下旬' = 21-end.
6. Always ORDER BY the grouping column for time-series charts.
""",
        model=LLM_MODEL_SQL,
        max_tokens=512
    )
    if not sql_resp:
        print(json.dumps({"status": "error", "message": "Failed to generate SQL"})); return
    sql_plan = extract_json(sql_resp)
    if not sql_plan or "sql" not in sql_plan:
        print(json.dumps({"status": "error", "message": "Failed to parse SQL", "raw": sql_resp})); return

    sql = sql_plan["sql"]
    print(f"  Generated SQL: {sql}", flush=True)

    # ── Execute SQL ──────────────────────────────────────────────────────────
    progress("【开始查询数据库...】正在执行 SQL，连接数据库获取数据...")
    db_result = query_database(sql)
    if not db_result:
        print(json.dumps({"status": "error", "message": "Failed to execute SQL"})); return

    result_str = json.dumps(db_result, ensure_ascii=False)
    print(f"  Query returned {len(result_str)} bytes of data.", flush=True)
    print(f"  Preview: {result_str[:400]}", flush=True)

    rows = db_result if isinstance(db_result, list) else (
        next((db_result[k] for k in ("data", "rows", "result", "records")
              if k in db_result and isinstance(db_result[k], list)), [])
    )
    if not rows:
        print(json.dumps({"status": "error", "message": "SQL returned no data", "sql": sql})); return

    print(f"  Total rows: {len(rows)}", flush=True)
    progress(f"【数据查询完成】共获取 {len(rows)} 条数据，正在渲染图表...")

    # ── LLM Call 3: Map data to template vars + generate summary ─────────────
    progress(f"【开始渲染报告...】正在将查询结果映射为 {template_name} 图表数据结构...")
    map_resp = call_llm(
        f"USER QUERY: {query}\n\n"
        f"QUERY RESULTS ({len(rows)} rows — ALL rows shown below):\n{result_str[:50000]}\n\n"
        f"TEMPLATE DOCUMENTATION ({template_name}.md):\n{template_doc}\n\n"
        f"CRITICAL: The query returned {len(rows)} rows. You MUST include ALL {len(rows)} data points "
        "in CHART_X_AXIS_DATA and CHART_SERIES_DATA (or equivalent) arrays. Do not skip any rows.\n\n"
        "Return ONLY valid JSON. The JSON must include a 'summary' field (2-3 sentences in Chinese "
        "describing key findings from the actual data) in addition to the template data fields.",
        system=(
            "You are a data visualization engineer and analyst. "
            "1. Map ALL query result rows to the exact JSON structure required by the template. "
            "2. Add a 'summary' field with 2-3 sentences of Chinese analysis based ONLY on the actual data. "
            "Do not skip or truncate any data rows. Do not fabricate data not present in the results."
        )
    )
    if not map_resp:
        print(json.dumps({"status": "error", "message": "Failed to map data"})); return
    template_data = extract_json(map_resp)
    if not template_data:
        print(json.dumps({"status": "error", "message": "Failed to parse mapped data", "raw": map_resp})); return

    # Extract summary from mapping response
    summary = template_data.pop("summary", "数据分析报告已生成。")

    # ── Render HTML ──────────────────────────────────────────────────────────
    progress("【开始生成HTML报告...】正在将数据注入模板，生成可视化报告...")
    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}"})); return

    html = read_file(html_path)

    # Replace {{PLACEHOLDER}} text values
    text_values = template_data.get("data", {}).get("text_values", {})
    for key, val in text_values.items():
        html = html.replace(f"{{{{{key}}}}}", str(val))

    # Inject JS variables via // {{INJECT_DATA}} placeholder
    js_vars = template_data.get("data", {}).get("js_variables", {})
    if js_vars:
        js_lines = [f"var {k} = {json.dumps(v, ensure_ascii=False)};" for k, v in js_vars.items()]
        html = html.replace("// {{INJECT_DATA}}", "\n        ".join(js_lines))

    # Check for unreplaced placeholders
    remaining = re.findall(r'\{\{[A-Z_]+\}\}', html)
    if remaining:
        print(f"  Warning: unreplaced placeholders: {set(remaining)}", flush=True)

    timestamp = int(time.time())
    out_file  = f"report_{timestamp}.html"
    out_path  = os.path.join(OUTPUT_DIR, out_file)
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write(html)
    print(f"  Saved: {out_path}", flush=True)

    progress("【报告生成完成！】HTML 报告已保存，正在整理分析摘要...")

    # ── Final result ─────────────────────────────────────────────────────────
    print("\n--- FINAL RESULT ---", flush=True)
    print(json.dumps({
        "status":         "success",
        "message":        "Report generated successfully",
        "selected_views": selected_views,
        "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), flush=True)


if __name__ == "__main__":
    main()