#!/usr/bin/env python3
"""AgentRE local launcher.

`chat` starts a persistent reverse-engineering session for one file, or a directory passed with `-d`.
`inspect` is kept as a compatibility alias for `chat`.
`triage` delegates to the existing bulk PE/ELF analyzer.
"""

from __future__ import annotations

import argparse
import json
import os
import shlex
import shutil
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import agentre_triage as triage


ROOT = Path(__file__).resolve().parent
LOCAL_AGENTRE_ENV = ROOT / "scripts" / "xref9b.env"


def read_shell_exports(path: Path) -> dict[str, str]:
    exports: dict[str, str] = {}
    if not path.exists():
        return exports
    for raw in path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = raw.strip()
        if not line.startswith("export "):
            continue
        try:
            parts = shlex.split(line)
        except ValueError:
            continue
        for part in parts[1:]:
            key, sep, value = part.partition("=")
            if sep and key:
                exports[key] = value
    return exports


def load_local_agentre_env() -> dict[str, str]:
    exports = read_shell_exports(LOCAL_AGENTRE_ENV)
    for key, value in exports.items():
        os.environ.setdefault(key, value)
    return exports


def now_stamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")


def print_top_help() -> None:
    print(
        """AgentRE local PE/ELF reverse-engineering CLI.

Usage:
  python3 agentre.py chat FILE [options]
  python3 agentre.py chat -d DIR [options]
  python3 agentre.py inspect FILE [options]
  python3 agentre.py inspect -d DIR [options]
  python3 agentre.py triage FILE_OR_DIR [options]

Commands:
  chat      Chat with the model while it uses static PE/ELF RE tools. Use -d for directories.
  inspect   Alias for chat.
  triage    Analyze one file or a directory of PE/ELF files and write reports.

Examples:
  python3 agentre.py chat holdout/mixed40/staged/sample_0003.elf
  python3 agentre.py chat -d holdout/mixed40/staged
  python3 agentre.py triage holdout/mixed40/staged --model ./xref-9b-q4_k_m.gguf
"""
    )


def build_backend(args: argparse.Namespace) -> tuple[str, Any]:
    backend_name = args.backend
    if backend_name == "auto":
        backend_name = "llama-cpp" if str(args.model).lower().endswith(".gguf") else "transformers"
    print(f"[xref 9b] backend={backend_name} model={args.model}", flush=True)
    if args.adapter and backend_name == "transformers":
        print(f"[xref 9b] adapter={args.adapter}", flush=True)
    elif args.adapter and backend_name == "llama-cpp":
        print("[xref 9b] note: --adapter is ignored by llama-cpp; use a merged GGUF or --llama-lora", flush=True)

    if backend_name == "transformers":
        return backend_name, triage.LocalQwenBackend(args.model, args.adapter or None, args.dtype, args.device_map)
    return backend_name, triage.LlamaCppBackend(
        args.model,
        args.llama_cli,
        args.ctx_size,
        args.threads,
        args.gpu_layers,
        args.llama_lora or None,
        args.llama_extra_arg,
    )


def build_inspect_messages(spec: str, sample_id: str, fmt: str, request: str) -> list[dict[str, Any]]:
    system = (
        spec.strip()
        + "\n\nRuntime rules:\n"
        + f"- You are in an interactive reverse-engineering session for staged sample `{sample_id}` only.\n"
        + f"- Detected format: {fmt}.\n"
        + "- The original path is hidden from the model to avoid path/name bias.\n"
        + "- Any tool path argument is ignored; tools are routed to the staged sample.\n"
        + "- Use static analysis only. Never ask to execute the sample.\n"
        + "- Use tools when they would materially improve the answer.\n"
        + "- Answer the user's follow-up questions directly and cite concrete evidence from tool output.\n"
        + "- If evidence is insufficient, say so plainly; classify as unknown or suspicious/unknown rather than guessing.\n"
        + "- For stripped, static, packed, encrypted, or sparse-string samples where confidence is low, explicitly recommend closer disassembly with Ghidra or another disassembler.\n"
        + "- Think internally, but do not expose chain-of-thought.\n"
        + "- Only call final_answer when the user asks for a final verdict or when your answer is a classification decision.\n"
    )
    user = (
        f"{request.strip()}\n\n"
        "Start by inspecting basic metadata and strings, then give a concise initial assessment. "
        "Do not require the user to know tool names."
    )
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def stage_one(source_path: Path, run_dir: Path, sample_id: str = "inspect_0001") -> tuple[Path, str]:
    fmt = triage.detect_format(source_path)
    suffix = ".elf" if fmt == "ELF" else ".exe" if fmt == "PE" else source_path.suffix
    stage_dir = run_dir / "staged"
    stage_dir.mkdir(parents=True, exist_ok=True)
    staged_path = (stage_dir / f"{sample_id}{suffix}").resolve()
    shutil.copy2(source_path, staged_path)
    return staged_path, fmt


def progress(state: dict[str, Any], message: str) -> None:
    if state.get("progress", True):
        print(f"[xref 9b] {message}", file=sys.stderr, flush=True)


def looks_truncated_reply(text: str) -> bool:
    stripped = text.rstrip()
    if len(stripped) < 600:
        return False
    if stripped.endswith((".", "!", "?", ")", "]", "}", "`")):
        return False
    tail = stripped.rsplit(maxsplit=1)[-1].lower().strip(",:;")
    return tail in {
        "a",
        "an",
        "and",
        "are",
        "as",
        "because",
        "but",
        "for",
        "from",
        "given",
        "in",
        "is",
        "of",
        "or",
        "that",
        "the",
        "this",
        "to",
        "with",
    } or not stripped.endswith((".", "!", "?", ":", ";"))


def requested_direct_tool(text: str) -> tuple[str, dict[str, Any]] | None:
    lowered = text.lower()
    if "ghidra" in lowered and any(
        word in lowered
        for word in ("analyze", "disassembl", "dive", "open", "run", "try", "use", "with")
    ):
        return "ghidra_summary", {"timeout": 180}
    return None


def format_final(answer: dict[str, Any] | None) -> str:
    if not answer:
        return "Final verdict was submitted, but it was not valid JSON."
    label = triage.normalize_prediction(answer)
    summary = answer.get("summary") or answer.get("rationale") or answer.get("reason") or ""
    evidence = answer.get("evidence") or answer.get("key_evidence") or []
    lines = [f"Final verdict: {label}"]
    if summary:
        lines.append(str(summary))
    if isinstance(evidence, list) and evidence:
        lines.append("Evidence:")
        for item in evidence[:6]:
            lines.append(f"- {item}")
    return "\n".join(lines)


def shorten_text(text: str, limit: int) -> str:
    stripped = "\n".join(line.rstrip() for line in text.strip().splitlines())
    if len(stripped) <= limit:
        return stripped
    return stripped[:limit].rstrip() + f"\n[truncated to {limit} chars]"


def deep_tool_sequence(fmt: str) -> list[tuple[str, dict[str, Any]]]:
    common: list[tuple[str, dict[str, Any]]] = [
        ("file", {}),
        ("entropy", {}),
        ("strings", {"min_length": 5}),
    ]
    if fmt == "ELF":
        return [
            ("file", {}),
            ("readelf", {}),
            ("nm", {}),
            ("entropy", {}),
            ("strings", {"min_length": 5}),
            ("objdump", {}),
            ("ghidra_summary", {"timeout": 240}),
        ]
    if fmt == "PE":
        return [
            ("file", {}),
            ("pe_headers", {}),
            ("pe_sections", {}),
            ("pe_imports", {}),
            ("pe_exports", {}),
            ("entropy", {}),
            ("strings", {"min_length": 5}),
            ("pe_disasm", {}),
            ("ghidra_summary", {"timeout": 240}),
        ]
    return common + [("ghidra_summary", {"timeout": 240})]


def compact_messages(
    messages: list[dict[str, Any]],
    tools: triage.StaticTools,
    state: dict[str, Any],
    max_tool_summaries: int = 12,
    max_assistant_summaries: int = 5,
) -> list[dict[str, Any]]:
    system_msg = next((msg for msg in messages if msg.get("role") == "system"), None)
    first_user = next((msg for msg in messages if msg.get("role") == "user"), None)

    tool_msgs = [msg for msg in messages if msg.get("role") == "tool"][-max_tool_summaries:]
    assistant_msgs = [
        msg for msg in messages
        if msg.get("role") == "assistant" and str(msg.get("content") or "").strip()
    ][-max_assistant_summaries:]

    lines = [
        "Session context was compacted to reduce prompt length.",
        f"Total tool calls before compaction: {state.get('tool_calls_seen', 0)}",
        "Tools used: " + (", ".join(tools.used_tools[-40:]) if tools.used_tools else "none"),
    ]
    if state.get("last_final"):
        lines.append("Last submitted final verdict JSON: " + json.dumps(state["last_final"], ensure_ascii=True)[:1600])

    if tool_msgs:
        lines.append("\nRecent tool observations:")
        for msg in tool_msgs:
            name = str(msg.get("name") or "tool")
            content = shorten_text(str(msg.get("content") or ""), 900)
            lines.append(f"\n[{name}]\n{content}")

    if assistant_msgs:
        lines.append("\nRecent assistant conclusions:")
        for msg in assistant_msgs:
            content = shorten_text(str(msg.get("content") or ""), 700)
            lines.append(f"- {content}")

    lines.append("\nIf exact old bytes or full tool output are needed, rerun a targeted tool or /deep.")

    compacted: list[dict[str, Any]] = []
    if system_msg:
        compacted.append(system_msg)
    if first_user:
        compacted.append(first_user)
    compacted.append({"role": "assistant", "content": "\n".join(lines), "reasoning_content": ""})
    return compacted


def save_compaction_snapshot(run_dir: Path, session: dict[str, Any]) -> Path:
    snapshot_dir = run_dir / "compaction_snapshots"
    snapshot_dir.mkdir(parents=True, exist_ok=True)
    sample_id = str(session.get("sample_id") or "sample")
    path = snapshot_dir / f"{sample_id}_{now_stamp()}_precompact.json"
    payload = {
        "sample_id": sample_id,
        "source_path": str(session.get("source_path")),
        "staged_path": str(session.get("staged_path")),
        "format": session.get("fmt"),
        "tools_used": session.get("tools").used_tools if session.get("tools") else [],
        "state": session.get("state", {}),
        "messages": session.get("messages", []),
        "saved_at": datetime.now(timezone.utc).isoformat(),
    }
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=True) + "\n", encoding="utf-8")
    return path


def append_direct_tool_result(
    canonical: str,
    args: dict[str, Any],
    tools: triage.StaticTools,
    messages: list[dict[str, Any]],
    cfg: triage.RuntimeConfig,
    state: dict[str, Any],
) -> None:
    call_id = f"manual_{int(time.time() * 1000)}"
    started = time.monotonic()
    state["tool_calls_seen"] = int(state.get("tool_calls_seen", 0)) + 1
    progress(state, f"direct tool {state['tool_calls_seen']}: {canonical}")
    try:
        content = getattr(tools, canonical)(**args)
    except Exception as exc:
        content = f"tool error: {type(exc).__name__}: {exc}"
    elapsed = time.monotonic() - started
    progress(state, f"tool done: {canonical} in {elapsed:.1f}s, observation={len(str(content))} chars")

    messages.append(
        {
            "role": "assistant",
            "content": "",
            "reasoning_content": "",
            "tool_calls": [
                {
                    "id": call_id,
                    "type": "function",
                    "function": {"name": canonical, "arguments": args},
                }
            ],
        }
    )
    messages.append(
        {
            "role": "tool",
            "tool_call_id": call_id,
            "name": canonical,
            "content": str(content)[: cfg.obs_limit],
        }
    )


def save_inspect_session(
    run_dir: Path,
    source_path: Path,
    staged_path: Path,
    fmt: str,
    messages: list[dict[str, Any]],
    tools: triage.StaticTools,
    state: dict[str, Any],
) -> None:
    sample_id = str(state.get("sample_id") or staged_path.stem)
    payload = {
        "sample_id": sample_id,
        "source_path": str(source_path),
        "staged_path": str(staged_path),
        "format": fmt,
        "tools_used": tools.used_tools,
        "tool_calls": state.get("tool_calls_seen", 0),
        "last_final": state.get("last_final"),
        "messages": messages,
        "saved_at": datetime.now(timezone.utc).isoformat(),
    }
    json_text = json.dumps(payload, indent=2, ensure_ascii=True) + "\n"
    (run_dir / "inspect_transcript.json").write_text(json_text, encoding="utf-8")
    (run_dir / f"{sample_id}_transcript.json").write_text(json_text, encoding="utf-8")

    lines = [
        "# AgentRE Chat Session",
        "",
        f"- Sample: `{sample_id}`",
        f"- Source: `{source_path}`",
        f"- Staged: `{staged_path}`",
        f"- Format: `{fmt}`",
        f"- Tool calls: `{state.get('tool_calls_seen', 0)}`",
        "",
        "## Conversation",
        "",
    ]
    for msg in messages:
        role = msg.get("role")
        if role not in {"user", "assistant"}:
            continue
        content = str(msg.get("content") or "").strip()
        if not content:
            continue
        lines.append(f"### {role}")
        lines.append("")
        lines.append(content[:8000])
        lines.append("")
    md_text = "\n".join(lines)
    (run_dir / "INSPECT.md").write_text(md_text, encoding="utf-8")
    (run_dir / f"{sample_id}.md").write_text(md_text, encoding="utf-8")


def run_tool_calls(
    tool_calls: list[dict[str, Any]],
    tools: triage.StaticTools,
    messages: list[dict[str, Any]],
    cfg: triage.RuntimeConfig,
    state: dict[str, Any],
    turn: int,
) -> str | None:
    final_text: str | None = None
    for idx, tc in enumerate(tool_calls):
        fn = tc.get("function", {})
        name = str(fn.get("name") or "")
        args = triage._parse_args(fn.get("arguments"))
        tcid = tc.get("id") or f"inspect_{turn}_{idx}"
        started = time.monotonic()

        if name in triage.FINAL_TOOLS:
            progress(state, "model submitted final_answer")
            answer = triage.final_answer_from_args(args)
            state["last_final"] = answer
            content = "final verdict submitted"
            final_text = format_final(answer)
        else:
            canonical = triage.TOOL_ALIASES.get(name)
            if canonical is None:
                progress(state, f"unknown tool requested: {name}")
                content = f"unknown tool {name}. Use listed tools or answer naturally."
            elif int(state.get("request_tool_calls_seen", 0)) >= cfg.max_tool_calls:
                progress(state, "tool budget exhausted for this request")
                content = "tool budget exhausted for this request; answer with the evidence already gathered, or ask the user to continue."
            else:
                state["request_tool_calls_seen"] = int(state.get("request_tool_calls_seen", 0)) + 1
                state["tool_calls_seen"] += 1
                progress(
                    state,
                    f"tool {state['request_tool_calls_seen']}/{cfg.max_tool_calls} this request "
                    f"(total {state['tool_calls_seen']}): {canonical}",
                )
                try:
                    content = getattr(tools, canonical)(**args)
                except Exception as exc:
                    content = f"tool error: {type(exc).__name__}: {exc}"
                elapsed = time.monotonic() - started
                progress(state, f"tool done: {canonical} in {elapsed:.1f}s, observation={len(str(content))} chars")

        messages.append(
            {
                "role": "tool",
                "tool_call_id": tcid,
                "name": name,
                "content": str(content)[: cfg.obs_limit],
            }
        )
    return final_text


def generate_visible_reply(
    backend: Any,
    messages: list[dict[str, Any]],
    cfg: triage.RuntimeConfig,
    tools: triage.StaticTools,
    state: dict[str, Any],
) -> str:
    final_text: str | None = None
    last_visible_reply = ""
    own_tool_budget = "request_tool_calls_seen" not in state
    if own_tool_budget:
        state["request_tool_calls_seen"] = 0
    for turn in range(cfg.max_turns):
        progress(state, f"model turn {turn + 1}/{cfg.max_turns}: generating")
        started = time.monotonic()
        generated, step_usage = backend.generate(
            messages,
            triage.tool_definitions(),
            cfg.max_tokens,
            cfg.temperature,
            cfg.thinking,
        )
        elapsed = time.monotonic() - started
        progress(state, f"model turn {turn + 1}/{cfg.max_turns}: generated in {elapsed:.1f}s")
        usage = state.setdefault("usage", {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
        for key, value in step_usage.items():
            usage[key] = usage.get(key, 0) + int(value)

        reasoning, visible = triage.split_thinking(generated)
        tool_calls = triage.parse_tool_calls(generated)
        if tool_calls:
            progress(state, f"model requested {len(tool_calls)} tool call(s)")
        assistant_content = triage.strip_tool_xml(visible)
        if assistant_content:
            last_visible_reply = assistant_content
        assistant_msg: dict[str, Any] = {"role": "assistant", "content": assistant_content}
        if reasoning and cfg.save_reasoning:
            assistant_msg["reasoning_content"] = reasoning
        elif reasoning:
            assistant_msg["reasoning_content"] = ""
        if tool_calls:
            assistant_msg["tool_calls"] = [
                triage.make_history_tool_call(tc, f"inspect_{turn}_{idx}")
                for idx, tc in enumerate(tool_calls)
            ]
        messages.append(assistant_msg)

        if not tool_calls:
            reply = assistant_content or generated.strip() or "(no response)"
            if looks_truncated_reply(reply):
                state["last_reply_truncated"] = True
                progress(state, "reply may have hit --max-tokens; use /continue or raise --max-tokens")
                reply += (
                    "\n\n[xref 9b] response may have hit --max-tokens before finishing; "
                    "use /continue or raise --max-tokens for longer answers."
                )
            if own_tool_budget:
                state.pop("request_tool_calls_seen", None)
            return reply

        maybe_final = run_tool_calls(tool_calls, tools, messages, cfg, state, turn)
        if maybe_final:
            final_text = maybe_final
            messages.append(
                {
                    "role": "user",
                    "content": "Summarize that final verdict for the user in concise natural language. Do not call another tool unless needed.",
                }
            )

    if final_text:
        if own_tool_budget:
            state.pop("request_tool_calls_seen", None)
        return final_text
    if last_visible_reply:
        if own_tool_budget:
            state.pop("request_tool_calls_seen", None)
        return (
            last_visible_reply
            + "\n\n[xref 9b] internal tool loop reached --max-turns before a final tool submission; "
            + "ask a follow-up or raise --max-turns if you want deeper analysis."
        )
    if own_tool_budget:
        state.pop("request_tool_calls_seen", None)
    return "[xref 9b] assistant turn stopped after the internal turn limit; ask a narrower follow-up or raise --max-turns."


def default_chat_model() -> str:
    packaged = ROOT / "xref-9b-q4_k_m.gguf"
    return (
        os.environ.get("XREF9B_GGUF")
        or os.environ.get("AGENTRE_GGUF")
        or os.environ.get("AGENTRE_MODEL")
        or (str(packaged) if packaged.exists() else triage.DEFAULT_BASE_MODEL)
    )


def default_llama_cli() -> str:
    return os.environ.get("XREF9B_LLAMA_CLI") or os.environ.get("AGENTRE_LLAMA_CLI") or triage.DEFAULT_LLAMA_CLI


def default_adapter_for_model(model: str) -> str:
    return "" if str(model).lower().endswith(".gguf") else triage.DEFAULT_ADAPTER


def sample_id_for_index(index: int, total: int) -> str:
    return "inspect_0001" if total == 1 else f"sample_{index + 1:04d}"


def display_path(path: Path, root: Path) -> str:
    try:
        base = root if root.is_dir() else root.parent
        return str(path.relative_to(base))
    except Exception:
        return str(path)


def save_chat_index(
    run_dir: Path,
    target_root: Path,
    targets: list[Path],
    active_idx: int,
    sessions: dict[int, dict[str, Any]],
) -> None:
    samples: list[dict[str, Any]] = []
    for idx, path in enumerate(targets):
        session = sessions.get(idx)
        state = session.get("state", {}) if session else {}
        samples.append(
            {
                "index": idx + 1,
                "sample_id": sample_id_for_index(idx, len(targets)),
                "path": str(path),
                "display_path": display_path(path, target_root),
                "format": session.get("fmt") if session else triage.detect_format(path),
                "opened": bool(session),
                "tool_calls": state.get("tool_calls_seen", 0),
                "last_final": state.get("last_final"),
            }
        )
    payload = {
        "target_root": str(target_root),
        "active_index": active_idx + 1,
        "samples": samples,
        "saved_at": datetime.now(timezone.utc).isoformat(),
    }
    (run_dir / "chat_state.json").write_text(json.dumps(payload, indent=2, ensure_ascii=True) + "\n", encoding="utf-8")

    lines = [
        "# AgentRE Chat Index",
        "",
        f"- Target: `{target_root}`",
        f"- Active sample: `{active_idx + 1}`",
        f"- Samples: `{len(targets)}`",
        "",
        "## Samples",
        "",
    ]
    for sample in samples:
        marker = "*" if sample["index"] == active_idx + 1 else "-"
        opened = " opened" if sample["opened"] else ""
        lines.append(
            f"{marker} `{sample['index']:04d}` `{sample['format']}`{opened} "
            f"tools={sample['tool_calls']} `{sample['display_path']}`"
        )
    (run_dir / "CHAT.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def print_sample_list(target_root: Path, targets: list[Path], active_idx: int, sessions: dict[int, dict[str, Any]], limit: int = 80) -> None:
    shown = targets[:limit]
    for idx, path in enumerate(shown):
        session = sessions.get(idx)
        fmt = session.get("fmt") if session else triage.detect_format(path)
        marker = "*" if idx == active_idx else " "
        opened = " opened" if session else ""
        print(f"{marker} {idx + 1:4d} {fmt:7s}{opened:8s} {display_path(path, target_root)}")
    if len(targets) > limit:
        print(f"... {len(targets) - limit} more samples not shown; use /open N by index.")


def resolve_sample_selector(selector: str, target_root: Path, targets: list[Path]) -> int:
    text = selector.strip()
    if not text:
        raise ValueError("usage: /open N or /open path-substring")
    if text.isdigit():
        idx = int(text) - 1
        if 0 <= idx < len(targets):
            return idx
        raise ValueError(f"sample index out of range: {text}")

    maybe_path = Path(text).expanduser()
    if maybe_path.exists():
        resolved = maybe_path.resolve()
        for idx, path in enumerate(targets):
            if path == resolved:
                return idx

    lowered = text.lower()
    matches = [
        idx
        for idx, path in enumerate(targets)
        if lowered in path.name.lower() or lowered in display_path(path, target_root).lower() or lowered in str(path).lower()
    ]
    if len(matches) == 1:
        return matches[0]
    if not matches:
        raise ValueError(f"no sample matched: {text}")
    preview = ", ".join(str(idx + 1) for idx in matches[:12])
    raise ValueError(f"ambiguous sample selector; matched indexes: {preview}")


def run_tool_smoke_targets(targets: list[Path], run_dir: Path, cfg: triage.RuntimeConfig) -> int:
    for idx, source_path in enumerate(targets):
        sample_id = sample_id_for_index(idx, len(targets))
        staged_path, fmt = stage_one(source_path, run_dir, sample_id)
        tools = triage.StaticTools(staged_path, fmt, cfg.obs_limit, cfg.ghidra_script_dir)
        print(f"== {idx + 1}: {source_path} ({fmt}) ==")
        print("\n[file]")
        print(tools.file())
        print("\n[strings]")
        print(tools.strings()[:2500])
        print(f"\n[xref 9b] staged={staged_path}")
    return 0


def run_chat(argv: list[str], command_name: str = "chat") -> int:
    parser = argparse.ArgumentParser(
        prog=f"agentre.py {command_name}",
        description="Chat with AgentRE about one PE/ELF sample, or a directory passed with -d, while it uses static RE tools.",
    )
    loaded_env = load_local_agentre_env()
    default_model = default_chat_model()
    default_cli = default_llama_cli()
    parser.add_argument("target", nargs="?", help="Binary file to inspect. Use -d/--directory for directory mode.")
    parser.add_argument("-d", "--directory", default="", help="Directory of PE/ELF files to inspect interactively.")
    parser.add_argument("--spec", default=str(triage.DEFAULT_SPEC), help="Unified reverse-engineering spec.")
    parser.add_argument(
        "--backend",
        choices=["auto", "transformers", "llama-cpp"],
        default="auto",
        help="Inference backend. auto selects llama-cpp for .gguf models, otherwise transformers.",
    )
    parser.add_argument("--model", default=default_model, help="HF base model path/name or GGUF path for llama.cpp.")
    parser.add_argument("--adapter", default=None, help="Optional HF LoRA adapter for transformers. Use '' to disable.")
    parser.add_argument("--output-dir", default="", help="Output directory. Defaults to runs/agentre_chat_<timestamp>.")
    parser.add_argument("--max-files", type=int, default=0, help="Directory mode limit; 0 means no limit.")
    parser.add_argument("--include-unknown", action="store_true", help="Include files that are not PE/ELF by magic bytes.")
    parser.add_argument("--max-turns", type=int, default=8, help="Internal assistant/tool rounds per user message.")
    parser.add_argument("--max-tool-calls", type=int, default=40, help="Model-selected tool calls allowed per user request.")
    parser.add_argument("--max-tokens", type=int, default=1200)
    parser.add_argument("--obs-limit", type=int, default=7000)
    parser.add_argument("--temperature", type=float, default=0.1)
    parser.add_argument("--dtype", default="auto", help="auto, bfloat16, float16, float32.")
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--llama-cli", default=default_cli, help="Path to llama.cpp llama-completion.")
    parser.add_argument("--ctx-size", type=int, default=65536, help="llama.cpp context size.")
    parser.add_argument("--threads", type=int, default=0, help="llama.cpp CPU threads; 0 lets llama.cpp choose.")
    parser.add_argument("--gpu-layers", type=int, default=-1, help="llama.cpp GPU layers; -1 leaves default.")
    parser.add_argument("--llama-lora", default="", help="Optional llama.cpp-compatible LoRA adapter.")
    parser.add_argument("--llama-extra-arg", action="append", default=[], help="Extra raw argument passed to llama-completion. Repeatable.")
    parser.add_argument("--no-thinking", action="store_true", help="Disable Qwen thinking template flag.")
    parser.add_argument("--quiet", action="store_true", help="Hide progress messages while the model/tools run.")
    parser.add_argument("--save-reasoning", action="store_true", help=argparse.SUPPRESS)
    parser.add_argument("--ghidra-script-dir", default=str(triage.DEFAULT_GHIDRA_SCRIPT_DIR))
    parser.add_argument(
        "--request",
        default="Is this malicious or benign? Reverse engineer it and use the available tools.",
        help="Initial user request.",
    )
    parser.add_argument("--one-shot", action="store_true", help="Run the initial inspection and exit instead of opening the prompt.")
    parser.add_argument("--tool-smoke", action="store_true", help="Run static tools and exit without loading a model.")
    args = parser.parse_args(argv)

    if loaded_env and not args.quiet:
        print(f"[xref 9b] loaded defaults from {LOCAL_AGENTRE_ENV}", flush=True)

    args.adapter = default_adapter_for_model(args.model) if args.adapter is None else args.adapter

    if args.directory and args.target:
        print("[xref 9b] provide either one file target or -d/--directory, not both.", file=sys.stderr)
        return 2
    if not args.directory and not args.target:
        print("[xref 9b] provide one binary file, or use -d/--directory DIR for directory mode.", file=sys.stderr)
        return 2

    directory_mode = bool(args.directory)
    target_root = Path(args.directory or args.target).expanduser().resolve()
    if directory_mode:
        if not target_root.is_dir():
            print(f"[xref 9b] -d/--directory expects a directory: {target_root}", file=sys.stderr)
            return 2
    elif target_root.is_dir():
        print(f"[xref 9b] refusing implicit directory chat target: {target_root}", file=sys.stderr)
        print("[xref 9b] use -d/--directory for directory mode, or pass the full file path.", file=sys.stderr)
        return 2

    try:
        targets = triage.collect_targets(target_root, include_unknown=args.include_unknown, max_files=args.max_files)
    except FileNotFoundError:
        print(f"[xref 9b] target not found: {target_root}", file=sys.stderr)
        return 2
    if not targets:
        print(f"[xref 9b] no PE/ELF targets found under {target_root}", file=sys.stderr)
        return 2

    run_dir = Path(args.output_dir).expanduser().resolve() if args.output_dir else (ROOT / "runs" / f"agentre_chat_{now_stamp()}").resolve()
    run_dir.mkdir(parents=True, exist_ok=True)
    cfg = triage.RuntimeConfig(
        max_turns=args.max_turns,
        max_tool_calls=args.max_tool_calls,
        max_tokens=args.max_tokens,
        obs_limit=args.obs_limit,
        temperature=args.temperature,
        thinking=not args.no_thinking,
        save_transcripts=True,
        save_reasoning=False,
        ghidra_script_dir=Path(args.ghidra_script_dir).expanduser().resolve(),
    )

    if args.tool_smoke:
        return run_tool_smoke_targets(targets, run_dir, cfg)

    spec = Path(args.spec).read_text(encoding="utf-8")
    sessions: dict[int, dict[str, Any]] = {}
    active_idx = 0

    print(f"[xref 9b] chat target={target_root}", flush=True)
    print(f"[xref 9b] discovered={len(targets)} output={run_dir}", flush=True)
    print(f"[xref 9b] thinking={'on' if cfg.thinking else 'off'}", flush=True)
    if len(targets) > 1:
        print("[xref 9b] use /samples to list files and /open N to switch samples", flush=True)
    _, backend = build_backend(args)

    def ensure_session(index: int) -> dict[str, Any]:
        if index in sessions:
            return sessions[index]
        source_path = targets[index]
        sample_id = sample_id_for_index(index, len(targets))
        staged_path, fmt = stage_one(source_path, run_dir, sample_id)
        session = {
            "source_path": source_path,
            "sample_id": sample_id,
            "staged_path": staged_path,
            "fmt": fmt,
            "tools": triage.StaticTools(staged_path, fmt, cfg.obs_limit, cfg.ghidra_script_dir),
            "messages": build_inspect_messages(spec, sample_id, fmt, args.request),
            "state": {"sample_id": sample_id, "tool_calls_seen": 0, "last_final": None, "progress": not args.quiet},
            "started": False,
        }
        sessions[index] = session
        return session

    def save_session(index: int) -> None:
        session = ensure_session(index)
        save_inspect_session(
            run_dir,
            session["source_path"],
            session["staged_path"],
            session["fmt"],
            session["messages"],
            session["tools"],
            session["state"],
        )
        save_chat_index(run_dir, target_root, targets, index, sessions)

    def activate_sample(index: int) -> dict[str, Any]:
        session = ensure_session(index)
        print(
            f"[xref 9b] active {index + 1}/{len(targets)} {session['sample_id']} "
            f"format={session['fmt']} path={display_path(session['source_path'], target_root)}",
            flush=True,
        )
        if not session["started"]:
            reply = generate_visible_reply(backend, session["messages"], cfg, session["tools"], session["state"])
            session["started"] = True
            print("\nxref 9b> " + reply.strip() + "\n")
        save_session(index)
        return session

    current = activate_sample(active_idx)
    if args.one_shot:
        print(f"[xref 9b] transcript={run_dir / 'INSPECT.md'}")
        return 0

    print("Commands: /help, /samples, /open N, /current, /tools, /deep, /ghidra, /verdict, /compact, /continue, /save, /quit")
    while True:
        try:
            user_text = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not user_text:
            continue
        if user_text in {"/q", "/quit", "quit", "exit"}:
            break
        if user_text == "/help":
            print("Ask natural-language RE questions about the active sample.")
            print("Examples: `is this packed?`, `run Ghidra`, `why malicious?`, `what strings matter?`")
            print("Slash commands: /samples, /open N, /current, /tools, /deep, /ghidra, /verdict, /compact, /continue, /save, /quit.")
            continue
        if user_text == "/samples":
            print_sample_list(target_root, targets, active_idx, sessions)
            continue
        if user_text.startswith("/open"):
            selector = user_text[len("/open") :].strip()
            try:
                active_idx = resolve_sample_selector(selector, target_root, targets)
            except ValueError as exc:
                print(f"[xref 9b] {exc}")
                continue
            current = activate_sample(active_idx)
            continue
        if user_text == "/current":
            state = current["state"]
            print(
                f"{active_idx + 1}/{len(targets)} {current['sample_id']} {current['fmt']} "
                f"tools={state.get('tool_calls_seen', 0)} {display_path(current['source_path'], target_root)}"
            )
            continue
        if user_text == "/tools":
            used = current["tools"].used_tools
            print(", ".join(used) if used else "(no tools used yet)")
            continue
        if user_text == "/save":
            save_session(active_idx)
            print(f"[xref 9b] transcript={run_dir / 'INSPECT.md'}")
            continue
        if user_text == "/compact":
            before = len(current["messages"])
            snapshot = save_compaction_snapshot(run_dir, current)
            current["messages"] = compact_messages(current["messages"], current["tools"], current["state"])
            after = len(current["messages"])
            save_session(active_idx)
            print(f"[xref 9b] compacted conversation {before}->{after} messages; precompact={snapshot}")
            continue
        if user_text == "/continue":
            user_text = "Continue the previous answer from where it stopped. Do not repeat earlier content."
        elif user_text == "/ghidra":
            user_text = "Run Ghidra headless on this sample and explain what it shows."

        if user_text == "/deep":
            state = current["state"]
            current["messages"].append({
                "role": "user",
                "content": "Run the deeper static reverse-engineering workflow for this sample, then summarize the evidence.",
            })
            sequence = deep_tool_sequence(current["fmt"])
            progress(state, f"deep analysis: running {len(sequence)} static tools")
            for canonical, tool_args in sequence:
                append_direct_tool_result(canonical, tool_args, current["tools"], current["messages"], cfg, state)
            current["messages"].append({
                "role": "user",
                "content": (
                    "Using the deep static tool results just provided, give a concise analyst summary. "
                    "Classify only if the evidence supports it; otherwise say unknown and name the missing evidence. "
                    "Do not repeat raw tool output and do not run the same broad tools again unless one targeted follow-up is essential."
                ),
            })
            state["request_tool_calls_seen"] = 0
            try:
                reply = generate_visible_reply(backend, current["messages"], cfg, current["tools"], state)
            finally:
                state.pop("request_tool_calls_seen", None)
            print("\nxref 9b> " + reply.strip() + "\n")
            save_session(active_idx)
            continue

        if user_text == "/verdict":
            state = current["state"]
            current["messages"].append({
                "role": "user",
                "content": (
                    "Give a concise final verdict now using the evidence already gathered. "
                    "Call final_answer with classification, confidence, summary, and key evidence. "
                    "Do not call additional tools; if evidence is insufficient, classify as unknown and explain why."
                ),
            })
            state["request_tool_calls_seen"] = cfg.max_tool_calls
            try:
                reply = generate_visible_reply(backend, current["messages"], cfg, current["tools"], state)
            finally:
                state.pop("request_tool_calls_seen", None)
            print("\nxref 9b> " + reply.strip() + "\n")
            save_session(active_idx)
            continue

        current["messages"].append({"role": "user", "content": user_text})
        state = current["state"]
        state["request_tool_calls_seen"] = 0
        try:
            direct_tool = requested_direct_tool(user_text)
            if direct_tool:
                canonical, tool_args = direct_tool
                append_direct_tool_result(canonical, tool_args, current["tools"], current["messages"], cfg, state)
            reply = generate_visible_reply(backend, current["messages"], cfg, current["tools"], state)
        finally:
            state.pop("request_tool_calls_seen", None)
        print("\nxref 9b> " + reply.strip() + "\n")
        save_session(active_idx)

    save_session(active_idx)
    print(f"[xref 9b] transcript={run_dir / 'INSPECT.md'}")
    print(f"[xref 9b] chat index={run_dir / 'CHAT.md'}")
    return 0


def main(argv: list[str] | None = None) -> int:
    args = list(sys.argv[1:] if argv is None else argv)
    if not args or args[0] in {"-h", "--help", "help"}:
        print_top_help()
        return 0
    command, rest = args[0], args[1:]
    if command == "triage":
        script = ROOT / "agentre_triage.py"
        os.execv(sys.executable, [sys.executable, str(script)] + rest)
    if command == "chat":
        return run_chat(rest, "chat")
    if command == "inspect":
        return run_chat(rest, "inspect")
    print(f"Unknown command: {command}\n", file=sys.stderr)
    print_top_help()
    return 2


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