import argparse
import csv
import json
import os
import random
import sys
from collections import defaultdict

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from agentic.data_analysis_env import AgenticToolEnv, get_agentic_tools


CHANNELS = ["抖音", "微信", "百度", "小红书", "自然流量"]
CATEGORIES = ["美妆", "数码", "服饰", "食品", "家居"]
PROVINCES = ["广东", "浙江", "北京", "上海", "四川", "湖北"]

METRIC_DOCS = [
    {
        "id": "metric_refund_amount",
        "title": "退款金额",
        "text": "退款金额使用 orders 表的 refund_amount 字段求和；按渠道、品类、省份或时间过滤后再聚合。",
    },
    {
        "id": "metric_net_revenue",
        "title": "净收入",
        "text": "净收入 = GMV - 退款金额。GMV 使用 orders.gmv，退款金额使用 orders.refund_amount。",
    },
    {
        "id": "metric_refund_rate",
        "title": "退款率",
        "text": "退款率 = 退款金额 / GMV。按品类或渠道对 refund_amount 与 gmv 分别求和后再相除。",
    },
    {
        "id": "metric_roas",
        "title": "ROAS",
        "text": "ROAS = GMV / 广告花费。GMV 来自 orders 表，广告花费来自 marketing 表的 spend 字段。",
    },
    {
        "id": "metric_new_users",
        "title": "新增用户",
        "text": "新增用户使用 users 表的 new_users 字段求和；活跃用户使用 active_users 字段求和。",
    },
    {
        "id": "schema_orders",
        "title": "orders 表字段",
        "text": "orders 包含 month、quarter、channel、category、province、gmv、refund_amount、orders_count。",
    },
    {
        "id": "schema_marketing",
        "title": "marketing 表字段",
        "text": "marketing 包含 month、quarter、channel、spend、impressions、clicks。",
    },
    {
        "id": "schema_users",
        "title": "users 表字段",
        "text": "users 包含 month、quarter、channel、province、active_users、new_users。",
    },
]


def mkdir(path):
    os.makedirs(path, exist_ok=True)


def write_csv(path, rows):
    if not rows:
        raise ValueError(f"empty csv rows: {path}")
    with open(path, "w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def relpath(path):
    cwd = os.getcwd()
    try:
        return os.path.relpath(path, cwd)
    except ValueError:
        return path


def quarter_of(month):
    return f"Q{(month - 1) // 3 + 1}"


def generate_tables(output_dir, seed):
    rng = random.Random(seed)
    table_dir = os.path.join(output_dir, "tables")
    mkdir(table_dir)

    orders = []
    for month in range(1, 13):
        season = 1.0 + (0.18 if month in {6, 11, 12} else 0.0) + (0.08 if month in {1, 2} else 0.0)
        for channel_idx, channel in enumerate(CHANNELS):
            for category_idx, category in enumerate(CATEGORIES):
                for province_idx, province in enumerate(PROVINCES):
                    base = 42000 + channel_idx * 4600 + category_idx * 3100 + province_idx * 1800
                    gmv = round((base + rng.uniform(-5200, 6200)) * season, 2)
                    refund_rate = 0.025 + 0.006 * category_idx + 0.003 * channel_idx + rng.uniform(0.0, 0.018)
                    refund_amount = round(gmv * refund_rate, 2)
                    orders_count = max(1, int(gmv / (90 + category_idx * 12 + rng.uniform(-8, 8))))
                    orders.append(
                        {
                            "month": month,
                            "quarter": quarter_of(month),
                            "channel": channel,
                            "category": category,
                            "province": province,
                            "gmv": gmv,
                            "refund_amount": refund_amount,
                            "orders_count": orders_count,
                        }
                    )

    marketing = []
    for month in range(1, 13):
        for channel_idx, channel in enumerate(CHANNELS):
            spend = round(18000 + channel_idx * 3500 + month * 420 + rng.uniform(-1800, 2200), 2)
            clicks = max(1, int(spend / (2.2 + channel_idx * 0.25 + rng.uniform(-0.2, 0.25))))
            impressions = clicks * rng.randint(18, 36)
            marketing.append(
                {
                    "month": month,
                    "quarter": quarter_of(month),
                    "channel": channel,
                    "spend": spend,
                    "impressions": impressions,
                    "clicks": clicks,
                }
            )

    users = []
    for month in range(1, 13):
        for channel_idx, channel in enumerate(CHANNELS):
            for province_idx, province in enumerate(PROVINCES):
                active_users = max(100, int(5200 + month * 110 + channel_idx * 520 + province_idx * 260 + rng.uniform(-420, 520)))
                new_users = max(20, int(active_users * (0.08 + channel_idx * 0.008 + rng.uniform(0.0, 0.04))))
                users.append(
                    {
                        "month": month,
                        "quarter": quarter_of(month),
                        "channel": channel,
                        "province": province,
                        "active_users": active_users,
                        "new_users": new_users,
                    }
                )

    paths = {
        "orders": os.path.join(table_dir, "orders.csv"),
        "marketing": os.path.join(table_dir, "marketing.csv"),
        "users": os.path.join(table_dir, "users.csv"),
    }
    write_csv(paths["orders"], orders)
    write_csv(paths["marketing"], marketing)
    write_csv(paths["users"], users)
    return {"orders": orders, "marketing": marketing, "users": users, "paths": paths}


def table_specs(paths):
    return [
        {
            "name": "orders",
            "path": relpath(paths["orders"]),
            "description": "订单聚合表，包含 GMV、退款金额和订单数。",
        },
        {
            "name": "marketing",
            "path": relpath(paths["marketing"]),
            "description": "投放聚合表，包含广告花费、曝光和点击。",
        },
        {
            "name": "users",
            "path": relpath(paths["users"]),
            "description": "用户活跃聚合表，包含活跃用户和新增用户。",
        },
    ]


def sum_field(rows, field, **filters):
    total = 0.0
    for row in rows:
        if all(row.get(key) == value for key, value in filters.items()):
            total += float(row[field])
    return round(total, 2)


def top_group(rows, group_key, value_field, filters):
    agg = defaultdict(float)
    for row in rows:
        if all(row.get(key) == value for key, value in filters.items()):
            agg[row[group_key]] += float(row[value_field])
    key, value = max(agg.items(), key=lambda item: item[1])
    return key, round(value, 2)


def refund_rate_by_category(rows, month):
    agg = defaultdict(lambda: {"gmv": 0.0, "refund": 0.0})
    for row in rows:
        if row["month"] == month:
            agg[row["category"]]["gmv"] += float(row["gmv"])
            agg[row["category"]]["refund"] += float(row["refund_amount"])
    rates = {
        category: values["refund"] / values["gmv"]
        for category, values in agg.items()
        if values["gmv"] > 0
    }
    category, rate = max(rates.items(), key=lambda item: item[1])
    return category, round(rate, 4)


def monthly_gmv_by_channel(orders, month, channel):
    return sum(float(row["gmv"]) for row in orders if row["month"] == month and row["channel"] == channel)


def marketing_spend(marketing, month, channel):
    return sum(float(row["spend"]) for row in marketing if row["month"] == month and row["channel"] == channel)


def build_tool_call(name, arguments):
    return {"name": name, "arguments": arguments}


def make_sample(sample_id, question, answer, checks, tables, expected_tools, actions, task_type, max_turns=4):
    sample = {
        "id": sample_id,
        "question": question,
        "tables": tables,
        "documents": METRIC_DOCS,
        "tools": ["retriever", "python_executor", "calculator"],
        "expected_tools": expected_tools,
        "answer": answer,
        "checks": checks,
        "max_turns": max_turns,
        "metadata": {"task_type": task_type},
    }
    env = AgenticToolEnv(sample)
    trajectory = []
    for action in actions:
        trajectory.append({"role": "assistant", "tool_call": action})
        result = env.execute(action["name"], action["arguments"])
        trajectory.append({"role": "tool", "name": action["name"], "content": result})
    trajectory.append({"role": "assistant", "content": answer})
    sample["expert_trajectory"] = trajectory
    sample["tool_schemas"] = get_agentic_tools(sample["tools"])
    return sample


def task_top_refund_channel(sample_id, tables, data, rng):
    quarter = rng.choice(["Q1", "Q2", "Q3", "Q4"])
    channel, value = top_group(data["orders"], "channel", "refund_amount", {"quarter": quarter})
    question = f"2024年{quarter}退款金额最高的渠道是什么？请给出退款金额。"
    code = (
        f'orders = pd.read_csv("{tables[0]["path"]}")\n'
        f'df = orders[orders["quarter"] == "{quarter}"]\n'
        'agg = df.groupby("channel")["refund_amount"].sum().sort_values(ascending=False)\n'
        'result = {"channel": agg.index[0], "refund_amount": round(float(agg.iloc[0]), 2)}'
    )
    answer = f"2024年{quarter}退款金额最高的渠道是{channel}，退款金额为{value:.2f}元。"
    return make_sample(
        sample_id,
        question,
        answer,
        {"contains": [channel], "number": value, "tolerance": 0.02},
        tables,
        ["retriever", "python_executor"],
        [
            build_tool_call("retriever", {"query": "退款金额 指标口径 orders refund_amount", "top_k": 2}),
            build_tool_call("python_executor", {"code": code}),
        ],
        "top_refund_channel",
    )


def task_net_revenue(sample_id, tables, data, rng):
    month = rng.randint(1, 12)
    channel = rng.choice(CHANNELS)
    gmv = sum_field(data["orders"], "gmv", month=month, channel=channel)
    refund = sum_field(data["orders"], "refund_amount", month=month, channel=channel)
    value = round(gmv - refund, 2)
    question = f"2024年{month}月{channel}渠道的净收入是多少？"
    code = (
        f'orders = pd.read_csv("{tables[0]["path"]}")\n'
        f'df = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]\n'
        'result = {"net_revenue": round(float(df["gmv"].sum() - df["refund_amount"].sum()), 2)}'
    )
    answer = f"2024年{month}月{channel}渠道的净收入为{value:.2f}元。"
    return make_sample(
        sample_id,
        question,
        answer,
        {"contains": [channel], "number": value, "tolerance": 0.02},
        tables,
        ["retriever", "python_executor"],
        [
            build_tool_call("retriever", {"query": "净收入 GMV 退款金额 计算口径", "top_k": 2}),
            build_tool_call("python_executor", {"code": code}),
        ],
        "net_revenue",
    )


def task_refund_rate_category(sample_id, tables, data, rng):
    month = rng.randint(1, 12)
    category, rate = refund_rate_by_category(data["orders"], month)
    percent = round(rate * 100, 2)
    question = f"2024年{month}月退款率最高的品类是什么？退款率是多少？"
    code = (
        f'orders = pd.read_csv("{tables[0]["path"]}")\n'
        f'df = orders[orders["month"] == {month}]\n'
        'agg = df.groupby("category")[["refund_amount", "gmv"]].sum()\n'
        'agg["refund_rate"] = agg["refund_amount"] / agg["gmv"]\n'
        'top = agg.sort_values("refund_rate", ascending=False).iloc[0]\n'
        'result = {"category": agg.sort_values("refund_rate", ascending=False).index[0], "refund_rate": round(float(top["refund_rate"] * 100), 2)}'
    )
    answer = f"2024年{month}月退款率最高的品类是{category}，退款率为{percent:.2f}%。"
    return make_sample(
        sample_id,
        question,
        answer,
        {"contains": [category], "number": percent, "tolerance": 0.02},
        tables,
        ["retriever", "python_executor"],
        [
            build_tool_call("retriever", {"query": "退款率 指标口径 refund_amount gmv", "top_k": 2}),
            build_tool_call("python_executor", {"code": code}),
        ],
        "refund_rate_category",
    )


def task_roas(sample_id, tables, data, rng):
    month = rng.randint(1, 12)
    channel = rng.choice([c for c in CHANNELS if c != "自然流量"])
    gmv = monthly_gmv_by_channel(data["orders"], month, channel)
    spend = marketing_spend(data["marketing"], month, channel)
    value = round(gmv / spend, 2)
    question = f"2024年{month}月{channel}渠道的ROAS是多少？"
    code = (
        f'orders = pd.read_csv("{tables[0]["path"]}")\n'
        f'marketing = pd.read_csv("{tables[1]["path"]}")\n'
        f'gmv = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]["gmv"].sum()\n'
        f'spend = marketing[(marketing["month"] == {month}) & (marketing["channel"] == "{channel}")]["spend"].sum()\n'
        'result = {"roas": round(float(gmv / spend), 2)}'
    )
    answer = f"2024年{month}月{channel}渠道的ROAS为{value:.2f}。"
    return make_sample(
        sample_id,
        question,
        answer,
        {"contains": [channel], "number": value, "tolerance": 0.02},
        tables,
        ["retriever", "python_executor"],
        [
            build_tool_call("retriever", {"query": "ROAS GMV 广告花费 指标口径", "top_k": 2}),
            build_tool_call("python_executor", {"code": code}),
        ],
        "roas",
    )


def task_new_users(sample_id, tables, data, rng):
    month = rng.randint(1, 12)
    channel = rng.choice(CHANNELS)
    province, value = top_group(data["users"], "province", "new_users", {"month": month, "channel": channel})
    question = f"2024年{month}月{channel}渠道新增用户最多的省份是哪个？新增用户数是多少？"
    code = (
        f'users = pd.read_csv("{tables[2]["path"]}")\n'
        f'df = users[(users["month"] == {month}) & (users["channel"] == "{channel}")]\n'
        'agg = df.groupby("province")["new_users"].sum().sort_values(ascending=False)\n'
        'result = {"province": agg.index[0], "new_users": int(agg.iloc[0])}'
    )
    answer = f"2024年{month}月{channel}渠道新增用户最多的省份是{province}，新增用户数为{int(value)}。"
    return make_sample(
        sample_id,
        question,
        answer,
        {"contains": [province], "number": int(value), "tolerance": 0.01},
        tables,
        ["retriever", "python_executor"],
        [
            build_tool_call("retriever", {"query": "新增用户 users new_users 字段口径", "top_k": 2}),
            build_tool_call("python_executor", {"code": code}),
        ],
        "new_users",
    )


def task_incremental_gmv(sample_id, tables, data, rng):
    month = rng.randint(1, 12)
    channel = rng.choice([c for c in CHANNELS if c != "自然流量"])
    pct = rng.choice([10, 12, 15, 20])
    gmv = monthly_gmv_by_channel(data["orders"], month, channel)
    spend = marketing_spend(data["marketing"], month, channel)
    roas = gmv / spend
    incremental = round(spend * pct / 100 * roas, 2)
    question = f"如果2024年{month}月{channel}渠道广告花费增加{pct}%，且ROAS保持不变，预计GMV增加多少？"
    code = (
        f'orders = pd.read_csv("{tables[0]["path"]}")\n'
        f'marketing = pd.read_csv("{tables[1]["path"]}")\n'
        f'gmv = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]["gmv"].sum()\n'
        f'spend = marketing[(marketing["month"] == {month}) & (marketing["channel"] == "{channel}")]["spend"].sum()\n'
        'result = {"spend": round(float(spend), 2), "roas": round(float(gmv / spend), 4)}'
    )
    expression = f"{spend:.2f} * {pct} / 100 * {roas:.6f}"
    answer = f"预计GMV增加{incremental:.2f}元。"
    return make_sample(
        sample_id,
        question,
        answer,
        {"contains": ["GMV"], "number": incremental, "tolerance": 0.05},
        tables,
        ["retriever", "python_executor", "calculator"],
        [
            build_tool_call("retriever", {"query": "ROAS GMV 广告花费 预算增加 计算口径", "top_k": 2}),
            build_tool_call("python_executor", {"code": code}),
            build_tool_call("calculator", {"expression": expression}),
        ],
        "incremental_gmv",
        max_turns=5,
    )


TASK_BUILDERS = [
    task_top_refund_channel,
    task_net_revenue,
    task_refund_rate_category,
    task_roas,
    task_new_users,
    task_incremental_gmv,
]


def write_jsonl(path, rows):
    with open(path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def build_split(name, size, tables, data, seed):
    rng = random.Random(seed)
    rows = []
    for i in range(size):
        task_fn = TASK_BUILDERS[i % len(TASK_BUILDERS)]
        sample_id = f"agentic_{name}_{i:06d}"
        rows.append(task_fn(sample_id, tables, data, rng))
    rng.shuffle(rows)
    return rows


def parse_args():
    parser = argparse.ArgumentParser(description="Build Agentic DataAnalysis tasks")
    parser.add_argument("--output_dir", default="dataset/agentic_data")
    parser.add_argument("--train_size", type=int, default=30000)
    parser.add_argument("--dev_size", type=int, default=1000)
    parser.add_argument("--test_size", type=int, default=1000)
    parser.add_argument("--seed", type=int, default=42)
    return parser.parse_args()


def main():
    args = parse_args()
    mkdir(args.output_dir)
    data = generate_tables(args.output_dir, args.seed)
    tables = table_specs(data["paths"])
    splits = {
        "train": build_split("train", args.train_size, tables, data, args.seed + 1),
        "dev": build_split("dev", args.dev_size, tables, data, args.seed + 2),
        "test": build_split("test", args.test_size, tables, data, args.seed + 3),
    }
    for split, rows in splits.items():
        write_jsonl(os.path.join(args.output_dir, f"{split}.jsonl"), rows)
    manifest = {
        "task": "AgenticDataAnalysis",
        "seed": args.seed,
        "splits": {split: len(rows) for split, rows in splits.items()},
        "tables": tables,
        "tools": [tool["function"]["name"] for tool in get_agentic_tools()],
        "task_types": [fn.__name__.replace("task_", "") for fn in TASK_BUILDERS],
    }
    with open(os.path.join(args.output_dir, "manifest.json"), "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)
    print(json.dumps(manifest, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()