"""Roll keyword_triager + memory_aware policies; emit (prompt, completion) JSONL for SFT.

Usage:
    python scripts/build_sft_traces.py --out data/sft_traces.jsonl --scenarios E0_quiet_day E1_release_week E2_spam_wave
"""
from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from eval.policies import build_policy
from server.opsguard_environment import OpsguardEnvironment


SYSTEM_PROMPT = """You triage GitHub issues for an open-source repository. Each step you see ONE issue or PR. You must take ONE action.

Available actions (emit STRICT JSON, single line):
{"action_type":"label","target_issue_id":N,"label":"<name>","reasoning":"..."}
{"action_type":"close_spam","target_issue_id":N,"reasoning":"..."}
{"action_type":"request_info","target_issue_id":N,"comment_body":"...","reasoning":"..."}
{"action_type":"link_duplicate","target_issue_id":N,"duplicate_of_id":M,"reasoning":"..."}
{"action_type":"assign","target_issue_id":N,"assignee_login":"...","reasoning":"..."}
{"action_type":"comment","target_issue_id":N,"comment_body":"...","reasoning":"..."}
{"action_type":"merge_pr","target_issue_id":N,"reasoning":"..."}
{"action_type":"query_history","query":"...","reasoning":"..."}
{"action_type":"wait"}

Goal: maximize action_match (matching the maintainer's real action) and spam_resistance (close_spam on synthetic spam, do NOT label/merge/assign spam). Avoid repeated actions on the same issue. Use query_history sparingly.
"""


def format_observation(obs) -> str:
    payload = {
        "scenario": obs.scenario_id,
        "step": obs.step,
        "step_budget": obs.step_budget,
        "queue_position": obs.queue_position,
        "queue_total": obs.queue_total,
        "feedback": obs.feedback,
        "memory_hits": obs.memory_hits,
        "recent_actions": obs.recent_actions,
    }
    if obs.current_issue is not None:
        ci = obs.current_issue
        payload["current_issue"] = {
            "issue_id": ci.issue_id,
            "number": ci.number,
            "title": ci.title,
            "body": ci.body[:1200],
            "is_pr": ci.is_pr,
            "author_login": ci.author_login,
            "author_pr_count": ci.author_pr_count,
            "author_account_age_days": ci.author_account_age_days,
            "available_labels": ci.available_labels[:25],
            "comments_preview": ci.comments_preview,
        }
    return json.dumps(payload, ensure_ascii=False)


def _build_reasoning_text(action, obs) -> str:
    at = action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type)
    ci = obs.current_issue
    rationale = action.reasoning or "applying triage heuristic"
    parts = []
    if ci is not None:
        parts.append(
            f"Issue #{ci.number} by {ci.author_login} "
            f"(prior PRs={ci.author_pr_count}, account_age_days={ci.author_account_age_days})."
        )
        if ci.is_pr and ci.pr_diff_preview:
            parts.append(f"PR diff preview indicates: {ci.pr_diff_preview[:160]}.")
        if ci.pr_dependency_changes:
            parts.append(f"Dependency changes: {ci.pr_dependency_changes[:2]}.")
    parts.append(f"Selecting action={at} because: {rationale}.")
    return " ".join(parts)


def _format_cot_completion(action, obs) -> str:
    reasoning = _build_reasoning_text(action, obs)
    comp = action.model_dump(mode="json", exclude_none=True)
    return f"REASONING: {reasoning}\nACTION: {json.dumps(comp, ensure_ascii=False)}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/sft_traces.jsonl")
    ap.add_argument("--scenarios", nargs="*",
                    default=["E0_normal_intake", "E1_typosquat_wave", "E2_social_eng_buildup", "E3_compromised_maintainer"])
    ap.add_argument("--seeds", nargs="*", type=int, default=[0, 1, 2])
    ap.add_argument("--policies", nargs="*", default=["keyword_security_triager", "memory_aware"])
    ap.add_argument("--cot-ratio", type=float, default=0.5,
                    help="Fraction of pairs emitted in REASONING+ACTION format (R2Vul-style).")
    ap.add_argument("--seed", type=int, default=17)
    args = ap.parse_args()

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    rng = random.Random(args.seed)

    n_pairs = 0
    n_cot = 0
    n_json = 0
    with open(out, "w", encoding="utf-8") as f:
        for pname in args.policies:
            policy = build_policy(pname)
            for sid in args.scenarios:
                for seed in args.seeds:
                    env = OpsguardEnvironment()
                    obs = env.reset(scenario_id=sid, seed=seed)
                    while not obs.done and obs.step < env._episode.scenario.step_budget:
                        action = policy(obs)
                        obs_text = format_observation(obs)
                        prompt = SYSTEM_PROMPT + "\n\nOBSERVATION:\n" + obs_text + "\n\nReturn ONE action as strict JSON."
                        if rng.random() < args.cot_ratio:
                            completion = _format_cot_completion(action, obs)
                            fmt = "reasoning_action"
                            n_cot += 1
                        else:
                            comp = action.model_dump(mode="json", exclude_none=True)
                            completion = json.dumps(comp, ensure_ascii=False)
                            fmt = "json_only"
                            n_json += 1
                        f.write(json.dumps({
                            "prompt": prompt,
                            "completion": completion,
                            "format": fmt,
                            "policy": pname, "scenario": sid, "seed": seed,
                        }) + "\n")
                        n_pairs += 1
                        obs = env.step(action)
    print(f"DONE: wrote {n_pairs} (prompt, completion) pairs to {out} "
          f"[reasoning_action={n_cot}, json_only={n_json}]")


if __name__ == "__main__":
    main()