#!/usr/bin/env python3
from __future__ import annotations

"""批量调用线上模型 generate 接口给数据打标。

输入通过 stdin 或 --input 传 JSON 对象，输出稳定 JSON 对象到 stdout。
脚本只依赖 Python 标准库，便于在本地、Linux runtime 用户或远程工作区迁移执行。
"""

import argparse
import csv
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from io import StringIO
from pathlib import Path
from typing import Any


DEFAULT_SYSTEM_PROMPT = "你是小爱同学，中文智能语音助手。"
DEFAULT_PARAMETERS = {"max_new_tokens": 64}


class LabelingError(ValueError):
    """输入参数或数据格式错误。"""


def main() -> int:
    try:
        payload = load_payload()
        result = run(payload)
        emit({"ok": True, **result})
        return 0
    except LabelingError as exc:
        emit({"ok": False, "error": str(exc), **getattr(exc, "extra", {})})
        return 1
    except Exception as exc:  # noqa: BLE001 - CLI 需要稳定 JSON 错误
        emit({"ok": False, "error": str(exc)})
        return 1


def load_payload() -> dict[str, Any]:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", "-i", help="JSON 参数文件；不传则从 stdin 读取")
    args = parser.parse_args()
    text = Path(args.input).read_text(encoding="utf-8") if args.input else sys.stdin.read()
    if not text.strip():
        raise LabelingError("input JSON is required")
    payload = json.loads(text)
    if not isinstance(payload, dict):
        raise LabelingError("input JSON must be an object")
    return payload


def emit(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, ensure_ascii=False, separators=(",", ":")))


def run(payload: dict[str, Any]) -> dict[str, Any]:
    input_path = require_string(payload, "input_path")
    model_url = require_string(payload, "model_url")
    dry_run = bool(payload.get("dry_run", False))
    include_prompt = bool(payload.get("include_prompt", False))
    max_records = optional_int(payload.get("max_records"))
    start_index = int(payload.get("start_index") or 0)
    timeout_seconds = float(payload.get("timeout_seconds") or 60)
    output_path = str(payload.get("output_path") or "output/model_predictions.jsonl")
    parameters = payload.get("parameters")
    if parameters is None:
        parameters = dict(DEFAULT_PARAMETERS)
    if not isinstance(parameters, dict):
        raise LabelingError("parameters must be an object")
    headers = payload.get("headers")
    if headers is None:
        headers = {"Content-Type": "application/json"}
    if not isinstance(headers, dict):
        raise LabelingError("headers must be an object")
    request_template = payload.get("request_template")
    if request_template is not None and not isinstance(request_template, dict):
        raise LabelingError("request_template must be an object")

    source = read_input_file(input_path)
    samples = normalize_samples(
        source,
        field_mapping=payload.get("field_mapping"),
        system_prompt=str(payload.get("system_prompt") or DEFAULT_SYSTEM_PROMPT),
        session_num=int(payload.get("session_num") or 5),
        session_time_minutes=int(payload.get("session_time_minutes") or 5),
    )
    if start_index:
        samples = samples[start_index:]
    if max_records is not None:
        samples = samples[:max_records]

    summary = {
        "input_path": input_path,
        "source_format": source["format"],
        "total_samples": len(samples),
        "sample_preview": preview_samples(samples),
    }
    if dry_run:
        return {**summary, "dry_run": True}

    if not model_url.startswith(("http://", "https://")):
        raise LabelingError("model_url must start with http:// or https://")
    results: list[dict[str, Any]] = []
    ok_count = 0
    output = resolve_runtime_path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8") as file:
        for sample in samples:
            item = request_one(
                sample,
                model_url=model_url,
                parameters=parameters,
                headers={str(k): str(v) for k, v in headers.items()},
                timeout_seconds=timeout_seconds,
                request_template=request_template,
                include_prompt=include_prompt,
            )
            if item.get("ok"):
                ok_count += 1
            results.append(item)
            file.write(json.dumps(item, ensure_ascii=False, separators=(",", ":")) + "\n")

    return {
        **summary,
        "dry_run": False,
        "output_path": str(output),
        "success_count": ok_count,
        "failure_count": len(results) - ok_count,
        "result_preview": results[:3],
    }


def request_one(
    sample: dict[str, Any],
    *,
    model_url: str,
    parameters: dict[str, Any],
    headers: dict[str, str],
    timeout_seconds: float,
    request_template: dict[str, Any] | None,
    include_prompt: bool,
) -> dict[str, Any]:
    prompt = str(sample["prompt"])
    body = build_request_body(prompt, parameters, request_template)
    started = time.time()
    base = {
        "index": sample["index"],
        "request_id": sample.get("request_id", ""),
        "query": sample.get("query", ""),
        "gold_label": sample.get("gold_label", ""),
        "source_format": sample.get("source_format", ""),
    }
    if include_prompt:
        base["prompt"] = prompt
    try:
        request = urllib.request.Request(
            model_url,
            data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
            headers=headers,
            method="POST",
        )
        with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
            text = response.read().decode("utf-8", errors="replace")
            status = getattr(response, "status", 200)
        latency_ms = int((time.time() - started) * 1000)
        decoded = try_json(text)
        return {
            **base,
            "prediction": extract_prediction(decoded, text),
            "ok": 200 <= int(status) < 300,
            "status": int(status),
            "latency_ms": latency_ms,
            "response": decoded if decoded is not None else text,
        }
    except urllib.error.HTTPError as exc:
        text = exc.read().decode("utf-8", errors="replace")
        return {
            **base,
            "prediction": extract_prediction(try_json(text), text),
            "ok": False,
            "status": exc.code,
            "latency_ms": int((time.time() - started) * 1000),
            "error": f"HTTP {exc.code}: {text[:500]}",
        }
    except Exception as exc:  # noqa: BLE001 - 单条失败不中断整体批次
        return {
            **base,
            "prediction": "",
            "ok": False,
            "latency_ms": int((time.time() - started) * 1000),
            "error": str(exc),
        }


def build_request_body(
    prompt: str,
    parameters: dict[str, Any],
    request_template: dict[str, Any] | None,
) -> dict[str, Any]:
    if request_template is None:
        return {"inputs": prompt, "parameters": parameters}
    return replace_placeholders(request_template, {"prompt": prompt, "parameters": parameters})


def replace_placeholders(value: Any, variables: dict[str, Any]) -> Any:
    if isinstance(value, str):
        if value == "{{prompt}}":
            return variables["prompt"]
        if value == "{{parameters}}":
            return variables["parameters"]
        return value.replace("{{prompt}}", str(variables["prompt"]))
    if isinstance(value, list):
        return [replace_placeholders(item, variables) for item in value]
    if isinstance(value, dict):
        return {key: replace_placeholders(item, variables) for key, item in value.items()}
    return value


def read_input_file(path: str) -> dict[str, Any]:
    file_path = resolve_runtime_path(path)
    if not file_path.exists():
        raise LabelingError(f"input_path not found: {path}")
    suffix = file_path.suffix.lower()
    text = file_path.read_text(encoding="utf-8", errors="replace")
    if suffix in {".csv", ".tsv"}:
        delimiter = "\t" if suffix == ".tsv" else ","
        rows = list(csv.DictReader(StringIO(text), delimiter=delimiter))
        return {"format": "table", "path": str(file_path), "rows": rows, "columns": list(rows[0].keys()) if rows else []}
    if suffix == ".jsonl":
        rows = [json.loads(line) for line in text.splitlines() if line.strip()]
        return {"format": "jsonl", "path": str(file_path), "rows": rows}
    decoded = json.loads(text)
    if isinstance(decoded, dict) and isinstance(decoded.get("records"), list):
        decoded = decoded["records"]
    if isinstance(decoded, list):
        return {"format": "json", "path": str(file_path), "rows": decoded}
    raise_with_mapping("JSON input must be an array or {records:[...]}", columns=list(decoded.keys()) if isinstance(decoded, dict) else [])


def normalize_samples(
    source: dict[str, Any],
    *,
    field_mapping: Any,
    system_prompt: str,
    session_num: int,
    session_time_minutes: int,
) -> list[dict[str, Any]]:
    rows = source.get("rows")
    if not isinstance(rows, list) or not rows:
        raise LabelingError("input file contains no rows")
    if isinstance(field_mapping, dict):
        return samples_from_mapping(rows, field_mapping, source["format"], system_prompt)
    first = rows[0]
    if not isinstance(first, dict):
        raise_with_mapping("rows must be objects")
    if is_canonical_record(first):
        return [
            sample_from_record(index, row, system_prompt, session_num, session_time_minutes)
            for index, row in enumerate(rows)
            if isinstance(row, dict)
        ]
    if is_training_row(first):
        return [sample_from_training(index, row, system_prompt) for index, row in enumerate(rows) if isinstance(row, dict)]
    if source["format"] == "table":
        columns = list(first.keys())
        if "newPrompt" in columns:
            return [sample_from_eval_row(index, row) for index, row in enumerate(rows) if isinstance(row, dict)]
        if "query" in columns and "function" in columns:
            return [sample_from_flow_row(index, row, system_prompt) for index, row in enumerate(rows) if isinstance(row, dict)]
        raise_with_mapping("unrecognized table format", columns=columns)
    raise_with_mapping("unrecognized JSON/JSONL format", columns=list(first.keys()))


def sample_from_record(
    index: int,
    record: dict[str, Any],
    system_prompt: str,
    session_num: int,
    session_time_minutes: int,
) -> dict[str, Any]:
    source = record.get("source") if isinstance(record.get("source"), dict) else {}
    turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
    return {
        "index": index,
        "request_id": str(source.get("request_id") or ""),
        "query": str(turn.get("query") or ""),
        "prompt": build_planning_prompt(record, system_prompt, session_num, session_time_minutes),
        "gold_label": combined_label(record),
        "source_format": "canonical_record_v1",
    }


def sample_from_training(index: int, row: dict[str, Any], system_prompt: str) -> dict[str, Any]:
    prompt = str(row.get("prompt") or "")
    if not prompt:
        system = str(row.get("system") or system_prompt)
        instruction = str(row.get("instruction") or row.get("input") or "")
        prompt = wrap_chat_prompt(system, instruction)
    return {
        "index": index,
        "request_id": str(row.get("request_id") or ""),
        "query": extract_query_from_prompt(prompt),
        "prompt": prompt,
        "gold_label": str(row.get("output") or row.get("target") or ""),
        "source_format": "training_jsonl",
    }


def sample_from_eval_row(index: int, row: dict[str, Any]) -> dict[str, Any]:
    gold = str(row.get("code标签") or row.get("function") or row.get("target") or "")
    complex_value = row.get("complex")
    if complex_value not in (None, "") and not gold.startswith("complex="):
        gold = f"complex={normalize_bool_literal(complex_value)}\n{gold}".rstrip()
    prompt = str(row.get("newPrompt") or row.get("prompt") or "")
    query = str(row.get("query") or "")
    if not prompt:
        prompt = build_planning_prompt(minimal_record(query, {}, [], gold), DEFAULT_SYSTEM_PROMPT, 5, 5)
    return {
        "index": index,
        "request_id": str(row.get("request_id") or ""),
        "query": query,
        "prompt": prompt,
        "gold_label": gold,
        "source_format": "eval_csv",
    }


def sample_from_flow_row(index: int, row: dict[str, Any], system_prompt: str) -> dict[str, Any]:
    record = record_from_flow_row(row)
    return {
        "index": index,
        "request_id": str(row.get("request_id") or ""),
        "query": str(row.get("query") or ""),
        "prompt": build_planning_prompt(record, system_prompt, 5, 5),
        "gold_label": str(row.get("function") or ""),
        "source_format": "flow_csv",
    }


def samples_from_mapping(
    rows: list[Any],
    mapping: dict[str, Any],
    source_format: str,
    system_prompt: str,
) -> list[dict[str, Any]]:
    query_field = str(mapping.get("query") or "")
    prompt_field = str(mapping.get("prompt") or "")
    label_field = str(mapping.get("label") or mapping.get("target") or "")
    request_id_field = str(mapping.get("request_id") or "")
    if not query_field and not prompt_field:
        raise LabelingError("field_mapping must provide query or prompt")
    samples: list[dict[str, Any]] = []
    for index, row in enumerate(rows):
        if not isinstance(row, dict):
            continue
        prompt = str(row.get(prompt_field) or "")
        query = str(row.get(query_field) or "")
        if not prompt:
            record = minimal_record(query, {}, [], str(row.get(label_field) or ""))
            prompt = build_planning_prompt(record, system_prompt, 5, 5)
        samples.append(
            {
                "index": index,
                "request_id": str(row.get(request_id_field) or ""),
                "query": query or extract_query_from_prompt(prompt),
                "prompt": prompt,
                "gold_label": str(row.get(label_field) or ""),
                "source_format": f"{source_format}_mapped",
            }
        )
    return samples


def is_canonical_record(row: dict[str, Any]) -> bool:
    return isinstance(row.get("turn"), dict) and bool(row["turn"].get("query"))


def is_training_row(row: dict[str, Any]) -> bool:
    return any(key in row for key in ("instruction", "system", "output", "prompt"))


def build_planning_prompt(
    record: dict[str, Any],
    system_prompt: str,
    session_num: int,
    session_time_minutes: int,
) -> str:
    instruction = build_training_instruction(record, session_num, session_time_minutes)
    return wrap_chat_prompt(system_prompt, instruction)


def wrap_chat_prompt(system_prompt: str, instruction: str) -> str:
    system_prompt = prompt_section_without_trailing_newline(system_prompt)
    instruction = prompt_section_with_one_trailing_newline(instruction)
    return (
        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
        f"<|im_start|>user\n{instruction}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )


def prompt_section_without_trailing_newline(text: str) -> str:
    return str(text or "").rstrip("\n")


def prompt_section_with_one_trailing_newline(text: str) -> str:
    return prompt_section_without_trailing_newline(text) + "\n"


def build_training_instruction(record: dict[str, Any], session_num: int, session_time_minutes: int) -> str:
    turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
    query = str(turn.get("query") or "")
    context = record.get("context") if isinstance(record.get("context"), dict) else {}
    prev_session = record.get("prev_session") if isinstance(record.get("prev_session"), list) else []
    current_ts = optional_int(turn.get("timestamp"))
    history = render_history(prev_session, current_ts, session_num, session_time_minutes)
    return (
        "请参考用户的[当前query]、[对话历史]、[知识注入]、[系统状态]识别出[当前query]的[function]结果，[function]是python的code形式。\n"
        "[知识注入]\n"
        f"{context_prompt_block(context)}\n"
        "[系统状态]\n"
        "{}\n"
        "[对话历史]\n"
        f"{history}"
        "[当前query]\n"
        f"用户: {query}\n"
        "[function]\n"
    )


def context_prompt_block(context: dict[str, Any]) -> str:
    fields = ("location", "rag")
    lines = ["{"]
    for index, field in enumerate(fields):
        comma = "," if index < len(fields) - 1 else ""
        value = str(context.get(field) or "")
        lines.append(f'"{field}": {json.dumps(value, ensure_ascii=False)}{comma}')
    lines.append("}")
    return "\n".join(lines)


def render_history(prev_session: list[Any], current_ts: int | None, session_num: int, session_time_minutes: int) -> str:
    usable: list[dict[str, Any]] = []
    for item in prev_session[-session_num:]:
        if not isinstance(item, dict):
            continue
        ts = optional_int(item.get("timestamp"))
        if current_ts is not None and ts is not None:
            if abs(current_ts - ts) > session_time_minutes * 60_000:
                continue
        usable.append(item)
    if not usable:
        return ""
    lines: list[str] = []
    for item in usable:
        query = str(item.get("query") or "").strip()
        tts = str(item.get("tts") or "").strip()
        if query:
            lines.append(f"用户: {query}")
        if tts:
            lines.append(f"小爱: {tts}")
    return "\n".join(lines) + ("\n" if lines else "")


def record_from_flow_row(row: dict[str, Any]) -> dict[str, Any]:
    prev_session = parse_json_cell(row.get("prev_session"), default=[])
    context = parse_json_cell(row.get("context"), default={})
    return minimal_record(
        str(row.get("query") or ""),
        context if isinstance(context, dict) else {},
        prev_session if isinstance(prev_session, list) else [],
        str(row.get("function") or ""),
        request_id=str(row.get("request_id") or ""),
        timestamp=optional_int(row.get("timestamp")),
    )


def minimal_record(
    query: str,
    context: dict[str, Any],
    prev_session: list[Any],
    target: str,
    *,
    request_id: str = "",
    timestamp: int | None = None,
) -> dict[str, Any]:
    return {
        "source": {"request_id": request_id, "timestamp": timestamp},
        "turn": {"query": query, "timestamp": timestamp},
        "prev_session": prev_session,
        "context": context,
        "label": {"target": target},
        "dimensions": {},
    }


def combined_label(record: dict[str, Any]) -> str:
    target = ""
    label = record.get("label") if isinstance(record.get("label"), dict) else {}
    if isinstance(label, dict):
        target = str(label.get("target") or "")
    dimensions = record.get("dimensions") if isinstance(record.get("dimensions"), dict) else {}
    complex_value = dimensions.get("complex") if isinstance(dimensions, dict) else None
    if isinstance(complex_value, bool):
        return f"complex={'true' if complex_value else 'false'}\n{target}".rstrip()
    return target


def extract_query_from_prompt(prompt: str) -> str:
    match = re.search(r"\[当前query\]\s*\n用户[:：]\s*(.+)", prompt)
    return match.group(1).strip() if match else ""


def normalize_bool_literal(value: Any) -> str:
    text = str(value).strip().lower()
    return "true" if text in {"true", "1", "yes", "y", "是", "复杂"} else "false"


def parse_json_cell(value: Any, default: Any) -> Any:
    if value in (None, ""):
        return default
    if isinstance(value, (dict, list)):
        return value
    try:
        return json.loads(str(value))
    except json.JSONDecodeError:
        return default


def resolve_runtime_path(path: str) -> Path:
    raw = Path(path).expanduser()
    if raw.is_absolute():
        return raw
    scratchpad = Path(str(Path.cwd()))
    if os.environ.get("PYTHON_EXEC_SCRATCHPAD"):
        scratchpad = Path(os.environ["PYTHON_EXEC_SCRATCHPAD"]).expanduser()
    parts = raw.parts
    if parts and parts[0] in {"output", "outputs"}:
        return scratchpad.parent / "output" / Path(*parts[1:])
    if parts and parts[0] in {"input", "inputs"}:
        return scratchpad.parent / "input" / Path(*parts[1:])
    if parts and parts[0] in {"scratchpad", "scratch"}:
        return scratchpad / Path(*parts[1:])
    return raw


def optional_int(value: Any) -> int | None:
    if value in (None, ""):
        return None
    try:
        return int(float(str(value)))
    except (TypeError, ValueError):
        return None


def extract_prediction(decoded: Any, text: str) -> str:
    if isinstance(decoded, dict):
        for key in ("generated_text", "text", "output", "response", "result"):
            value = decoded.get(key)
            if isinstance(value, str):
                return value.strip()
        outputs = decoded.get("outputs")
        if isinstance(outputs, list) and outputs:
            first = outputs[0]
            if isinstance(first, str):
                return first.strip()
            if isinstance(first, dict):
                return extract_prediction(first, json.dumps(first, ensure_ascii=False))
        choices = decoded.get("choices")
        if isinstance(choices, list) and choices:
            first = choices[0]
            if isinstance(first, dict):
                message = first.get("message")
                if isinstance(message, dict) and isinstance(message.get("content"), str):
                    return message["content"].strip()
                if isinstance(first.get("text"), str):
                    return first["text"].strip()
    return text.strip()


def try_json(text: str) -> Any:
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return None


def preview_samples(samples: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [
        {
            "index": item.get("index"),
            "query": item.get("query"),
            "gold_label": item.get("gold_label"),
            "prompt_preview": str(item.get("prompt") or "")[:200],
        }
        for item in samples[:3]
    ]


def require_string(payload: dict[str, Any], key: str) -> str:
    value = payload.get(key)
    if not isinstance(value, str) or not value.strip():
        raise LabelingError(f"{key} is required")
    return value.strip()


def raise_with_mapping(message: str, columns: list[str] | None = None) -> None:
    exc = LabelingError(message)
    exc.extra = {"needs_mapping": True, "columns": columns or []}  # type: ignore[attr-defined]
    raise exc


if __name__ == "__main__":
    raise SystemExit(main())
