qwen.py

"""Qwen stream agent: pairwise vision comparison with optional class memory."""

from __future__ import annotations

import argparse
import base64
import io
import json
import os
import random
import re
import shutil
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from PIL import Image

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None  # type: ignore[misc, assignment]

from imagenet_hs import (
    LABEL_TO_SYNSET,
    SOURCE_CORR,
    SOURCE_INR,
    SOURCE_TINY,
    TASK_SPECS,
    SampleEntry,
    TaskSpec,
    build_task_entries,
    build_task_sequence,
)

REPO_ROOT = Path(__file__).resolve().parent
RESULTS_DIR = REPO_ROOT / "results"

# Dataset roots under repo data/
DATA_ROOT = REPO_ROOT / "data"
SOURCE_PATHS: dict[str, Path] = {
    SOURCE_TINY: DATA_ROOT / "tiny-imagenet-200",
    SOURCE_INR: DATA_ROOT / "imagenet-r",
    SOURCE_CORR: DATA_ROOT / "tiny-imagenet-200-corr",
}

SPLIT = "train"
IMAGE_SIZE = 256
QWEN_MODEL = "qwen3.6-flash"
SEED = 0

MIN_IMAGES = 4
MAX_IMAGES = 10

CHANGE_TYPE_TO_OPTION: dict[str, str | None] = {
    "initial": None,
    "new_class": "A",
    "domain_shift": "B",
    "corruption": "C",
}

OPTION_ALIASES: dict[str, str] = {
    "A": "A",
    "B": "B",
    "C": "C",
    "D": "D",
    "(A)": "A",
    "(B)": "B",
    "(C)": "C",
    "(D)": "D",
    "NEW CLASS": "A",
    "NEW CLASSES": "A",
    "NEW CATEGORIES": "A",
    "NEW_CATEGORY": "A",
    "DOMAIN SHIFT": "B",
    "DOMAIN_SHIFT": "B",
    "CORRUPTION": "C",
    "CORRUPTION / NOISE": "C",
    "DATA CORRUPTION": "C",
    "NOISE": "C",
    "NO SIGNIFICANT CHANGE": "D",
    "NO CHANGE": "D",
    "FALSE ALARM": "D",
}

SHIFT_ATTRIBUTION_PROMPT = """\
A distribution shift is detected in the online task-free continual data stream. \
Please serve as an expert in attribution analysis and identify the specific type of shift:
(a) new classes
(b) domain shift
(c) data corruption such as noise
(d) no significant change

Compare **only** the immediately previous and current data batches below (images only). \
You have no access to non-adjacent batches.

Decision order (first match wins):
1. (c) — Categories overlap the previous batch and/or your discovered-category memory, \
but the current batch is dominated by heavy noise, blur, or artifacts.
2. (b) — Categories overlap the previous batch and/or discovered-category memory, \
but the visual style or medium changed strongly (e.g. photo vs sketch, tattoo, line art).
3. (a) — The current batch introduces categories not seen in the previous batch \
and not already listed in discovered-category memory.
4. (d) — No meaningful shift between batches.

Output format:
Shift type: <a, b, c, or d>
Reasoning: <one short sentence>
New categories: <comma-separated category names newly introduced in the current batch vs previous and memory, or none>
"""

TASK_FREE_MEMORY_GUIDE = """\
Task-free memory: You may use the discovered categories and prior attributions below \
from earlier stream steps. Do not assume hidden grouping or non-adjacent batches. \
If the current batch reuses categories already discovered earlier and only style or \
noise changes, prefer (b) or (c) over (a) — even when the immediately previous batch \
looks different.
"""


@dataclass(frozen=True)
class RunConfig:
    model: str
    seed: int
    n_per_batch: int
    shuffle_tasks: bool

    @property
    def tag(self) -> str:
        model_safe = re.sub(r"[^\w.-]+", "_", self.model).strip("_")
        order = "shuffle" if self.shuffle_tasks else "ordered"
        return f"{model_safe}_seed{self.seed}_n{self.n_per_batch}_{order}"

    def report_path(self, results_dir: Path = RESULTS_DIR) -> Path:
        return results_dir / f"qwen_{self.tag}_stream.jsonl"

    def dryrun_dir(self, results_dir: Path = RESULTS_DIR) -> Path:
        return results_dir / f"qwen_{self.tag}_dryrun_batches"

    def as_metadata(self) -> dict[str, str | int | bool]:
        return {
            "model": self.model,
            "seed": self.seed,
            "n_per_batch": self.n_per_batch,
            "shuffle_tasks": self.shuffle_tasks,
            "run_tag": self.tag,
        }


@dataclass
class TaskBatchSample:
    task_id: int
    name: str
    images: list[Image.Image]
    categories: list[str]
    labels: list[int] = field(default_factory=list)


@dataclass
class DetectionRecord:
    from_task: int
    to_task: int
    predicted: str | None
    reasoning: str
    new_categories: list[str] = field(default_factory=list)


@dataclass
class TransitionResult:
    from_task: int
    to_task: int
    gt_change: str
    expected: str | None
    predicted: str | None
    raw_response: str
    correct: bool | None
    prev_batch_categories: list[str]
    curr_batch_categories: list[str]
    new_categories: list[str] = field(default_factory=list)
    reasoning: str = ""


def _normalize_shift_letter(raw: str) -> str | None:
    letter = raw.strip().upper().strip("()")
    if letter in ("A", "B", "C", "D"):
        return letter
    return None


def format_discovered_categories(categories: list[str]) -> str:
    if not categories:
        return ""
    return (
        "Discovered categories (new classes reported in earlier stream steps):\n"
        f"{', '.join(categories)}\n\n"
    )


def format_detection_history(records: list[DetectionRecord]) -> str:
    if not records:
        return ""
    lines = ["Prior attributions in this stream:"]
    for rec in records:
        pred = f"({rec.predicted.lower()})" if rec.predicted else "(?)"
        lines.append(
            f"- Step T{rec.from_task}→T{rec.to_task}: shift {pred} — {rec.reasoning}"
        )
        if rec.new_categories:
            lines.append(f"  New categories: {', '.join(rec.new_categories)}")
    return "\n".join(lines) + "\n\n"


def _append_discovered(discovered: list[str], names: list[str]) -> None:
    seen = {name.lower() for name in discovered}
    for name in names:
        cleaned = name.strip()
        key = cleaned.lower()
        if not cleaned or key in seen or key in ("none", "n/a", "na"):
            continue
        discovered.append(cleaned)
        seen.add(key)


def label_to_category(label: int) -> str:
    return LABEL_TO_SYNSET[label]


def resize_pil(img: Image.Image, size: int) -> Image.Image:
    if img.mode not in ("RGB", "RGBA"):
        img = img.convert("RGB")
    resample = getattr(Image, "Resampling", Image).LANCZOS
    return img.resize((size, size), resample=resample)


def _pick_entries_stratified(
    entries: list[SampleEntry],
    k: int,
    rng: random.Random,
) -> list[SampleEntry]:
    by_label: dict[int, list[SampleEntry]] = {}
    for e in entries:
        by_label.setdefault(e.label, []).append(e)

    label_order = list(by_label.keys())
    rng.shuffle(label_order)

    picked: list[SampleEntry] = []
    used: set[Path] = set()

    for lab in label_order:
        if len(picked) >= k:
            break
        e = rng.choice(by_label[lab])
        picked.append(e)
        used.add(e.path)

    if len(picked) < k:
        remaining = [e for e in entries if e.path not in used]
        rng.shuffle(remaining)
        for e in remaining:
            if len(picked) >= k:
                break
            picked.append(e)
            used.add(e.path)

    rng.shuffle(picked)
    return picked


def sample_task_batch(
    spec: TaskSpec,
    *,
    source_paths: dict[str, Path] | None,
    n_per_batch: int,
    seed: int,
    image_size: int,
    split: str = "train",
) -> TaskBatchSample:
    entries = build_task_entries(
        spec, source_paths=source_paths, splits=(split,)
    )[split]
    if not entries:
        raise RuntimeError(f"T{spec.task_id}: no images in {split}")

    rng = random.Random(seed + spec.task_id * 1000)
    k = min(n_per_batch, len(entries))
    picked = _pick_entries_stratified(entries, k, rng)

    images = [resize_pil(Image.open(e.path).convert("RGB"), image_size) for e in picked]
    labels = [e.label for e in picked]
    cats = [label_to_category(l) for l in labels]
    name = (
        f"T{spec.task_id} group={spec.group} source={spec.source} "
        f"change={spec.change_type} n={len(images)}"
    )
    return TaskBatchSample(
        task_id=spec.task_id,
        name=name,
        images=images,
        categories=cats,
        labels=labels,
    )


def build_client():
    if OpenAI is None:
        raise ImportError("Install openai: pip install openai")
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise EnvironmentError("Set DASHSCOPE_API_KEY before running.")
    return OpenAI(
        api_key=api_key,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )


def pil_to_base64(img: Image.Image, *, fmt: str = "PNG") -> str:
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    mime = "image/png" if fmt.upper() == "PNG" else "image/jpeg"
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    return f"data:{mime};base64,{b64}"


def images_to_content(images: list[Image.Image]) -> list[dict]:
    return [
        {"type": "image_url", "image_url": {"url": pil_to_base64(img)}}
        for img in images
    ]


def _parse_category_list(raw: str) -> list[str]:
    text = raw.strip()
    if not text:
        return []
    lower = text.lower()
    if lower in ("none", "n/a", "na", "(none)", "-", "no new categories", "no new category"):
        return []
    parts = re.split(r"[,;]", text)
    out: list[str] = []
    for p in parts:
        p = p.strip().strip(".")
        if p and p.lower() not in ("none", "n/a"):
            out.append(p)
    return out


def parse_detection_response(raw: str) -> tuple[str | None, str, list[str]]:
    text = raw.strip()
    cause = None
    for pattern in (
        r"SHIFT\s*TYPE\s*:\s*\(?([ABCDabcd])\)?\b",
        r"CHANGE\s*CAUSE\s*:\s*\(?([ABCDabcd])\)?\b",
    ):
        m = re.search(pattern, text, re.I)
        if m:
            cause = _normalize_shift_letter(m.group(1))
            break
    if cause is None:
        m2 = re.search(r"\(([ABCDabcd])\)", text)
        if m2:
            cause = _normalize_shift_letter(m2.group(1))
    if cause is None:
        m3 = re.search(r"\b([ABCD])\b", text.upper())
        if m3:
            cause = m3.group(1).upper()

    reasoning = ""
    rm = re.search(
        r"REASONING\s*:\s*(.+?)(?=\n\s*NEW\s*CATEGORIES\s*:|$)",
        text,
        re.I | re.S,
    )
    if rm:
        reasoning = rm.group(1).strip()

    new_categories: list[str] = []
    nm = re.search(r"NEW\s*CATEGORIES\s*:\s*(.+?)(?:\n\n|\Z)", text, re.I | re.S)
    if nm:
        new_categories = _parse_category_list(nm.group(1))

    return cause, reasoning, new_categories


def parse_detection(raw: str) -> str | None:
    text = raw.strip()
    for pattern in (
        r"SHIFT\s*TYPE\s*:\s*\(?([ABCDabcd])\)?\b",
        r"CHANGE\s*CAUSE\s*:\s*\(?([ABCDabcd])\)?\b",
    ):
        m = re.search(pattern, text, re.I)
        if m:
            return _normalize_shift_letter(m.group(1))
    m = re.search(r"\(([ABCDabcd])\)", text)
    if m:
        return _normalize_shift_letter(m.group(1))
    m = re.search(r"\b([ABCD])\b", text.upper())
    if m:
        return m.group(1).upper()
    upper = text.upper()
    for phrase, letter in OPTION_ALIASES.items():
        if phrase in upper and len(phrase) > 1:
            return letter
    return None


def build_pair_message(
    prev: TaskBatchSample,
    curr: TaskBatchSample,
    *,
    discovered_categories: list[str],
    detection_history: list[DetectionRecord],
) -> list[dict]:
    memory_block = format_discovered_categories(discovered_categories)
    history_block = format_detection_history(detection_history)
    preamble = (
        f"{memory_block}{history_block}{TASK_FREE_MEMORY_GUIDE}\n\n{SHIFT_ATTRIBUTION_PROMPT}"
    )

    content: list[dict] = [
        {"type": "text", "text": preamble},
        {"type": "text", "text": "<Previous data batch>:\n"},
        *images_to_content(prev.images),
        {"type": "text", "text": "\n<Current data batch>:\n"},
        *images_to_content(curr.images),
    ]
    return [{"role": "user", "content": content}]


def detect_pair(
    client,
    prev: TaskBatchSample,
    curr: TaskBatchSample,
    *,
    discovered_categories: list[str],
    detection_history: list[DetectionRecord],
    model: str,
    temperature: float,
) -> str:
    messages = build_pair_message(
        prev,
        curr,
        discovered_categories=discovered_categories,
        detection_history=detection_history,
    )
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
    )
    return completion.choices[0].message.content or ""


def expected_option(spec: TaskSpec) -> str | None:
    if spec.change_type in ("initial", "new_class"):
        return "A"
    return CHANGE_TYPE_TO_OPTION.get(spec.change_type)


@dataclass
class StreamAgent:
    """Pairwise Qwen agent with optional shuffled stream and class memory."""

    n_per_batch: int = 8
    seed: int = SEED
    shuffle_tasks: bool = False
    model: str = QWEN_MODEL
    run_config: RunConfig | None = None
    client: object | None = None

    def __post_init__(self) -> None:
        if self.run_config is None:
            self.run_config = RunConfig(
                model=self.model,
                seed=self.seed,
                n_per_batch=self.n_per_batch,
                shuffle_tasks=self.shuffle_tasks,
            )

    def _client(self):
        if self.client is None:
            self.client = build_client()
        return self.client

    def _task_sequence(self) -> list[TaskSpec]:
        return build_task_sequence(shuffle=self.shuffle_tasks, seed=self.seed)

    def run_pair(
        self,
        prev_spec: TaskSpec,
        curr_spec: TaskSpec,
        prev_batch: TaskBatchSample,
        curr_batch: TaskBatchSample,
        *,
        discovered_categories: list[str] | None = None,
        detection_history: list[DetectionRecord] | None = None,
        call_api: bool = True,
    ) -> TransitionResult:
        expected = expected_option(curr_spec)
        raw = ""
        predicted: str | None = None
        reasoning = ""
        new_categories: list[str] = []

        if call_api:
            raw = detect_pair(
                self._client(),
                prev_batch,
                curr_batch,
                discovered_categories=list(discovered_categories or []),
                detection_history=list(detection_history or []),
                model=self.model,
                temperature=0.0,
            )
            predicted, reasoning, new_categories = parse_detection_response(raw)
            if predicted is None:
                predicted = parse_detection(raw)

        correct: bool | None = None
        if expected is not None and predicted is not None:
            correct = predicted == expected

        return TransitionResult(
            from_task=prev_spec.task_id,
            to_task=curr_spec.task_id,
            gt_change=curr_spec.change_type,
            expected=expected,
            predicted=predicted,
            raw_response=raw,
            correct=correct,
            prev_batch_categories=prev_batch.categories,
            curr_batch_categories=curr_batch.categories,
            new_categories=new_categories,
            reasoning=reasoning,
        )

    def run_stream(
        self,
        *,
        call_api: bool = True,
        report_path: Path | None = None,
        dryrun_batch_dir: Path | None = None,
    ) -> list[TransitionResult]:
        specs = self._task_sequence()
        results: list[TransitionResult] = []
        prev_spec: TaskSpec | None = None
        prev_batch: TaskBatchSample | None = None
        discovered_categories: list[str] = []
        detection_history: list[DetectionRecord] = []

        n_img = self.n_per_batch
        run = self.run_config
        shuffle_note = ", interleaved task order" if self.shuffle_tasks else ""
        print(
            f"Stream agent: model={self.model}, run_tag={run.tag if run else ''}, "
            f"{len(specs)} tasks, n={n_img} per task "
            f"({n_img * 2} images per adjacent pairwise call), "
            f"task-free discovered-category memory{shuffle_note}"
        )
        if self.shuffle_tasks:
            order = " → ".join(f"T{s.task_id}:{s.group}" for s in specs)
            print(f"  task order: {order}")

        for spec in specs:
            batch = sample_task_batch(
                spec,
                source_paths=SOURCE_PATHS,
                n_per_batch=self.n_per_batch,
                seed=self.seed,
                image_size=IMAGE_SIZE,
                split=SPLIT,
            )
            if dryrun_batch_dir is not None:
                _save_task_batch_images(spec=spec, batch=batch, output_root=dryrun_batch_dir)
            print(f"\n--- T{spec.task_id} {spec.change_type} ({spec.description}) ---")
            print(
                f"  sampled labels: {batch.labels} "
                f"synsets: {batch.categories[:3]}..."
            )

            if prev_spec is not None and prev_batch is not None:
                tr = self.run_pair(
                    prev_spec,
                    spec,
                    prev_batch,
                    batch,
                    discovered_categories=discovered_categories,
                    detection_history=detection_history,
                    call_api=call_api,
                )
                results.append(tr)
                mark = "✓" if tr.correct else ("✗" if tr.correct is False else "?")
                print(
                    f"  pair T{tr.from_task}→T{tr.to_task} "
                    f"gt={tr.gt_change} expected={tr.expected} "
                    f"predicted={tr.predicted} {mark}"
                )
                if tr.new_categories:
                    print(f"  new_categories: {tr.new_categories}")
                if tr.reasoning:
                    print(f"  reasoning: {tr.reasoning}")
                self._append_report(tr, report_path)
                detection_history.append(
                    DetectionRecord(
                        from_task=tr.from_task,
                        to_task=tr.to_task,
                        predicted=tr.predicted,
                        reasoning=tr.reasoning,
                        new_categories=tr.new_categories,
                    )
                )
                _append_discovered(discovered_categories, tr.new_categories)
            else:
                print("  (stream start — no previous batch)")

            prev_spec = spec
            prev_batch = batch

        if results:
            n_ok = sum(1 for r in results if r.correct is True)
            n_bad = sum(1 for r in results if r.correct is False)
            n_unk = sum(1 for r in results if r.correct is None)
            print(
                f"\nVerification: {n_ok} correct, {n_bad} wrong, {n_unk} unparsed "
                f"(of {len(results)} pairs)"
            )
        return results

    def _append_report(self, tr: TransitionResult, report_path: Path | None) -> None:
        if report_path is None:
            return
        report_path.parent.mkdir(parents=True, exist_ok=True)
        meta = self.run_config.as_metadata() if self.run_config else {}
        row = {**asdict(tr), **meta, "ts": datetime.now(timezone.utc).isoformat()}
        with report_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def _sources_for_specs(specs: list[TaskSpec]) -> set[str]:
    return {s.source for s in specs}


def check_source_paths(
    paths: dict[str, Path],
    *,
    required_sources: set[str] | None = None,
) -> None:
    """Ensure dataset roots exist (all three by default)."""
    needed = required_sources or {SOURCE_TINY, SOURCE_INR, SOURCE_CORR}
    missing = [paths[src] for src in sorted(needed) if not paths[src].is_dir()]
    if missing:
        raise FileNotFoundError(
            "Missing ImageNet-HS data directories:\n"
            + "\n".join(f"  - {p}" for p in missing)
            + f"\nExpected under DATA_ROOT={DATA_ROOT}"
        )


def _save_task_batch_images(
    *,
    spec: TaskSpec,
    batch: TaskBatchSample,
    output_root: Path,
) -> None:
    """Save sampled task images and metadata for plotting."""
    task_dir = output_root / f"T{spec.task_id:02d}_{spec.group}_{spec.source}_{spec.change_type}"
    if task_dir.exists():
        shutil.rmtree(task_dir)
    task_dir.mkdir(parents=True, exist_ok=True)

    rows: list[dict[str, str | int]] = []
    for idx, (img, label, synset) in enumerate(zip(batch.images, batch.labels, batch.categories)):
        file_name = f"{idx:02d}_label{label:02d}_{synset}.png"
        img.save(task_dir / file_name, format="PNG")
        rows.append({
            "index": idx,
            "file": file_name,
            "label": label,
            "synset": synset,
        })

    meta = {
        "task_id": spec.task_id,
        "group": spec.group,
        "source": spec.source,
        "change_type": spec.change_type,
        "description": spec.description,
        "num_images": len(batch.images),
        "rows": rows,
    }
    (task_dir / "meta.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=(
            "Qwen pairwise stream: compare consecutive task batches with class memory."
        )
    )
    p.add_argument(
        "-n",
        "--num-images",
        dest="n_per_batch",
        type=int,
        default=8,
        help=f"Images per task; pairwise call sends 2×n ({MIN_IMAGES}–{MAX_IMAGES} each).",
    )
    p.add_argument(
        "--dry-run",
        action="store_true",
        help="Sample batches only; skip Qwen API.",
    )
    p.add_argument(
        "--pair",
        type=str,
        default=None,
        help="Run one pairwise step at stream task id N (uses T(N-1) and TN), e.g. 4 or 3-4.",
    )
    p.add_argument(
        "--shuffle-tasks",
        action="store_true",
        help=(
            "Randomly interleave tasks across groups; within each group the Tiny "
            "source batch (new classes) always precedes domain shift / corruption."
        ),
    )
    p.add_argument(
        "--model",
        type=str,
        default=QWEN_MODEL,
        help=f"DashScope vision model name (default: {QWEN_MODEL}).",
    )
    p.add_argument(
        "--seed",
        type=int,
        default=SEED,
        help="Random seed for image sampling and task-order shuffle.",
    )
    return p.parse_args()


def _spec_by_id(task_id: int, specs: list[TaskSpec] | None = None) -> TaskSpec:
    seq = specs or list(TASK_SPECS)
    for spec in seq:
        if spec.task_id == task_id:
            return spec
    raise ValueError(f"Unknown task id {task_id}")


def main() -> None:
    args = parse_args()
    if not MIN_IMAGES <= args.n_per_batch <= MAX_IMAGES:
        raise SystemExit(f"-n must be in [{MIN_IMAGES}, {MAX_IMAGES}]")

    agent = StreamAgent(
        n_per_batch=args.n_per_batch,
        seed=args.seed,
        shuffle_tasks=args.shuffle_tasks,
        model=args.model,
    )
    run = agent.run_config
    assert run is not None
    report_path = run.report_path()
    dryrun_dir = run.dryrun_dir()
    specs = agent._task_sequence()

    if args.pair:
        m = re.match(r"(\d+)\s*[-–]\s*(\d+)", args.pair.strip())
        if m:
            to_id = int(m.group(2))
        elif args.pair.strip().isdigit():
            to_id = int(args.pair.strip())
        else:
            raise SystemExit("--pair format: 4 or 3-4 (detect at arrival of T4)")
        if to_id <= 1:
            raise SystemExit("T1 has no previous task; use --pair with id >= 2")
        to_spec = _spec_by_id(to_id, specs)
        prev_spec = _spec_by_id(to_id - 1, specs)
        check_source_paths(
            SOURCE_PATHS,
            required_sources=_sources_for_specs([prev_spec, to_spec]),
        )
        prev_batch = sample_task_batch(
            prev_spec,
            source_paths=SOURCE_PATHS,
            n_per_batch=agent.n_per_batch,
            seed=agent.seed,
            image_size=IMAGE_SIZE,
            split=SPLIT,
        )
        curr_batch = sample_task_batch(
            to_spec,
            source_paths=SOURCE_PATHS,
            n_per_batch=agent.n_per_batch,
            seed=agent.seed,
            image_size=IMAGE_SIZE,
            split=SPLIT,
        )
        tr = agent.run_pair(
            prev_spec,
            to_spec,
            prev_batch,
            curr_batch,
            discovered_categories=[],
            detection_history=[],
            call_api=not args.dry_run,
        )
        if args.dry_run:
            dryrun_dir.mkdir(parents=True, exist_ok=True)
            _save_task_batch_images(spec=prev_spec, batch=prev_batch, output_root=dryrun_dir)
            _save_task_batch_images(spec=to_spec, batch=curr_batch, output_root=dryrun_dir)
            print(f"Saved dry-run batches to: {dryrun_dir}")
        print(json.dumps(asdict(tr), indent=2, ensure_ascii=False))
        if not args.dry_run:
            agent._append_report(tr, report_path)
            print(f"Appended result to: {report_path}")
        return

    check_source_paths(SOURCE_PATHS, required_sources=_sources_for_specs(specs))

    if args.dry_run:
        dryrun_dir.mkdir(parents=True, exist_ok=True)
        print(f"Dry-run output: {dryrun_dir}")

    agent.run_stream(
        call_api=not args.dry_run,
        report_path=None if args.dry_run else report_path,
        dryrun_batch_dir=dryrun_dir if args.dry_run else None,
    )
    if args.dry_run:
        print(f"Saved dry-run batches to: {dryrun_dir}")
    else:
        print(f"Results saved to: {report_path}")


if __name__ == "__main__":
    main()