#!/usr/bin/env python3
"""ARIS human-in-the-loop refinement helper for the human-in-loop-refine skill.

A refinement round lives in `_refine/<round>/` at the project root:
    INPUTS.md       verbatim human inputs and the materials read
    papers/*.md     one note per reference paper (converted full texts go in papers/raw/)
    REVISION.md     revised plan, change list, responses, landing plan
    STATUS.json     round state; backup/ holds originals copied before a direct landing
Mainline files are not written before landing. Every command prints one JSON object.
"""

import argparse
import json
import os
import re
import shutil
import sys
from datetime import datetime
from pathlib import Path

STAGE = "_refine"
VERIFY_ANCHOR = "_verify-idea-stage"
STATUS_FILE = "STATUS.json"
SNAPSHOT_FILE = ".mainline_snapshot.json"
TEMPLATE_DIR = Path(__file__).resolve().parent.parent / "templates"
# 主线 ARIS 产物：落地前应保持不变
MAINLINE_DIRS = ("idea-stage", "refine-logs", "review-stage", "paper", "research-wiki", "pilot", "src", "results")
MAINLINE_FILES = ("IDEA_REPORT.md", "EXPERIMENT_PLAN.md", "EXPERIMENT_LOG.md", "NARRATIVE_REPORT.md",
                  "AUTO_REVIEW.md", "REVIEW_STATE.json", "findings.md")
SNAPSHOT_CAP = 5000  # 单个目录超过这么多文件时不做逐文件比对，记为 unchecked
SKIP_PARTS = {".git", "__pycache__", "node_modules", ".venv", "venv", ".ipynb_checkpoints"}
ROUND_RE = re.compile(r"^r\d{2,}$")
SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9-]{0,63}$")
# 改动清单的一行：- [ ] C01 [依据|验证|否决] ... ｜目标：<路径>#<章节> ｜验证：<slug>
CHANGE_RE = re.compile(r"^\s*-\s*\[([ xX-])\]\s*(C\d{2,})\s*\[(依据|验证|否决)\](.*)$")
VERIFY_RE = re.compile(r"验证[:：]\s*([a-z0-9][a-z0-9-]{0,63})")
TARGET_RE = re.compile(r"目标[:：]\s*([^\s｜|#]+)")
TODO = "【TODO】"
CLARIFY = "【待澄清】"
STATE = {" ": "pending", "x": "adopted", "X": "adopted", "-": "dropped"}
LAND_TARGETS = ("research-refine", "research-refine-pipeline", "experiment-plan", "experiment-bridge", "paper-write",
                "auto-paper-improvement-loop", "idea-discovery", "direct")
NEXT_HINT = {
    "collect": "fill INPUTS.md (drafts read, verbatim F items, paper list) and every paper note",
    "triage": "write REVISION.md: change list with [依据]/[验证]/[否决], revised plan, responses, landing plan",
    "review": "report the triage to the user; after they approve run: approve <round> --stage triage",
    "verify": "run /verify-manager for each [验证] change and append ｜验证：<slug> to its line",
    "reconcile": "update REVISION.md for every item listed in issues",
    "confirm": "show the final revision; only after the user confirms run: approve <round> --stage landing",
    "land": "run: land <round> --to <skill|direct>, then hand REVISION.md to that skill (or back up and edit)",
    "closed": "nothing left to do",
}


class RefineError(Exception):
    pass


def now():
    return datetime.now().astimezone().isoformat(timespec="seconds")


def resolve_root(raw):
    root = Path(raw).expanduser().resolve()
    if not root.is_dir():
        raise RefineError(f"project root does not exist: {root}")
    if root.parent == root:
        raise RefineError(f"refusing to use a filesystem root as project root: {root}")
    # ARIS 仓库本身不是研究项目，往里写修订产物会污染上游克隆
    if (root / "AGENT_GUIDE.md").is_file() and (root / "skills").is_dir() and (root / "tools").is_dir():
        raise RefineError(f"{root} looks like the ARIS repository itself, not a research project; pass --root <project>")
    return root


def round_dir(root, rid):
    if not ROUND_RE.match(rid):
        raise RefineError(f"invalid round id: {rid!r} (expected r01, r02, ...)")
    return root / STAGE / rid


def rel(root, path):
    return path.relative_to(root).as_posix()


def read(path):
    return path.read_text(encoding="utf-8") if path.is_file() else None


def render(name, fields):
    path = TEMPLATE_DIR / name
    if not path.is_file():
        raise RefineError(f"template missing: {path}")
    text = path.read_text(encoding="utf-8")
    for key, value in fields.items():
        text = text.replace("{{" + key + "}}", str(value))
    return text


def load(root, rid):
    path = round_dir(root, rid) / STATUS_FILE
    if not path.is_file():
        raise RefineError(f"no refinement round {rid!r} (expected {rel(root, path)})")
    return json.loads(path.read_text(encoding="utf-8"))


def save(root, state):
    path = round_dir(root, state["round"]) / STATUS_FILE
    tmp = path.with_name(STATUS_FILE + ".tmp")
    tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    os.replace(tmp, path)


def scan(root):
    """Record size and mtime of mainline files; oversized trees are returned as unchecked."""
    files, unchecked = {}, []
    for name in MAINLINE_DIRS + MAINLINE_FILES:
        target = root / name
        if target.is_file():
            st = target.stat()
            files[name] = [st.st_size, st.st_mtime_ns]
            continue
        if not target.is_dir():
            continue
        found = {}
        for dirpath, dirnames, filenames in os.walk(target):
            dirnames[:] = [d for d in dirnames if d not in SKIP_PARTS]
            for fn in filenames:
                p = Path(dirpath) / fn
                try:
                    st = p.stat()
                except OSError:
                    continue
                found[rel(root, p)] = [st.st_size, st.st_mtime_ns]
            if len(found) > SNAPSHOT_CAP:
                break
        if len(found) > SNAPSHOT_CAP:
            unchecked.append(name)
        else:
            files.update(found)
    return files, unchecked


def mainline_changes(root, rid):
    snap = json.loads((round_dir(root, rid) / SNAPSHOT_FILE).read_text(encoding="utf-8"))
    files, _ = scan(root)
    files = {k: v for k, v in files.items() if k.split("/")[0] not in snap["unchecked"]}
    before = snap["files"]
    return sorted({k for k in before.keys() & files.keys() if before[k] != files[k]}
                  | (files.keys() - before.keys()) | (before.keys() - files.keys()))


def verify_status(root, slug):
    path = root / VERIFY_ANCHOR / slug / STATUS_FILE
    try:
        s = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    return {"status": s.get("status"), "verdict": s.get("verdict")}


def parse_changes(text):
    changes = []
    for line in (text or "").splitlines():
        m = CHANGE_RE.match(line)
        if not m:
            continue
        v, t = VERIFY_RE.search(m.group(4)), TARGET_RE.search(m.group(4))
        changes.append({"id": m.group(2), "class": m.group(3), "state": STATE[m.group(1)],
                        "verify": v.group(1) if v else None, "target": t.group(1) if t else None})
    return changes


def progress(root, rid, state):
    """Document completeness, change triage, linked verdicts, and the lifecycle phase they imply."""
    rd = round_dir(root, rid)
    docs = {}
    for key, name in (("inputs", "INPUTS.md"), ("revision", "REVISION.md")):
        text = read(rd / name)
        docs[key] = {"path": rel(root, rd / name), "exists": text is not None,
                     "todo": text.count(TODO) if text else 0, "clarify": text.count(CLARIFY) if text else 0}
    papers = [{"path": rel(root, p), "todo": (read(p) or "").count(TODO)} for p in sorted((rd / "papers").glob("*.md"))]
    changes = parse_changes(read(rd / "REVISION.md"))
    issues, waiting, notes = [], [], []
    for c in changes:
        cid, cls, st = c["id"], c["class"], c["state"]
        if cls == "依据" and st == "pending":
            issues.append({"change": cid, "do": "write it into the revised plan and mark [x]"})
        elif cls == "否决" and st != "dropped":
            issues.append({"change": cid, "do": "mark [-] and give the reason"})
        elif cls == "验证":
            vs = verify_status(root, c["verify"]) if c["verify"] else None
            c["verification"] = vs
            if not c["verify"]:
                if st == "pending":
                    waiting.append({"change": cid, "wait": "no verification linked yet"})
                elif st == "adopted":
                    notes.append({"change": cid, "note": "adopted without a verification"})
            elif vs is None:
                issues.append({"change": cid, "do": f"verification {c['verify']!r} not found; fix the slug"})
            elif vs["status"] == "open":
                waiting.append({"change": cid, "wait": f"verification {c['verify']} still open"})
            elif vs["verdict"] == "supported" and st == "pending":
                issues.append({"change": cid, "do": "supported: write it into the revised plan and mark [x]"})
            elif vs["verdict"] == "refuted" and st != "dropped":
                issues.append({"change": cid, "do": "refuted: mark [-] and record why in the responses"})
            elif vs["verdict"] != "supported" and st == "pending":
                issues.append({"change": cid, "do": "inconclusive: ask the user to adopt, drop, or run a follow-up verification"})
            elif vs["verdict"] != "supported" and st == "adopted":
                notes.append({"change": cid, "note": f"adopted although the verification was {vs['verdict']}"})

    def unfinished(d):
        return not d["exists"] or d["todo"] > 0 or d["clarify"] > 0

    if state["status"] in ("landed", "closed"):
        phase = "closed"
    elif unfinished(docs["inputs"]) or any(p["todo"] for p in papers):
        phase = "collect"
    elif unfinished(docs["revision"]) or not changes:
        phase = "triage"
    elif not state["triage_approved"]:
        phase = "review"
    elif waiting:
        phase = "verify"
    elif issues:
        phase = "reconcile"
    elif not state["landing_approved"]:
        phase = "confirm"
    else:
        phase = "land"
    counts = {"total": len(changes)}
    for key in ("依据", "验证", "否决"):
        counts[key] = sum(1 for c in changes if c["class"] == key)
    for key in ("pending", "adopted", "dropped"):
        counts[key] = sum(1 for c in changes if c["state"] == key)
    return {"phase": phase, "next": NEXT_HINT[phase], "docs": docs, "papers": papers, "counts": counts,
            "changes": changes[:100], "issues": issues, "waiting": waiting, "notes": notes}


def cmd_new(args, root):
    title = (args.title or "").strip() or "方案修订"
    base = root / STAGE
    numbers = [int(p.name[1:]) for p in base.iterdir() if p.is_dir() and ROUND_RE.match(p.name)] if base.is_dir() else []
    rid = f"r{(max(numbers) + 1) if numbers else 1:02d}"
    rd = round_dir(root, rid)
    created = now()
    fields = {"title": title, "round": rid, "created": created, "inputs": rel(root, rd / "INPUTS.md")}
    # 先渲染全部模板，模板缺失时不留下半成品目录
    rendered = {name: render(name, fields) for name in ("INPUTS.md", "REVISION.md")}
    (rd / "papers").mkdir(parents=True)
    for name, text in rendered.items():
        (rd / name).write_text(text, encoding="utf-8")
    files, unchecked = scan(root)
    (rd / SNAPSHOT_FILE).write_text(json.dumps({"files": files, "unchecked": unchecked}, ensure_ascii=False), encoding="utf-8")
    save(root, {"round": rid, "title": title, "status": "open", "created": created, "updated": created,
                "triage_approved": False, "landing_approved": False, "papers": [], "landing": None,
                "history": [{"at": created, "event": "created"}]})
    out = {"round": rid, "phase": "collect", "next": NEXT_HINT["collect"],
           "docs": {name: rel(root, rd / name) for name in rendered}}
    if not any((root / name).exists() for name in MAINLINE_DIRS + MAINLINE_FILES):
        out["warning"] = "no ARIS mainline found at this root; confirm it is the intended project"
    return out


def cmd_paper(args, root):
    state = load(root, args.round)
    if not SLUG_RE.match(args.slug):
        raise RefineError(f"invalid paper slug {args.slug!r}: use lowercase letters, digits and hyphens")
    path = round_dir(root, args.round) / "papers" / f"{args.slug}.md"
    if path.exists():
        raise RefineError(f"paper note already exists: {rel(root, path)}")
    path.write_text(render("PAPER_NOTE.md", {"paper": args.title or args.slug, "source": args.source}), encoding="utf-8")
    state["papers"].append({"slug": args.slug, "source": args.source})
    state["updated"] = now()
    save(root, state)
    return {"round": args.round, "note": rel(root, path)}


def cmd_check(args, root):
    state = load(root, args.round)
    out = {"round": args.round, "status": state["status"], **progress(root, args.round, state)}
    if state["status"] not in ("landed", "closed"):
        touched = mainline_changes(root, args.round)
        out.update(mainline_touched=bool(touched), mainline_changes=touched[:20], mainline_changes_total=len(touched))
    return out


def cmd_approve(args, root):
    state = load(root, args.round)
    phase = progress(root, args.round, state)["phase"]
    if args.stage == "triage":
        if phase in ("collect", "triage", "closed"):
            raise RefineError(f"the triage is not ready for approval (phase={phase})")
        state["triage_approved"] = True
    else:
        if phase != "confirm":
            raise RefineError(f"landing can be approved only in the confirm phase (phase={phase})")
        state["landing_approved"] = True
    when = now()
    state["updated"] = when
    state["history"].append({"at": when, "event": f"approved {args.stage}"})
    save(root, state)
    return {"round": args.round, "approved": args.stage, "phase": progress(root, args.round, state)["phase"]}


def mainline_path(root, raw):
    path = (root / raw).resolve()
    try:
        r = path.relative_to(root).as_posix()
    except ValueError:
        raise RefineError(f"{raw} is outside the project root") from None
    top = r.split("/")[0]
    if top == STAGE or top.startswith("_verify-"):
        raise RefineError(f"{raw} is not a mainline path")
    return path, r


def cmd_backup(args, root):
    state = load(root, args.round)
    if state["status"] in ("landed", "closed"):
        raise RefineError(f"round {args.round} is already {state['status']}")
    targets = []
    for raw in args.paths:
        path, r = mainline_path(root, raw)
        if not path.exists():
            raise RefineError(f"{raw} does not exist")
        targets.append((path, r))
    done, skipped = [], []
    for path, r in targets:
        dest = round_dir(root, args.round) / "backup" / r
        if dest.exists():
            skipped.append(r)  # 保留第一次备份的原文，不被后来的修改覆盖
            continue
        dest.parent.mkdir(parents=True, exist_ok=True)
        if path.is_dir():
            shutil.copytree(path, dest)
        else:
            shutil.copy2(path, dest)
        done.append(r)
    when = now()
    state["updated"] = when
    state["history"].append({"at": when, "event": "backup", "paths": done})
    save(root, state)
    return {"round": args.round, "backed_up": done, "already_backed_up": skipped}


def cmd_land(args, root):
    state = load(root, args.round)
    if state["status"] in ("landed", "closed"):
        raise RefineError(f"round {args.round} is already {state['status']}")
    if not state["landing_approved"]:
        raise RefineError("landing needs the user's confirmation first: approve <round> --stage landing")
    prog = progress(root, args.round, state)
    if prog["phase"] != "land":
        raise RefineError(f"round is not ready to land (phase={prog['phase']}): {prog['next']}")
    targets = list(dict.fromkeys(args.to))
    if "direct" in targets:
        missing = set()
        for c in prog["changes"]:
            if c["state"] != "adopted" or not c["target"]:
                continue
            path, r = mainline_path(root, c["target"])
            if path.exists() and not (round_dir(root, args.round) / "backup" / r).exists():
                missing.add(r)
        if missing:
            raise RefineError("back up these files before a direct landing: " + ", ".join(sorted(missing)))
    when = now()
    state.update(status="landed", updated=when, landing={"to": targets, "at": when})
    state["history"].append({"at": when, "event": "landed", "to": targets})
    save(root, state)
    return {"round": args.round, "landed_to": targets,
            "next": "edit the backed-up files per the change list" if targets == ["direct"]
            else "run the mainline skill(s) with REVISION.md as input"}


def cmd_list(args, root):
    items = []
    base = root / STAGE
    if base.is_dir():
        for d in sorted(base.iterdir()):
            if not (d.is_dir() and ROUND_RE.match(d.name) and (d / STATUS_FILE).is_file()):
                continue
            try:
                s = load(root, d.name)
            except ValueError as exc:
                items.append({"round": d.name, "error": f"unreadable {STATUS_FILE}: {exc}"})
                continue
            if s["status"] == "closed" and not args.all:
                continue
            p = progress(root, d.name, s)
            items.append({"round": d.name, "title": s["title"], "status": s["status"], "phase": p["phase"],
                          "changes": p["counts"]["total"], "updated": s["updated"]})
    return {"count": len(items), "rounds": items}


def cmd_close(args, root):
    state = load(root, args.round)
    if state["status"] == "closed":
        return {"round": args.round, "status": "closed", "note": "already closed"}
    when = now()
    state.update(status="closed", updated=when, close_reason=args.reason)
    state["history"].append({"at": when, "event": "closed", "reason": args.reason})
    save(root, state)
    return {"round": args.round, "status": "closed"}


def main(argv=None):
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--root", default=".", help="research project root (default: current directory)")
    parser = argparse.ArgumentParser(prog="refine_round.py", description="ARIS human-in-the-loop refinement rounds")
    sub = parser.add_subparsers(dest="command", required=True)
    p = sub.add_parser("new", parents=[common], help="start a refinement round")
    p.add_argument("--title")
    p = sub.add_parser("paper", parents=[common], help="create a note for one reference paper")
    p.add_argument("round")
    p.add_argument("--slug", required=True)
    p.add_argument("--source", required=True)
    p.add_argument("--title")
    p = sub.add_parser("check", parents=[common], help="report progress and the current phase")
    p.add_argument("round")
    p = sub.add_parser("approve", parents=[common], help="record the user's approval")
    p.add_argument("round")
    p.add_argument("--stage", required=True, choices=("triage", "landing"))
    p = sub.add_parser("backup", parents=[common], help="copy mainline files before a direct landing")
    p.add_argument("round")
    p.add_argument("paths", nargs="+")
    p = sub.add_parser("land", parents=[common], help="record where the revision lands")
    p.add_argument("round")
    p.add_argument("--to", required=True, action="append", choices=LAND_TARGETS)
    p = sub.add_parser("list", parents=[common], help="list rounds")
    p.add_argument("--all", action="store_true", help="include closed rounds")
    p = sub.add_parser("close", parents=[common], help="abandon a round")
    p.add_argument("round")
    p.add_argument("--reason")
    args = parser.parse_args(argv)
    handlers = {"new": cmd_new, "paper": cmd_paper, "check": cmd_check, "approve": cmd_approve,
                "backup": cmd_backup, "land": cmd_land, "list": cmd_list, "close": cmd_close}
    try:
        result = handlers[args.command](args, resolve_root(args.root))
    except RefineError as exc:
        print(json.dumps({"ok": False, "error": str(exc)}, ensure_ascii=False))
        return 1
    print(json.dumps({"ok": True, **result}, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())