import ast
import contextlib
import io
import json
import math
import multiprocessing as mp
import os
import re
import statistics
from collections import Counter
from types import SimpleNamespace

import numpy as np
import pandas as pd


AGENTIC_SYSTEM_PROMPT = (
    "你是一个面向运营数据分析的 Agent。你需要根据用户问题选择合适工具，"
    "优先检索指标口径，再用 Python 对给定 CSV 数据进行计算；最终答案必须简洁，"
    "包含关键结论和必要数值。工具调用必须使用 <tool_call> 包裹的 JSON。"
)


AGENTIC_TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "retriever",
            "description": "检索指标口径、表字段说明或业务规则，返回最相关的知识片段。",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "检索查询"},
                    "top_k": {"type": "integer", "description": "返回片段数量", "default": 3},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "python_executor",
            "description": "在受限 Python 环境中使用 pandas/numpy 读取给定 CSV 并完成数据分析。",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "Python 代码。把最终结果赋值给 result 或 answer。",
                    }
                },
                "required": ["code"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "执行确定性数学表达式计算，适合简单四则运算、比例和幂运算。",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "数学表达式，如 1280 * 0.15 * 3.2"}
                },
                "required": ["expression"],
            },
        },
    },
]


def get_agentic_tools(names=None):
    if names is None:
        names = ["retriever", "python_executor", "calculator"]
    name_set = set(names)
    return [tool for tool in AGENTIC_TOOL_SCHEMAS if tool["function"]["name"] in name_set]


def json_dumps(data):
    return json.dumps(data, ensure_ascii=False, default=str)


def format_agentic_user_prompt(sample):
    table_lines = []
    for table in sample.get("tables", []):
        table_lines.append(
            f"- {table.get('name', '')}: path={table.get('path', '')}; {table.get('description', '')}"
        )
    tool_names = ", ".join(sample.get("tools", ["retriever", "python_executor", "calculator"]))
    return (
        f"用户问题：{sample.get('question', '')}\n\n"
        f"可用数据表：\n" + "\n".join(table_lines) + "\n\n"
        f"可用工具：{tool_names}\n"
        "要求：先确认指标口径，再调用工具计算；最终答案只输出结论，不要编造数据。"
    )


def parse_tool_calls(text):
    calls = []
    for match in re.findall(r"<tool_call>(.*?)</tool_call>", text or "", re.DOTALL):
        raw = match.strip()
        try:
            data = json.loads(raw)
            args = data.get("arguments", {})
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except Exception:
                    args = {}
            calls.append(
                {
                    "name": data.get("name", ""),
                    "arguments": args if isinstance(args, dict) else {},
                    "raw": raw,
                    "valid_json": True,
                }
            )
        except Exception:
            calls.append({"name": "", "arguments": {}, "raw": raw, "valid_json": False})
    return calls


def strip_thinking(text):
    text = text or ""
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()


def extract_final_answer(turn_outputs):
    if isinstance(turn_outputs, str):
        turn_outputs = [turn_outputs]
    if not turn_outputs:
        return ""
    text = strip_thinking(turn_outputs[-1])
    if "</tool_call>" in text:
        text = text.split("</tool_call>")[-1]
    text = re.sub(r"<tool_call>.*?</tool_call>", "", text, flags=re.DOTALL)
    return text.strip()


def _tokenize(text):
    return re.findall(r"[\w\u4e00-\u9fff]+", str(text).lower())


def _overlap_score(query, text):
    q = Counter(_tokenize(query))
    d = Counter(_tokenize(text))
    return sum(min(q[token], d[token]) for token in q)


def _safe_number(value, default=0.0):
    try:
        return float(value)
    except Exception:
        return default


class _SafeArithmetic(ast.NodeVisitor):
    allowed_binops = {
        ast.Add: lambda a, b: a + b,
        ast.Sub: lambda a, b: a - b,
        ast.Mult: lambda a, b: a * b,
        ast.Div: lambda a, b: a / b,
        ast.FloorDiv: lambda a, b: a // b,
        ast.Mod: lambda a, b: a % b,
        ast.Pow: lambda a, b: a ** b,
    }
    allowed_unary = {
        ast.UAdd: lambda a: +a,
        ast.USub: lambda a: -a,
    }
    allowed_funcs = {
        "abs": abs,
        "round": round,
        "sqrt": math.sqrt,
        "log": math.log,
        "log10": math.log10,
        "ceil": math.ceil,
        "floor": math.floor,
        "pow": pow,
    }

    def visit_Expression(self, node):
        return self.visit(node.body)

    def visit_Constant(self, node):
        if isinstance(node.value, (int, float)):
            return float(node.value)
        raise ValueError("calculator only supports numeric constants")

    def visit_BinOp(self, node):
        op = type(node.op)
        if op not in self.allowed_binops:
            raise ValueError(f"operator not allowed: {op.__name__}")
        return self.allowed_binops[op](self.visit(node.left), self.visit(node.right))

    def visit_UnaryOp(self, node):
        op = type(node.op)
        if op not in self.allowed_unary:
            raise ValueError(f"operator not allowed: {op.__name__}")
        return self.allowed_unary[op](self.visit(node.operand))

    def visit_Call(self, node):
        if not isinstance(node.func, ast.Name) or node.func.id not in self.allowed_funcs:
            raise ValueError("function not allowed")
        args = [self.visit(arg) for arg in node.args]
        return self.allowed_funcs[node.func.id](*args)

    def generic_visit(self, node):
        raise ValueError(f"node not allowed: {type(node).__name__}")


def safe_calculate(expression):
    normalized = (
        str(expression)
        .replace("^", "**")
        .replace("×", "*")
        .replace("÷", "/")
        .replace("−", "-")
        .replace("（", "(")
        .replace("）", ")")
    )
    tree = ast.parse(normalized, mode="eval")
    return _SafeArithmetic().visit(tree)


def _jsonable(value):
    if isinstance(value, pd.DataFrame):
        return value.to_dict("records")
    if isinstance(value, pd.Series):
        return value.to_dict()
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, dict):
        return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [_jsonable(v) for v in value]
    return value


def _validate_python_ast(code):
    banned_nodes = (ast.Import, ast.ImportFrom, ast.Global, ast.Nonlocal, ast.Lambda, ast.ClassDef)
    banned_names = {
        "__import__",
        "eval",
        "exec",
        "compile",
        "open",
        "input",
        "globals",
        "locals",
        "vars",
        "dir",
        "getattr",
        "setattr",
        "delattr",
        "os",
        "sys",
        "subprocess",
        "socket",
        "requests",
        "urllib",
        "pathlib",
        "shutil",
    }
    banned_attrs = {
        "to_csv",
        "to_excel",
        "to_pickle",
        "to_parquet",
        "write",
        "remove",
        "unlink",
        "rmdir",
        "mkdir",
        "rename",
        "replace",
        "system",
        "popen",
    }
    tree = ast.parse(code, mode="exec")
    for node in ast.walk(tree):
        if isinstance(node, banned_nodes):
            raise ValueError(f"python_executor does not allow {type(node).__name__}")
        if isinstance(node, ast.Name) and node.id in banned_names:
            raise ValueError(f"name not allowed: {node.id}")
        if isinstance(node, ast.Attribute) and node.attr in banned_attrs:
            raise ValueError(f"attribute not allowed: {node.attr}")
    return tree


class _PandasProxy:
    def __init__(self, repo_root, allowed_paths):
        self.repo_root = repo_root
        self.allowed_paths = {os.path.abspath(path) for path in allowed_paths}

    def read_csv(self, path, *args, **kwargs):
        candidates = []
        raw = str(path)
        if os.path.isabs(raw):
            candidates.append(os.path.abspath(raw))
        candidates.append(os.path.abspath(raw))
        candidates.append(os.path.abspath(os.path.join(self.repo_root, raw)))
        allowed_by_basename = {os.path.basename(path): path for path in self.allowed_paths}
        if raw in allowed_by_basename:
            candidates.append(allowed_by_basename[raw])
        for candidate in candidates:
            if candidate in self.allowed_paths:
                return pd.read_csv(candidate, *args, **kwargs)
        raise PermissionError(f"read_csv path is not in sample tables: {raw}")

    def __getattr__(self, name):
        return getattr(pd, name)


def _python_worker(code, repo_root, allowed_paths, queue):
    try:
        tree = _validate_python_ast(code)
        output = io.StringIO()
        safe_builtins = {
            "abs": abs,
            "all": all,
            "any": any,
            "bool": bool,
            "dict": dict,
            "enumerate": enumerate,
            "float": float,
            "int": int,
            "len": len,
            "list": list,
            "max": max,
            "min": min,
            "pow": pow,
            "print": print,
            "range": range,
            "round": round,
            "set": set,
            "sorted": sorted,
            "str": str,
            "sum": sum,
            "tuple": tuple,
            "zip": zip,
        }
        local_vars = {}
        global_vars = {
            "__builtins__": safe_builtins,
            "pd": _PandasProxy(repo_root, allowed_paths),
            "np": np,
            "math": math,
            "json": json,
            "statistics": statistics,
        }
        with contextlib.redirect_stdout(output):
            exec(compile(tree, "<agentic_python_executor>", "exec"), global_vars, local_vars)
        value = local_vars.get("result", local_vars.get("answer", output.getvalue().strip()))
        queue.put({"ok": True, "result": _jsonable(value), "error": "", "artifacts": []})
    except Exception as exc:
        queue.put({"ok": False, "result": "", "error": str(exc)[:500], "artifacts": []})


class AgenticToolEnv:
    def __init__(self, sample, repo_root=None, timeout=3):
        self.sample = sample
        self.repo_root = os.path.abspath(repo_root or os.getcwd())
        self.timeout = timeout
        self.allowed_paths = [self._resolve_table_path(table.get("path", "")) for table in sample.get("tables", [])]

    def _resolve_table_path(self, path):
        if os.path.isabs(path):
            return os.path.abspath(path)
        return os.path.abspath(os.path.join(self.repo_root, path))

    def execute(self, name, arguments=None):
        arguments = arguments or {}
        if name == "retriever":
            return self._retriever(arguments)
        if name == "calculator":
            return self._calculator(arguments)
        if name == "python_executor":
            return self._python_executor(arguments)
        return {"ok": False, "result": "", "error": f"unknown tool: {name}", "artifacts": []}

    def _retriever(self, arguments):
        query = str(arguments.get("query", self.sample.get("question", "")))
        top_k = max(1, min(int(arguments.get("top_k", 3) or 3), 8))
        docs = self.sample.get("documents", [])
        scored = []
        for doc in docs:
            text = f"{doc.get('title', '')} {doc.get('text', '')}"
            scored.append((_overlap_score(query, text), doc))
        scored.sort(key=lambda item: item[0], reverse=True)
        result = []
        for score, doc in scored[:top_k]:
            result.append(
                {
                    "id": doc.get("id", ""),
                    "title": doc.get("title", ""),
                    "text": doc.get("text", ""),
                    "score": score,
                }
            )
        return {"ok": True, "result": result, "error": "", "artifacts": []}

    def _calculator(self, arguments):
        expression = arguments.get("expression", "")
        try:
            value = safe_calculate(expression)
        except Exception as exc:
            return {"ok": False, "result": "", "error": str(exc)[:500], "artifacts": []}
        return {"ok": True, "result": value, "error": "", "artifacts": []}

    def _python_executor(self, arguments):
        code = str(arguments.get("code", ""))
        if not code.strip():
            return {"ok": False, "result": "", "error": "empty code", "artifacts": []}
        queue = mp.Queue(maxsize=1)
        proc = mp.Process(target=_python_worker, args=(code, self.repo_root, self.allowed_paths, queue))
        proc.start()
        proc.join(self.timeout)
        if proc.is_alive():
            proc.terminate()
            proc.join()
            return {"ok": False, "result": "", "error": f"python_executor timeout>{self.timeout}s", "artifacts": []}
        if queue.empty():
            return {"ok": False, "result": "", "error": "python_executor returned no result", "artifacts": []}
        return queue.get()


def validate_tool_call(call, available_names):
    if not call.get("valid_json", False):
        return False
    name = call.get("name", "")
    args = call.get("arguments", {})
    if name not in available_names:
        return False
    if name == "retriever":
        return bool(args.get("query"))
    if name == "calculator":
        return bool(args.get("expression"))
    if name == "python_executor":
        return bool(args.get("code"))
    return False


def _f1(pred_items, gold_items):
    pred = set(pred_items)
    gold = set(gold_items)
    if not pred and not gold:
        return 1.0
    if not pred or not gold:
        return 0.0
    inter = len(pred & gold)
    if inter == 0:
        return 0.0
    precision = inter / len(pred)
    recall = inter / len(gold)
    return 2 * precision * recall / (precision + recall)


def _numbers(text):
    return [float(x.replace(",", "")) for x in re.findall(r"[-+]?\d+(?:,\d{3})*(?:\.\d+)?", str(text))]


def score_answer(final_answer, sample):
    checks = sample.get("checks", {}) or {}
    score_parts = []
    contains = checks.get("contains") or []
    if contains:
        score_parts.append(float(all(str(item).lower() in final_answer.lower() for item in contains)))
    if "number" in checks:
        target = _safe_number(checks.get("number"))
        tolerance = _safe_number(checks.get("tolerance"), 1e-3)
        nums = _numbers(final_answer)
        score_parts.append(float(any(abs(num - target) <= tolerance for num in nums)))
    if not score_parts and sample.get("answer"):
        score_parts.append(float(str(sample["answer"]).lower() in final_answer.lower()))
    return sum(score_parts) / len(score_parts) if score_parts else 0.0


def repetition_penalty(text, n=3, cap=0.5):
    toks = _tokenize(text)
    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0


def score_agentic_trajectory(turn_outputs, sample, repo_root=None, unfinished=False, execute_tools=True):
    if isinstance(turn_outputs, str):
        turn_outputs = [turn_outputs]
    turn_outputs = turn_outputs or []
    available_names = set(sample.get("tools", ["retriever", "python_executor", "calculator"]))
    expected_tools = sample.get("expected_tools", [])
    env = AgenticToolEnv(sample, repo_root=repo_root)

    calls = []
    tag_errors = 0
    for turn in turn_outputs:
        tag_errors += abs(str(turn).count("<tool_call>") - str(turn).count("</tool_call>"))
        calls.extend(parse_tool_calls(turn))

    valid_flags = [validate_tool_call(call, available_names) for call in calls]
    valid_count = sum(valid_flags)
    invalid_count = len(calls) - valid_count
    selected_tools = [call.get("name", "") for call, ok in zip(calls, valid_flags) if ok]
    tool_f1 = _f1(selected_tools, expected_tools)

    exec_success = 1.0
    if execute_tools and calls:
        successes = []
        for call, ok in zip(calls, valid_flags):
            if not ok:
                successes.append(False)
                continue
            result = env.execute(call["name"], call["arguments"])
            successes.append(bool(result.get("ok")))
        exec_success = sum(successes) / len(successes) if successes else 0.0

    final_answer = extract_final_answer(turn_outputs)
    answer_score = 0.0 if unfinished else score_answer(final_answer, sample)
    task_success = float(answer_score >= 0.999)
    schema_valid_rate = valid_count / len(calls) if calls else 0.0
    invalid_action_rate = invalid_count / len(calls) if calls else 0.0
    format_score = 1.0 if tag_errors == 0 and all(call.get("valid_json") for call in calls) else 0.0
    turn_count = len(turn_outputs)
    expected_turns = max(1, int(sample.get("max_turns", 4)))
    turn_efficiency = max(0.0, 1.0 - max(0, turn_count - expected_turns) / expected_turns)
    repeat_penalty = repetition_penalty(final_answer)

    reward = (
        2.5 * task_success
        + 1.2 * answer_score
        + 0.8 * tool_f1
        + 0.6 * schema_valid_rate
        + 0.6 * exec_success
        + 0.4 * format_score
        + 0.2 * turn_efficiency
        - 0.8 * invalid_action_rate
        - 0.3 * tag_errors
        - repeat_penalty
        - (0.5 if unfinished else 0.0)
    )
    reward = max(min(reward, 5.0), -3.0)
    return reward, {
        "reward_total": reward,
        "task_success": task_success,
        "answer_score": answer_score,
        "tool_selection_f1": tool_f1,
        "schema_valid_rate": schema_valid_rate,
        "exec_success_rate": exec_success,
        "invalid_action_rate": invalid_action_rate,
        "format_score": format_score,
        "turn_efficiency": turn_efficiency,
        "repetition_penalty": repeat_penalty,
        "num_tool_calls": len(calls),
        "avg_turns": turn_count,
    }


def average_agentic_metrics(metrics):
    if not metrics:
        return {}
    keys = sorted({key for row in metrics for key in row})
    return {
        key: sum(float(row.get(key, 0.0)) for row in metrics) / len(metrics)
        for key in keys
    }