from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
import urllib.request
from dataclasses import dataclass
from datetime import date, datetime, timezone, timedelta
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
RUNS_DIR = ROOT / "runs"
SOURCE_RAW = RUNS_DIR / "run-2026-06-04-mcp-run301-llm-scale10x-v4-pro-prose-natural-followup-800tok.raw_result.json"
RUN_ID = "run-2026-06-05-honcho-run301-retrieval-baseline"
HONCHO_CONFIG = Path.home() / ".honcho" / "config.json"
ANTHROPIC_ENV = ROOT / ".env.anthropic"
ANTHROPIC_BASE_URL = os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com")
DEFAULT_SONNET_MODEL = os.environ.get("ANTHROPIC_MODEL", "claude-sonnet-4-6")

START_DATE = date(2026, 1, 5)
COMPANY_COUNT = 80
TIMELINE_DAYS = 120
ROLES = ("account owner", "technical owner", "finance reviewer")
TOPICS = (
    "contract review",
    "data migration",
    "billing reconciliation",
    "API rollout",
    "support escalation",
    "weekly steering",
)
SEGMENTS = (
    "logistics",
    "healthcare",
    "retail",
    "payments",
    "biotech",
    "construction",
    "insurance",
    "energy",
    "education",
    "manufacturing",
)
MILESTONE_DAYS = (1, 30, 60, 90, 120)
LIMIT = 8
SONNET_PRICING_USD_PER_M = {
    "claude-sonnet-4-6": {"input": 3.0, "output": 15.0},
}


@dataclass(frozen=True)
class BaselineRecord:
    record_id: str
    tier: str
    text: str
    payload: dict[str, Any]


def iso_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf8"))


def load_env(path: Path) -> None:
    if not path.exists():
        return
    for line in path.read_text(encoding="utf8").splitlines():
        if not line or line.strip().startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        os.environ.setdefault(key.strip(), value.strip())


def write_json(path: Path, value: Any) -> None:
    path.write_text(json.dumps(value, indent=2, sort_keys=True, default=str) + "\n", encoding="utf8")


def slug(index: int) -> str:
    return f"scale10-chrono-company-{index:03d}"


def display(index: int) -> str:
    return f"Chrono Company {index:03d}"


def segment(index: int) -> str:
    return SEGMENTS[(index - 1) % len(SEGMENTS)]


def topic(day_index: int) -> str:
    return TOPICS[(day_index - 1) % len(TOPICS)]


def day_date(day_index: int) -> str:
    return (START_DATE + timedelta(days=day_index - 1)).isoformat()


def marker(index: int, day_index: int) -> str:
    return f"scale10chronoscale10chronocompany{index:03d}day{day_index:03d}"


def record(record_id: str, tier: str, payload: dict[str, Any]) -> BaselineRecord:
    text = json.dumps({"id": record_id, "tier": tier, "body": payload}, sort_keys=True, default=str)
    return BaselineRecord(record_id=record_id, tier=tier, text=text, payload=payload)


def build_records(*, company_count: int, timeline_days: int, company_limit: int | None = None) -> list[BaselineRecord]:
    records: list[BaselineRecord] = []
    max_company = min(company_limit or company_count, company_count)
    for index in range(1, max_company + 1):
        company_slug = slug(index)
        company_display = display(index)
        records.append(
            record(
                f"run278-{company_slug}",
                "entity",
                {
                    "category": "company",
                    "display_name": company_display,
                    "segment": segment(index),
                    "timeline_start": START_DATE.isoformat(),
                    "timeline_days": timeline_days,
                },
            )
        )
        for role_index, role in enumerate(ROLES, start=1):
            person_id = f"run278-person-{index:03d}-{role_index:02d}"
            records.append(
                record(
                    person_id,
                    "entity",
                    {
                        "category": "person",
                        "display_name": f"{company_display} Contact {role_index}",
                        "role": role,
                        "company": f"run278-{company_slug}",
                    },
                )
            )
            records.append(
                record(
                    f"run278-rel-{index:03d}-{role_index:02d}",
                    "entity",
                    {
                        "category": "relationship",
                        "source": person_id,
                        "target": f"run278-{company_slug}",
                        "type": "works_on_account",
                        "description": f"{person_id} is the {role} for {company_display}.",
                    },
                )
            )
        for day_index in range(1, timeline_days + 1):
            event_marker = marker(index, day_index)
            records.append(
                record(
                    f"run278-{company_slug}-day-{day_index:03d}",
                    "journal",
                    {
                        "kind": "scale10_chronology_daily_update",
                        "company": company_display,
                        "company_slug": company_slug,
                        "day_index": day_index,
                        "date": day_date(day_index),
                        "topic": topic(day_index),
                        "marker": event_marker,
                        "summary": f"{company_display} day {day_index:03d}: {topic(day_index)} update recorded.",
                    },
                )
            )
            if day_index in MILESTONE_DAYS:
                records.append(
                    record(
                        f"run278-{company_slug}-day-{day_index:03d}-milestone",
                        "entity",
                        {
                            "category": "timeline_milestone",
                            "company": f"run278-{company_slug}",
                            "company_display": company_display,
                            "day_index": day_index,
                            "date": day_date(day_index),
                            "topic": topic(day_index),
                            "marker": event_marker,
                            "decision": f"{company_display} checkpoint day {day_index:03d} accepted for {topic(day_index)}.",
                        },
                    )
                )
            if day_index % 10 == 0 or day_index == timeline_days:
                records.append(
                    record(
                        f"run278-current-status-{company_slug}",
                        "state",
                        {
                            "company": company_display,
                            "last_day_index": day_index,
                            "last_date": day_date(day_index),
                            "last_marker": event_marker,
                            "status": "active",
                        },
                    )
                )
    return records


def approx_tokens(value: Any) -> int:
    return max(1, len(json.dumps(value, sort_keys=True, default=str)) // 4)


def contains_all(context: list[dict[str, Any]], expected: list[str]) -> bool:
    if not expected:
        return True
    text = json.dumps(context, sort_keys=True, default=str).lower()
    return all(item.lower() in text for item in expected)


def contains_any(context: list[dict[str, Any]], forbidden: list[str]) -> bool:
    if not forbidden:
        return False
    text = json.dumps(context, sort_keys=True, default=str).lower()
    return any(item.lower() in text for item in forbidden)


def local_keyword_search(records: list[BaselineRecord], query: str, limit: int) -> list[dict[str, Any]]:
    terms = [part.lower() for part in query.split() if part.strip()]
    scored: list[tuple[int, BaselineRecord]] = []
    for item in records:
        text = item.text.lower()
        score = sum(text.count(term) for term in terms)
        if score:
            scored.append((score, item))
    scored.sort(key=lambda row: (-row[0], row[1].record_id))
    return [
        {"tier": item.tier, "key": item.record_id, "body": item.payload, "score": score}
        for score, item in scored[:limit]
    ]


def load_honcho() -> Any:
    try:
        from honcho import Honcho  # type: ignore
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "Honcho SDK missing. Install it in your active environment, then rerun with --execute. "
            "Expected import: from honcho import Honcho"
        ) from exc
    return Honcho


def serialize_honcho_result(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        return value
    for attr in ("model_dump", "dict"):
        method = getattr(value, attr, None)
        if callable(method):
            try:
                dumped = method()
                if isinstance(dumped, dict):
                    return dumped
            except TypeError:
                pass
    data = {"repr": repr(value)}
    for attr in ("id", "content", "text", "message", "metadata"):
        if hasattr(value, attr):
            try:
                data[attr] = getattr(value, attr)
            except Exception:
                pass
    return data


def honcho_client(workspace_id: str | None) -> Any:
    Honcho = load_honcho()
    file_config: dict[str, Any] = {}
    if HONCHO_CONFIG.exists():
        try:
            loaded = json.loads(HONCHO_CONFIG.read_text(encoding="utf8"))
            file_config = loaded if isinstance(loaded, dict) else {}
        except json.JSONDecodeError:
            file_config = {}
    kwargs: dict[str, Any] = {}
    if workspace_id:
        kwargs["workspace_id"] = workspace_id
    base_url = os.environ.get("HONCHO_BASE_URL") or file_config.get("environmentUrl")
    if base_url:
        kwargs["base_url"] = base_url
    api_key = os.environ.get("HONCHO_API_KEY") or file_config.get("apiKey")
    if api_key:
        kwargs["api_key"] = api_key
    return Honcho(**kwargs)


def ingest_honcho(records: list[BaselineRecord], *, workspace_id: str | None, session_id: str, peer_id: str, batch_size: int) -> None:
    honcho = honcho_client(workspace_id)
    session = honcho.session(session_id)
    peer = honcho.peer(peer_id)
    total_batches = (len(records) + batch_size - 1) // batch_size
    for start in range(0, len(records), batch_size):
        batch = records[start : start + batch_size]
        batch_index = start // batch_size + 1
        print(
            f"[ingest] batch {batch_index}/{total_batches} records {start + 1}-{start + len(batch)}",
            file=sys.stderr,
            flush=True,
        )
        messages = [
            peer.message(
                item.text,
                metadata={
                    "source_run": "run301",
                    "record_id": item.record_id,
                    "tier": item.tier,
                },
            )
            for item in batch
        ]
        session.add_messages(messages)


def search_honcho(*, workspace_id: str | None, session_id: str, query: str, limit: int) -> list[dict[str, Any]]:
    honcho = honcho_client(workspace_id)
    session = honcho.session(session_id)
    try:
        results = session.search(query, limit=limit)
    except TypeError:
        results = session.search(query)
    return [serialize_honcho_result(item) for item in list(results)[:limit]]


def row_for_question(question: dict[str, Any], honcho_context: list[dict[str, Any]]) -> dict[str, Any]:
    expected = [str(item) for item in question.get("expected_contains", [])]
    forbidden = [str(item) for item in question.get("forbidden_contains", [])]
    expected_unknown = bool(question.get("expected_unknown"))
    if expected_unknown:
        passed = not contains_any(honcho_context, forbidden)
    else:
        passed = contains_all(honcho_context, expected)
    return {
        "id": question["id"],
        "query": question["query"],
        "retrieval_query": question["retrieval_query"],
        "expected_contains": expected,
        "expected_unknown": expected_unknown,
        "honcho_context_count": len(honcho_context),
        "honcho_context_tokens_estimate": approx_tokens(honcho_context),
        "honcho_retrieval_passed": passed,
        "honcho_context": honcho_context,
        "sibyl_context_count": int(question.get("context_count") or len(question.get("context") or [])),
        "sibyl_context_tokens_estimate": approx_tokens(question.get("context") or []),
        "sibyl_retrieval_passed": bool(question.get("score", {}).get("retrieval_has_expected")),
        "sibyl_final_answer_passed": bool(question.get("score", {}).get("passed")),
    }


def build_prompt(question: str, context: list[dict[str, Any]]) -> tuple[str, str]:
    system = (
        "You answer strictly from the provided memory context. "
        "Answer in plain English in one or two short sentences. "
        "Include the exact company name, dates, status, topic, segment, role, person, or marker when relevant. "
        "If the answer is not supported by the context, say exactly: "
        "I don't have that information in the provided memory."
    )
    user_content = json.dumps({"question": question, "context": context}, sort_keys=True, default=str)
    return system, user_content


def parse_model_content(content: str, usage: dict[str, Any]) -> dict[str, Any]:
    unsupported = "i don't have that information" in content.lower()
    return {"response": {"answer": content, "unknown": unsupported}, "usage": usage}


def estimate_sonnet_cost(model: str, usage: dict[str, Any]) -> dict[str, Any]:
    pricing = SONNET_PRICING_USD_PER_M.get(model, SONNET_PRICING_USD_PER_M["claude-sonnet-4-6"])
    input_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)
    output_tokens = int(usage.get("output_tokens") or usage.get("completion_tokens") or 0)
    return {
        "prompt_tokens": input_tokens,
        "completion_tokens": output_tokens,
        "pricing_usd_per_m": pricing,
        "estimated_input_usd": input_tokens / 1_000_000 * pricing["input"],
        "estimated_output_usd": output_tokens / 1_000_000 * pricing["output"],
        "estimated_total_usd": input_tokens / 1_000_000 * pricing["input"] + output_tokens / 1_000_000 * pricing["output"],
    }


def ask_sonnet(question: str, context: list[dict[str, Any]], *, model: str, max_tokens: int) -> dict[str, Any]:
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        raise SystemExit("ANTHROPIC_API_KEY missing. Add it to .env.anthropic or export it before --answer-llm.")
    system, user_content = build_prompt(question, context)
    payload = {
        "model": model,
        "system": system,
        "messages": [{"role": "user", "content": user_content}],
        "temperature": 0,
        "max_tokens": max_tokens,
    }
    request = urllib.request.Request(
        f"{ANTHROPIC_BASE_URL.rstrip('/')}/v1/messages",
        data=json.dumps(payload).encode("utf8"),
        headers={
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=120) as response:
        raw = json.loads(response.read().decode("utf8"))
    content = "".join(part.get("text", "") for part in raw.get("content", []) if isinstance(part, dict))
    parsed = parse_model_content(content, raw.get("usage", {}))
    parsed["cost"] = estimate_sonnet_cost(model, parsed["usage"])
    return parsed


def build_natural_followup_question(original_question: str, previous_answer: dict[str, Any]) -> str:
    previous = json.dumps(previous_answer, sort_keys=True, default=str)
    return (
        "Can you be more precise for the same question? "
        "Use only the provided memory context. "
        "If present in memory, include the exact company name, day number, ISO date, status, topic, segment, role, "
        "person, or marker that answers the question. "
        f"Original question: {original_question}\n"
        f"Previous answer: {previous}"
    )


def score_answer(question: dict[str, Any], answer: dict[str, Any], context: list[dict[str, Any]]) -> dict[str, Any]:
    combined = json.dumps(answer, sort_keys=True, default=str).lower()
    context_text = json.dumps(context, sort_keys=True, default=str).lower()
    expected_contains = [str(item).lower() for item in question.get("expected_contains", [])]
    forbidden_contains = [str(item).lower() for item in question.get("forbidden_contains", [])]
    contains_ok = all(item in combined for item in expected_contains)
    forbidden_ok = not any(item in combined for item in forbidden_contains)
    unknown_ok = True
    if question.get("expected_unknown") is True:
        unknown_ok = any(
            marker in combined
            for marker in (
                "i don't have that information",
                "not supported",
                "provided memory",
                "unknown",
                "no information",
            )
        )
    retrieval_has_expected = all(item in context_text for item in expected_contains) if expected_contains else True
    return {
        "passed": contains_ok and forbidden_ok and unknown_ok,
        "contains_ok": contains_ok,
        "forbidden_ok": forbidden_ok,
        "unknown_ok": unknown_ok,
        "retrieval_has_expected": retrieval_has_expected,
    }


def add_cost(total: dict[str, Any], cost: dict[str, Any]) -> None:
    for key in ("prompt_tokens", "completion_tokens", "estimated_input_usd", "estimated_output_usd", "estimated_total_usd"):
        total[key] += cost.get(key, 0)


def answer_with_sonnet(question: dict[str, Any], context: list[dict[str, Any]], *, model: str, max_tokens: int, followup: bool) -> dict[str, Any]:
    initial = ask_sonnet(question["query"], context, model=model, max_tokens=max_tokens)
    initial_score = score_answer(question, initial["response"], context)
    final_answer = initial["response"]
    final_score = dict(initial_score)
    followup_result: dict[str, Any] | None = None
    total_cost = {
        "prompt_tokens": 0,
        "completion_tokens": 0,
        "estimated_input_usd": 0.0,
        "estimated_output_usd": 0.0,
        "estimated_total_usd": 0.0,
    }
    add_cost(total_cost, initial["cost"])
    if followup and not initial_score["passed"] and initial_score["retrieval_has_expected"] and question.get("expected_unknown") is not True:
        followup_query = build_natural_followup_question(question["query"], initial["response"])
        second = ask_sonnet(followup_query, context, model=model, max_tokens=max_tokens)
        final_answer = {"initial": initial["response"], "followup": second["response"]}
        final_score = score_answer(question, final_answer, context)
        add_cost(total_cost, second["cost"])
        followup_result = {
            "query": followup_query,
            "answer": second["response"],
            "usage": second["usage"],
            "cost": second["cost"],
            "score_after_followup": final_score,
        }
    return {
        "initial_answer": initial["response"],
        "initial_usage": initial["usage"],
        "initial_cost": initial["cost"],
        "initial_score": initial_score,
        "answer": final_answer,
        "score": final_score,
        "followup": followup_result,
        "cost": total_cost,
    }


def build_markdown(raw: dict[str, Any]) -> str:
    summary = raw["summary"]
    answer_cols = ""
    answer_header = ""
    answer_separator = ""
    answer_summary = ""
    if raw.get("answer_mode") == "sonnet":
        answer_header = " | Honcho + Sonnet"
        answer_separator = "|---:"
        answer_cols = " | {answer}"
        answer_summary = f"""| Honcho + Sonnet answered | {summary["honcho_sonnet_answer_passed"]} / {summary["total"]} |
| Honcho + Sonnet follow-ups | {summary["honcho_sonnet_followups"]} |
| Honcho + Sonnet prompt tokens | {summary["honcho_sonnet_cost"]["prompt_tokens"]} |
| Honcho + Sonnet completion tokens | {summary["honcho_sonnet_cost"]["completion_tokens"]} |
| Honcho + Sonnet estimated cost USD | {summary["honcho_sonnet_cost"]["estimated_total_usd"]:.6f} |
"""
    rows = "\n".join(
        f"| {item['id']} | {item['sibyl_retrieval_passed']} | {item['honcho_retrieval_passed']} | "
        f"{item['sibyl_context_tokens_estimate']} | {item['honcho_context_tokens_estimate']}"
        f"{answer_cols.format(answer=item.get('honcho_sonnet', {}).get('score', {}).get('passed', 'n/a'))} |"
        for item in raw["questions"]
    )
    failed = [
        {
            "id": item["id"],
            "retrieval_query": item["retrieval_query"],
            "expected_contains": item["expected_contains"],
            "honcho_context_sample": item["honcho_context"][:2],
        }
        for item in raw["questions"]
        if not item["honcho_retrieval_passed"]
    ]
    return f"""# Honcho Retrieval Baseline

Run: `{raw["run_id"]}`

## Scope

- Source baseline: `{raw["source_run_id"]}`
- Dataset: {raw["dataset"]["record_count"]} generated memory records, {raw["dataset"]["company_count"]} companies, {raw["dataset"]["stakeholder_count"]} stakeholders, {raw["dataset"]["timeline_days"]} simulated days.
- Questions: {summary["total"]} from the existing run301 artifact.
- Metric: retrieval contained the expected facts, before any LLM answer generation.
- Honcho mode: `{raw["honcho_mode"]}`.

## Summary

| Metric | Value |
|---|---:|
| Sibyl retrieval passed | {summary["sibyl_retrieval_passed"]} / {summary["total"]} |
| Honcho retrieval passed | {summary["honcho_retrieval_passed"]} / {summary["total"]} |
| Avg Sibyl context tokens | {summary["avg_sibyl_context_tokens"]:.2f} |
| Avg Honcho context tokens | {summary["avg_honcho_context_tokens"]:.2f} |
{answer_summary}

## Questions

| Question | Sibyl retrieval | Honcho retrieval | Sibyl ctx tokens | Honcho ctx tokens{answer_header} |
|---|---:|---:|---:|---:{answer_separator}|
{rows}

## Honcho Failures

```json
{json.dumps(failed[:10], indent=2, sort_keys=True, default=str)}
```
"""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run an isolated Honcho retrieval baseline for the existing run301 benchmark.")
    parser.add_argument("--source-raw", default=str(SOURCE_RAW), help="Existing run301 raw result JSON.")
    parser.add_argument("--execute", action="store_true", help="Use the Honcho SDK. Without this, only a local dry-run baseline is produced.")
    parser.add_argument("--write-dry-run", action="store_true", help="Write dry-run local keyword artifacts. Do not use these as Honcho results.")
    parser.add_argument("--ingest", action="store_true", help="Ingest the generated run301 records into Honcho before querying.")
    parser.add_argument("--workspace-id", default=os.environ.get("HONCHO_WORKSPACE_ID"), help="Honcho workspace id.")
    parser.add_argument("--session-id", default=os.environ.get("HONCHO_SESSION_ID", "sibyl-run301-baseline"), help="Isolated Honcho session id.")
    parser.add_argument("--peer-id", default=os.environ.get("HONCHO_PEER_ID", "sibyl-run301-source"), help="Honcho peer id used for ingestion.")
    parser.add_argument("--batch-size", type=int, default=100, help="Honcho ingestion batch size.")
    parser.add_argument("--limit", type=int, default=LIMIT, help="Search results per question.")
    parser.add_argument("--company-count", type=int, default=COMPANY_COUNT, help="Total companies in the generated chronology dataset.")
    parser.add_argument("--timeline-days", type=int, default=TIMELINE_DAYS, help="Total simulated days per company.")
    parser.add_argument("--company-limit", type=int, default=None, help="Only generate records for the first N companies.")
    parser.add_argument("--question-limit", type=int, default=None, help="Only run the first N filtered questions.")
    parser.add_argument("--answer-llm", action="store_true", help="Answer from Honcho context with Sonnet and score like run301.")
    parser.add_argument("--answer-model", default=DEFAULT_SONNET_MODEL, help="Anthropic Sonnet model name.")
    parser.add_argument("--answer-max-tokens", type=int, default=800, help="Max output tokens for Sonnet answers.")
    parser.add_argument("--no-answer-followup", action="store_true", help="Disable natural follow-up repair for Sonnet answers.")
    parser.add_argument("--output-run-id", default=RUN_ID, help="Output run id.")
    return parser.parse_args()


def company_index_from_case_id(case_id: str) -> int | None:
    match = re.search(r"company_(\d{3})", case_id)
    return int(match.group(1)) if match else None


def filter_questions(questions: list[dict[str, Any]], *, company_limit: int | None, question_limit: int | None) -> list[dict[str, Any]]:
    filtered: list[dict[str, Any]] = []
    for question in questions:
        index = company_index_from_case_id(str(question.get("id", "")))
        if company_limit is not None and index is not None and index > company_limit:
            continue
        filtered.append(question)
    return filtered[:question_limit] if question_limit is not None else filtered


def main() -> None:
    args = parse_args()
    if args.answer_llm and not args.execute:
        raise SystemExit("--answer-llm requires --execute so paid LLM calls are never made during dry-run.")
    if args.answer_llm:
        load_env(ANTHROPIC_ENV)
    source_raw = read_json(Path(args.source_raw))
    records = build_records(
        company_count=args.company_count,
        timeline_days=args.timeline_days,
        company_limit=args.company_limit,
    )
    questions = filter_questions(
        source_raw["questions"],
        company_limit=args.company_limit,
        question_limit=args.question_limit,
    )

    if not args.execute:
        print(
            json.dumps(
                {
                    "mode": "dry_run",
                    "source_run": source_raw["run_id"],
                    "records_to_ingest": len(records),
                    "questions": len(questions),
                    "company_count": args.company_count,
                    "timeline_days": args.timeline_days,
                    "company_limit": args.company_limit,
                    "question_limit": args.question_limit,
                    "next": "Set HONCHO_API_KEY/HONCHO_BASE_URL if needed, then run with --execute --ingest.",
                },
                indent=2,
            )
        )

    if args.execute and args.ingest:
        ingest_honcho(
            records,
            workspace_id=args.workspace_id,
            session_id=args.session_id,
            peer_id=args.peer_id,
            batch_size=args.batch_size,
        )
        # Honcho processes memory asynchronously. Give small local/cloud setups a short window.
        time.sleep(3)

    rows: list[dict[str, Any]] = []
    for question in questions:
        if args.execute:
            context = search_honcho(
                workspace_id=args.workspace_id,
                session_id=args.session_id,
                query=question["retrieval_query"],
                limit=args.limit,
            )
        else:
            context = local_keyword_search(records, question["retrieval_query"], args.limit)
        row = row_for_question(question, context)
        if args.answer_llm:
            print(f"[answer] {question['id']}", file=sys.stderr, flush=True)
            row["honcho_sonnet"] = answer_with_sonnet(
                question,
                context,
                model=args.answer_model,
                max_tokens=args.answer_max_tokens,
                followup=not args.no_answer_followup,
            )
        rows.append(row)

    summary = {
        "total": len(rows),
        "sibyl_retrieval_passed": sum(1 for item in rows if item["sibyl_retrieval_passed"]),
        "honcho_retrieval_passed": sum(1 for item in rows if item["honcho_retrieval_passed"]),
        "sibyl_final_answer_passed": sum(1 for item in rows if item["sibyl_final_answer_passed"]),
        "avg_sibyl_context_tokens": sum(item["sibyl_context_tokens_estimate"] for item in rows) / len(rows),
        "avg_honcho_context_tokens": sum(item["honcho_context_tokens_estimate"] for item in rows) / len(rows),
    }
    answer_mode = "none"
    if args.answer_llm:
        answer_mode = "sonnet"
        answer_cost = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "estimated_input_usd": 0.0,
            "estimated_output_usd": 0.0,
            "estimated_total_usd": 0.0,
        }
        for item in rows:
            add_cost(answer_cost, item["honcho_sonnet"]["cost"])
        summary["honcho_sonnet_answer_passed"] = sum(1 for item in rows if item["honcho_sonnet"]["score"]["passed"])
        summary["honcho_sonnet_initial_passed"] = sum(1 for item in rows if item["honcho_sonnet"]["initial_score"]["passed"])
        summary["honcho_sonnet_followups"] = sum(1 for item in rows if item["honcho_sonnet"]["followup"] is not None)
        summary["honcho_sonnet_model"] = args.answer_model
        summary["honcho_sonnet_cost"] = answer_cost
    raw = {
        "run_id": args.output_run_id,
        "ts_started": iso_now(),
        "ts_completed": iso_now(),
        "source_run_id": source_raw["run_id"],
        "source_raw": str(Path(args.source_raw)),
        "honcho_mode": "sdk" if args.execute else "local_keyword_dry_run",
        "answer_mode": answer_mode,
        "honcho_workspace_id": args.workspace_id,
        "honcho_session_id": args.session_id,
        "honcho_peer_id": args.peer_id,
        "dataset": {
            "record_count": len(records),
            "company_count": args.company_limit or args.company_count,
            "stakeholder_count": (args.company_limit or args.company_count) * len(ROLES),
            "timeline_days": args.timeline_days,
        },
        "summary": summary,
        "questions": rows,
    }

    if args.execute or args.write_dry_run:
        RUNS_DIR.mkdir(parents=True, exist_ok=True)
        raw_path = RUNS_DIR / f"{args.output_run_id}.raw_result.json"
        md_path = RUNS_DIR / f"{args.output_run_id}.md"
        write_json(raw_path, raw)
        md_path.write_text(build_markdown(raw), encoding="utf8")
        print(json.dumps({"summary": summary, "raw": str(raw_path), "md": str(md_path)}, indent=2))
    else:
        print(json.dumps({"summary": summary, "artifacts": "not_written_dry_run"}, indent=2))


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)
