#!/usr/bin/env python3
"""
Step 1: Select template based on user query and data structure.
"""
import os
import sys
import json
import argparse
import urllib.request
import re

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

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")

# ---------------------------------------------------------------------------
# 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():
    output_dir = os.path.join(SKILL_DIR, "output")
    os.makedirs(output_dir, exist_ok=True)
    log_path = os.path.join(output_dir, "token_usage.jsonl")
    record = {
        "time":   _dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "script": "step1_select",
        **_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)

def call_llm(prompt, system="You are a helpful assistant.", model=None, max_tokens=None):
    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(json.dumps({"status": "error", "message": f"LLM error: {e}"}))
        sys.exit(1)

def extract_json(text):
    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
    brace_match = re.search(r'(\{.*\})', text, re.DOTALL)
    if brace_match:
        try: return json.loads(brace_match.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 get_template_list():
    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 main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--query", required=True)
    parser.add_argument("--data", required=True)
    args = parser.parse_args()
    
    try:
        data = json.loads(args.data)
    except Exception as e:
        print(json.dumps({"status": "error", "message": f"Invalid JSON data: {e}"}))
        sys.exit(1)
    
    templates = get_template_list()
    templates_summary = "\n\n".join([f"[{t['name']}]\n{t['summary']}" for t in templates])
    data_sample = json.dumps(data[:5] if isinstance(data, list) else data, ensure_ascii=False)[:1000]

    step1_resp = call_llm(
        f"USER QUERY: {args.query}\n\nDATA SAMPLE: {data_sample}\n\nAVAILABLE TEMPLATES:\n{templates_summary}\n\nReturn ONLY valid JSON.",
        system="""You are a data visualization expert. Given the user's query and data sample, 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:
{"template_name": "templateName", "reason": "one sentence explanation"}
Use EXACT template names from the list.""",
        max_tokens=512
    )
    
    if not step1_resp:
        print(json.dumps({"status": "error", "message": "Failed to get response from LLM"}))
        sys.exit(1)
        
    step1 = extract_json(step1_resp)
    if not step1 or "template_name" not in step1:
        print(json.dumps({"status": "error", "message": "Failed to parse step 1 JSON", "raw": step1_resp}))
        sys.exit(1)

    # Output pure JSON for the next step to consume
    print(json.dumps({
        "status": "success",
        "template_name": step1["template_name"],
        "reason": step1.get("reason", "")
    }, ensure_ascii=False))

if __name__ == "__main__":
    main()