#!/usr/bin/env python3
"""wx-quote-assistant — 价格计算引擎

加载定价规则（自然语言），结合产品明细和用户需求，输出结构化价格结果。
核心计算由 LLM 辅助执行，输出标记「AI 辅助生成，请人工复核」。

CLI 用法:
  python3 price_calculator.py calculate \
    --products <json_path> \
    --rules <txt_path> \
    [--tax-rate <float>] \
    [--shipping <float>] \
    --output <json_path>
"""

import argparse
import json
import os
import sys


def load_rules(rules_path: str) -> str:
    """加载自然语言定价规则。"""
    if not os.path.exists(rules_path):
        return ""
    with open(rules_path, "r", encoding="utf-8") as f:
        return f.read().strip()


def load_products(products_path: str) -> list[dict]:
    """加载产品明细 JSON。"""
    with open(products_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    # 支持两种格式: {"products": [...]} 或直接 [...]
    if isinstance(data, dict) and "products" in data:
        return data["products"]
    if isinstance(data, list):
        return data
    return []


def compute_subtotals(products: list[dict]) -> list[dict]:
    """计算每个产品的小计 (数量×单价×折扣)。"""
    result = []
    for p in products:
        qty = float(p.get("qty", 0))
        unit_price = float(p.get("unit_price", 0))
        discount = float(p.get("discount", 1.0))  # 默认无折扣=1.0

        p["subtotal"] = round(qty * unit_price * discount, 2)
        p["discount_amount"] = round(qty * unit_price * (1 - discount), 2)
        result.append(p)
    return result


def build_calculation_context(products: list[dict], rules_text: str,
                              tax_rate: float = 0.13,
                              shipping: float = 0.0) -> dict:
    """构建传给 LLM 的价格计算上下文。

    此函数输出的 JSON 将被传回 Skill，由 Skill 调用 LLM 进行定价规则解析和最终计算。
    输出包含：
    - 原始产品明细（含小计）
    - 定价规则文本
    - 建议的计算步骤
    - 一个标记，要求 LLM 输出结构化结果
    """
    products_with_subtotals = compute_subtotals(products)
    subtotal_sum = round(sum(p["subtotal"] for p in products_with_subtotals), 2)
    discount_sum = round(sum(p["discount_amount"] for p in products_with_subtotals), 2)

    context = {
        "meta": {
            "ai_note": "AI 辅助生成，请人工复核",
            "calculation_mode": "llm_assisted",
        },
        "products": products_with_subtotals,
        "subtotal_sum": subtotal_sum,
        "individual_discount_sum": discount_sum,
        "tax_rate": tax_rate,
        "shipping_fee": shipping,
        "pricing_rules": rules_text if rules_text else "无特殊定价规则",
        "expected_output": {
            "instructions": (
                "请根据以上定价规则和产品明细，计算最终价格。输出以下结构：\n"
                "1. discount_desc: 适用的折扣说明（文字）\n"
                "2. discount_amount: 整单折扣额（数字）\n"
                "3. discounted_total: 折扣后合计（数字）= subtotal_sum - discount_amount\n"
                "4. tax_amount: 税费 = discounted_total x tax_rate（保留2位小数）\n"
                "5. shipping_fee: 运费（数字）\n"
                "6. grand_total: 总计 = discounted_total + tax_amount + shipping_fee（保留2位小数）\n"
                "7. remarks: 计算过程中的注意事项\n"
                "如定价规则未明确说明，折扣额为 0，折扣说明为[无]。"
            ),
            "fields": [
                "discount_desc", "discount_amount", "discounted_total",
                "tax_amount", "shipping_fee", "grand_total", "remarks"
            ],
        },
    }
    return context


def cmd_calculate(args) -> None:
    """执行价格计算（生成计算上下文供 LLM 使用）。"""
    products = load_products(args.products)
    if not products:
        print(json.dumps({"error": "No products found in input"}, ensure_ascii=False), file=sys.stderr)
        sys.exit(1)

    rules_text = load_rules(args.rules) if args.rules else ""
    tax_rate = float(args.tax_rate) if args.tax_rate else 0.13
    shipping = float(args.shipping) if args.shipping else 0.0

    context = build_calculation_context(products, rules_text, tax_rate, shipping)

    # 确保输出目录存在
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    with open(args.output, "w", encoding="utf-8") as f:
        json.dump(context, f, ensure_ascii=False, indent=2)

    print(json.dumps({
        "ok": True,
        "output": args.output,
        "product_count": len(products),
        "subtotal_sum": context["subtotal_sum"],
        "note": "计算上下文已生成。请将 pricing_context 传给 LLM 进行定价规则解析，然后调用 save-result 写入最终结果。"
    }, ensure_ascii=False, indent=2))


def cmd_save_result(args) -> None:
    """保存 LLM 返回的定价结果。"""
    with open(args.result, "r", encoding="utf-8") as f:
        result = json.load(f)

    # 确保必需字段存在
    required = ["discount_desc", "discount_amount", "discounted_total",
                "tax_amount", "shipping_fee", "grand_total", "remarks"]
    for field in required:
        if field not in result:
            result[field] = 0 if "amount" in field or field == "total" else ""

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    with open(args.output, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)

    print(json.dumps({"ok": True, "output": args.output}, ensure_ascii=False))


def main():
    parser = argparse.ArgumentParser(description="wx-quote-assistant Price Calculator")
    sub = parser.add_subparsers(dest="command", required=True)

    p_calc = sub.add_parser("calculate", help="Build pricing context for LLM")
    p_calc.add_argument("--products", required=True, help="Products JSON file path")
    p_calc.add_argument("--rules", default=None, help="Pricing rules text file path")
    p_calc.add_argument("--tax-rate", default=None, type=float, help="Tax rate (default: 0.13)")
    p_calc.add_argument("--shipping", default=None, type=float, help="Shipping fee (default: 0)")
    p_calc.add_argument("--output", required=True, help="Output JSON path for pricing context")
    p_calc.set_defaults(func=cmd_calculate)

    p_save = sub.add_parser("save-result", help="Save LLM pricing result")
    p_save.add_argument("--result", required=True, help="LLM result JSON file")
    p_save.add_argument("--output", required=True, help="Final priced result JSON path")
    p_save.set_defaults(func=cmd_save_result)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
